Convert OME into a function transform

Author: meg-gupta

lib/SILOptimizer/Mandatory/OwnershipModelEliminator.cpp

@@ -370,7 +370,7 @@ static void prepareSILFunctionForOptimization(ModuleDecl *, SILFunction *F) {
 
 namespace {
 
-struct OwnershipModelEliminator : SILModuleTransform {
+struct OwnershipModelEliminator : SILFunctionTransform {
   bool SkipTransparent;
   bool SkipStdlibModule;
 
@@ -379,53 +379,50 @@ struct OwnershipModelEliminator : SILModuleTransform {
 
   void run() override {
     if (DumpBefore.size()) {
-      getModule()->dump(DumpBefore.c_str());
+      getFunction()->dump(DumpBefore.c_str());
     }
 
-    auto &Mod = *getModule();
+    auto *F = getFunction();
+    auto &Mod = getFunction()->getModule();
 
     // If we are supposed to skip the stdlib module and we are in the stdlib
     // module bail.
     if (SkipStdlibModule && Mod.isStdlibModule()) {
       return;
     }
 
-    for (auto &F : Mod) {
-      // If F does not have ownership, skip it. We have no further work to do.
-      if (!F.hasOwnership())
-        continue;
-
-      // If we were asked to not strip ownership from transparent functions in
-      // /our/ module, continue.
-      if (SkipTransparent && F.isTransparent())
-        continue;
-
-      // Verify here to make sure ownership is correct before we strip.
-      {
-        // Add a pretty stack trace entry to tell users who see a verification
-        // failure triggered by this verification check that they need to re-run
-        // with -sil-verify-all to actually find the pass that introduced the
-        // verification error.
-        //
-        // DISCUSSION: This occurs due to the crash from the verification
-        // failure happening in the pass itself. This causes us to dump the
-        // SILFunction and emit a msg that this pass (OME) is the culprit. This
-        // is generally correct for most passes, but not for OME since we are
-        // verifying before we have even modified the function to ensure that
-        // all ownership invariants have been respected before we lower
-        // ownership from the function.
-        llvm::PrettyStackTraceString silVerifyAllMsgOnFailure(
-            "Found verification error when verifying before lowering "
-            "ownership. Please re-run with -sil-verify-all to identify the "
-            "actual pass that introduced the verification error.");
-        F.verify();
-      }
-
-      if (stripOwnership(F)) {
-        auto InvalidKind =
-            SILAnalysis::InvalidationKind::BranchesAndInstructions;
-        invalidateAnalysis(&F, InvalidKind);
-      }
+    if (!F->hasOwnership())
+      return;
+
+    // If we were asked to not strip ownership from transparent functions in
+    // /our/ module, return.
+    if (SkipTransparent && F->isTransparent())
+      return;
+
+    // Verify here to make sure ownership is correct before we strip.
+    {
+      // Add a pretty stack trace entry to tell users who see a verification
+      // failure triggered by this verification check that they need to re-run
+      // with -sil-verify-all to actually find the pass that introduced the
+      // verification error.
+      //
+      // DISCUSSION: This occurs due to the crash from the verification
+      // failure happening in the pass itself. This causes us to dump the
+      // SILFunction and emit a msg that this pass (OME) is the culprit. This
+      // is generally correct for most passes, but not for OME since we are
+      // verifying before we have even modified the function to ensure that
+      // all ownership invariants have been respected before we lower
+      // ownership from the function.
+      llvm::PrettyStackTraceString silVerifyAllMsgOnFailure(
+          "Found verification error when verifying before lowering "
+          "ownership. Please re-run with -sil-verify-all to identify the "
+          "actual pass that introduced the verification error.");
+      F->verify();
+    }
+
+    if (stripOwnership(*F)) {
+      auto InvalidKind = SILAnalysis::InvalidationKind::BranchesAndInstructions;
+      invalidateAnalysis(InvalidKind);
     }
 
     // If we were asked to strip transparent, we are at the beginning of the