Register deserialization notification handlers in ome only once

Author: meg-gupta

include/swift/SIL/SILModule.h

@@ -263,6 +263,15 @@ class SILModule {
   /// to ensure that the module is serialized only once.
   bool serialized;
 
+  /// Set if we have registered a deserialization notification handler for
+  /// lowering ownership in non transparent functions.
+  /// This gets set in NonTransparent OwnershipModelEliminator pass.
+  bool regDeserializationNotificationHandlerForNonTransparentFuncOME;
+  /// Set if we have registered a deserialization notification handler for
+  /// lowering ownership in transparent functions.
+  /// This gets set in OwnershipModelEliminator pass.
+  bool regDeserializationNotificationHandlerForAllFuncOME;
+
   /// Action to be executed for serializing the SILModule.
   ActionCallback SerializeSILAction;
 
@@ -301,6 +310,19 @@ class SILModule {
     deserializationNotificationHandlers.erase(handler);
   }
 
+  bool hasRegisteredDeserializationNotificationHandlerForNonTransparentFuncOME() {
+    return regDeserializationNotificationHandlerForNonTransparentFuncOME;
+  }
+  bool hasRegisteredDeserializationNotificationHandlerForAllFuncOME() {
+    return regDeserializationNotificationHandlerForAllFuncOME;
+  }
+  void setRegisteredDeserializationNotificationHandlerForNonTransparentFuncOME() {
+    regDeserializationNotificationHandlerForNonTransparentFuncOME = true;
+  }
+  void setRegisteredDeserializationNotificationHandlerForAllFuncOME() {
+    regDeserializationNotificationHandlerForAllFuncOME = true;
+  }
+
   /// Add a delete notification handler \p Handler to the module context.
   void registerDeleteNotificationHandler(DeleteNotificationHandler* Handler);
 

lib/SIL/IR/SILModule.cpp

@@ -95,6 +95,8 @@ class SILModule::SerializationCallback final
 SILModule::SILModule(llvm::PointerUnion<FileUnit *, ModuleDecl *> context,
                      Lowering::TypeConverter &TC, const SILOptions &Options)
     : Stage(SILStage::Raw), Options(Options), serialized(false),
+      regDeserializationNotificationHandlerForNonTransparentFuncOME(false),
+      regDeserializationNotificationHandlerForAllFuncOME(false),
       SerializeSILAction(), Types(TC) {
   assert(!context.isNull());
   if (auto *file = context.dyn_cast<FileUnit *>()) {

lib/SILOptimizer/Mandatory/OwnershipModelEliminator.cpp

@@ -432,12 +432,19 @@ struct OwnershipModelEliminator : SILFunctionTransform {
         FunctionBodyDeserializationNotificationHandler;
     std::unique_ptr<DeserializationNotificationHandler> ptr;
     if (SkipTransparent) {
-      ptr.reset(new NotificationHandlerTy(
-          prepareNonTransparentSILFunctionForOptimization));
+      if (!Mod.hasRegisteredDeserializationNotificationHandlerForNonTransparentFuncOME()) {
+        ptr.reset(new NotificationHandlerTy(
+            prepareNonTransparentSILFunctionForOptimization));
+        Mod.registerDeserializationNotificationHandler(std::move(ptr));
+        Mod.setRegisteredDeserializationNotificationHandlerForNonTransparentFuncOME();
+      }
     } else {
-      ptr.reset(new NotificationHandlerTy(prepareSILFunctionForOptimization));
+      if (!Mod.hasRegisteredDeserializationNotificationHandlerForAllFuncOME()) {
+        ptr.reset(new NotificationHandlerTy(prepareSILFunctionForOptimization));
+        Mod.registerDeserializationNotificationHandler(std::move(ptr));
+        Mod.setRegisteredDeserializationNotificationHandlerForAllFuncOME();
+      }
     }
-    Mod.registerDeserializationNotificationHandler(std::move(ptr));
   }
 };