Extend SILModule with an API to serialize its content

The `serialize` method can be called multiple times, but it will perform the actual serialization only the first time.
By means of this API we get the flexibility to serialize the SILModule not only after all the optimizations, but  e.g. at any time during optimizations.

Author: swiftix

include/swift/SIL/SILModule.h

@@ -100,6 +100,7 @@ class SILModule {
   using DefaultWitnessTableListType = llvm::ilist<SILDefaultWitnessTable>;
   using CoverageMapListType = llvm::ilist<SILCoverageMap>;
   using LinkingMode = SILOptions::LinkingMode;
+  using ActionCallback = std::function<void()>;
 
 private:
   friend KeyPathPattern;
@@ -213,6 +214,13 @@ class SILModule {
   /// The options passed into this SILModule.
   SILOptions &Options;
 
+  /// Set if the SILModule was serialized already. It is used
+  /// to ensure that the module is serialized only once.
+  bool serialized;
+
+  /// Action to be executed for serializing the SILModule.
+  ActionCallback SerializeSILAction;
+
   /// A list of clients that need to be notified when an instruction
   /// invalidation message is sent.
   llvm::SetVector<DeleteNotificationHandler*> NotificationHandlers;
@@ -251,6 +259,17 @@ class SILModule {
   /// registered handlers. The order of handlers is deterministic but arbitrary.
   void notifyDeleteHandlers(SILNode *node);
 
+  /// Set a serialization action.
+  void setSerializeSILAction(ActionCallback SerializeSILAction);
+  ActionCallback getSerializeSILAction() const;
+
+  /// Set a flag indicating that this module is serialized already.
+  void setSerialized() { serialized = true; }
+  bool isSerialized() const { return serialized; }
+
+  /// Serialize a SIL module using the configured SerializeSILAction.
+  void serialize();
+
   /// \brief This converts Swift types to SILTypes.
   mutable Lowering::TypeConverter Types;
 

lib/FrontendTool/FrontendTool.cpp

@@ -854,6 +854,46 @@ static bool performCompile(CompilerInstance &Instance,
     SM->verify();
   }
 
+  // This is the action to be used to serialize SILModule.
+  // It may be invoked multiple times, but it will perform
+  // serialization only once. The serialization may either happen
+  // after high-level optimizations or after all optimizations are
+  // done, depending on the compiler setting.
+
+  auto SerializeSILModuleAction = [&]() {
+    if (!opts.ModuleOutputPath.empty() || !opts.ModuleDocOutputPath.empty()) {
+      auto DC = PrimarySourceFile ? ModuleOrSourceFile(PrimarySourceFile)
+                                  : Instance.getMainModule();
+      if (!opts.ModuleOutputPath.empty()) {
+        SerializationOptions serializationOpts;
+        serializationOpts.OutputPath = opts.ModuleOutputPath.c_str();
+        serializationOpts.DocOutputPath = opts.ModuleDocOutputPath.c_str();
+        serializationOpts.GroupInfoPath = opts.GroupInfoPath.c_str();
+        if (opts.SerializeBridgingHeader)
+          serializationOpts.ImportedHeader = opts.ImplicitObjCHeaderPath;
+        serializationOpts.ModuleLinkName = opts.ModuleLinkName;
+        serializationOpts.ExtraClangOptions =
+            Invocation.getClangImporterOptions().ExtraArgs;
+        serializationOpts.EnableNestedTypeLookupTable =
+            opts.EnableSerializationNestedTypeLookupTable;
+        if (!IRGenOpts.ForceLoadSymbolName.empty())
+          serializationOpts.AutolinkForceLoad = true;
+
+        // Options contain information about the developer's computer,
+        // so only serialize them if the module isn't going to be shipped to
+        // the public.
+        serializationOpts.SerializeOptionsForDebugging =
+            !moduleIsPublic || opts.AlwaysSerializeDebuggingOptions;
+
+        serialize(DC, serializationOpts, SM.get());
+      }
+    }
+  };
+
+  // Set the serialization action, so that the SIL module
+  // can be serialized at any moment, e.g. during the optimization pipeline.
+  SM->setSerializeSILAction(SerializeSILModuleAction);
+
   // Perform SIL optimization passes if optimizations haven't been disabled.
   // These may change across compiler versions.
   {
@@ -921,32 +961,9 @@ static bool performCompile(CompilerInstance &Instance,
   }
 
   if (!opts.ModuleOutputPath.empty() || !opts.ModuleDocOutputPath.empty()) {
-    auto DC = PrimarySourceFile ? ModuleOrSourceFile(PrimarySourceFile) :
-                                  Instance.getMainModule();
-    if (!opts.ModuleOutputPath.empty()) {
-      SerializationOptions serializationOpts;
-      serializationOpts.OutputPath = opts.ModuleOutputPath.c_str();
-      serializationOpts.DocOutputPath = opts.ModuleDocOutputPath.c_str();
-      serializationOpts.GroupInfoPath = opts.GroupInfoPath.c_str();
-      if (opts.SerializeBridgingHeader)
-        serializationOpts.ImportedHeader = opts.ImplicitObjCHeaderPath;
-      serializationOpts.ModuleLinkName = opts.ModuleLinkName;
-      serializationOpts.ExtraClangOptions =
-          Invocation.getClangImporterOptions().ExtraArgs;
-      serializationOpts.EnableNestedTypeLookupTable =
-          opts.EnableSerializationNestedTypeLookupTable;
-      if (!IRGenOpts.ForceLoadSymbolName.empty())
-        serializationOpts.AutolinkForceLoad = true;
-
-      // Options contain information about the developer's computer,
-      // so only serialize them if the module isn't going to be shipped to
-      // the public.
-      serializationOpts.SerializeOptionsForDebugging =
-          !moduleIsPublic || opts.AlwaysSerializeDebuggingOptions;
-
-      serialize(DC, serializationOpts, SM.get());
-    }
-
+    // Serialize the SILModule if it was not serialized yet.
+    if (!SM.get()->isSerialized())
+      SM.get()->serialize();
     if (Action == FrontendOptions::MergeModules ||
         Action == FrontendOptions::EmitModuleOnly) {
       if (shouldIndex) {

lib/SIL/SILModule.cpp

@@ -87,7 +87,8 @@ SILModule::SILModule(ModuleDecl *SwiftModule, SILOptions &Options,
                      const DeclContext *associatedDC, bool wholeModule)
     : TheSwiftModule(SwiftModule), AssociatedDeclContext(associatedDC),
       Stage(SILStage::Raw), Callback(new SILModule::SerializationCallback()),
-      wholeModule(wholeModule), Options(Options), Types(*this) {}
+      wholeModule(wholeModule), Options(Options), serialized(false),
+      SerializeSILAction(), Types(*this) {}
 
 SILModule::~SILModule() {
   // Decrement ref count for each SILGlobalVariable with static initializers.
@@ -762,3 +763,20 @@ bool SILModule::isOptimizedOnoneSupportModule() const {
   return getOptions().Optimization >= SILOptions::SILOptMode::Optimize &&
          isOnoneSupportModule();
 }
+
+void SILModule::setSerializeSILAction(SILModule::ActionCallback Action) {
+  assert(!SerializeSILAction && "Serialization action can be set only once");
+  SerializeSILAction = Action;
+}
+
+SILModule::ActionCallback SILModule::getSerializeSILAction() const {
+  return SerializeSILAction;
+}
+
+void SILModule::serialize() {
+  assert(SerializeSILAction && "Serialization action should be set");
+  assert(!isSerialized() && "The module was serialized already");
+  SerializeSILAction();
+  setSerialized();
+}
+

tools/sil-opt/SILOpt.cpp

@@ -383,6 +383,9 @@ int main(int argc, char **argv) {
   if (VerifyMode)
     enableDiagnosticVerifier(CI.getSourceMgr());
 
+  if (CI.getSILModule())
+    CI.getSILModule()->setSerializeSILAction([]{});
+
   if (OptimizationGroup == OptGroup::Diagnostics) {
     runSILDiagnosticPasses(*CI.getSILModule());
   } else if (OptimizationGroup == OptGroup::Performance) {