PassManager: cleanup the analysis-invalidation mechanism for swift passes

This fixes a bug with module passes.

Author: eeckstein

include/swift/SILOptimizer/PassManager/PassManager.h

@@ -61,6 +61,10 @@ class SwiftPassInvocation {
   /// Non-null if this is an instruction pass, invoked from SILCombine.
   SILCombiner *silCombiner = nullptr;
 
+  /// Change notifications, collected during a pass run.
+  SILAnalysis::InvalidationKind changeNotifications =
+      SILAnalysis::InvalidationKind::Nothing;
+
   /// All slabs, allocated by the pass.
   SILModule::SlabList allocatedSlabs;
 
@@ -126,17 +130,9 @@ class SwiftPassInvocation {
   /// Called by the SILCombiner when the instruction pass has finished.
   void finishedInstructionPassRun();
 
-  void beginTransformFunction(SILFunction *function) {
-    assert(!this->function && transform && "not running a module pass");
-    this->function = function;
-  }
+  void beginTransformFunction(SILFunction *function);
 
-  void endTransformFunction() {
-    assert(function && "beginTransformFunction not called");
-    function = nullptr;
-    assert(numBlockSetsAllocated == 0 && "Not all BasicBlockSets deallocated");
-    assert(numNodeSetsAllocated == 0 && "Not all NodeSets deallocated");
-  }
+  void endTransformFunction();
 };
 
 /// The SIL pass manager.
@@ -184,10 +180,6 @@ class SILPassManager {
   /// For invoking Swift passes.
   SwiftPassInvocation swiftPassInvocation;
 
-  /// Change notifications, collected during a bridged pass run.
-  SILAnalysis::InvalidationKind changeNotifications =
-      SILAnalysis::InvalidationKind::Nothing;
-
   /// A mask which has one bit for each pass. A one for a pass-bit means that
   /// the pass doesn't need to run, because nothing has changed since the
   /// previous run of that pass.
@@ -341,11 +333,6 @@ class SILPassManager {
     CompletedPassesMap[F].reset();
   }
 
-  void notifyPassChanges(SILAnalysis::InvalidationKind invalidationKind) {
-    changeNotifications = (SILAnalysis::InvalidationKind)
-        (changeNotifications | invalidationKind);
-  }
-
   /// Reset the state of the pass manager and remove all transformation
   /// owned by the pass manager. Analysis passes will be kept.
   void resetAndRemoveTransformations();
@@ -440,7 +427,8 @@ class SILPassManager {
 
 inline void SwiftPassInvocation::
 notifyChanges(SILAnalysis::InvalidationKind invalidationKind) {
-  passManager->notifyPassChanges(invalidationKind);
+    changeNotifications = (SILAnalysis::InvalidationKind)
+        (changeNotifications | invalidationKind);
 }
 
 } // end namespace swift

lib/SILOptimizer/PassManager/PassManager.cpp

@@ -530,9 +530,6 @@ void SILPassManager::runPassOnFunction(unsigned TransIdx, SILFunction *F) {
     forcePrecomputeAnalyses(F);
   }
 
-  assert(changeNotifications == SILAnalysis::InvalidationKind::Nothing
-         && "change notifications not cleared");
-
   llvm::sys::TimePoint<> startTime = std::chrono::system_clock::now();
   std::chrono::nanoseconds duration(0);
 
@@ -553,24 +550,21 @@ void SILPassManager::runPassOnFunction(unsigned TransIdx, SILFunction *F) {
     // Run it!
     SFT->run();
 
-    if (CurrentPassHasInvalidated ||
-        changeNotifications != SILAnalysis::InvalidationKind::Nothing) {
+    swiftPassInvocation.finishedFunctionPassRun();
+
+    if (CurrentPassHasInvalidated) {
       // Pause time measurement while invalidating analysis and restoring the snapshot.
       duration += (std::chrono::system_clock::now() - startTime);
 
       if (runIdx < numRepeats - 1) {
         invalidateAnalysis(F, SILAnalysis::InvalidationKind::Everything);
         F->restoreFromSnapshot(SnapshotID);
-      } else if (changeNotifications != SILAnalysis::InvalidationKind::Nothing) {
-        invalidateAnalysis(F, changeNotifications);
       }
-      changeNotifications = SILAnalysis::InvalidationKind::Nothing;
 
       // Continue time measurement (including flushing deleted instructions).
       startTime = std::chrono::system_clock::now();
     }
     Mod->flushDeletedInsts();
-    swiftPassInvocation.finishedFunctionPassRun();
   }
 
   duration += (std::chrono::system_clock::now() - startTime);
@@ -1289,9 +1283,9 @@ void SwiftPassInvocation::startModulePassRun(SILModuleTransform *transform) {
 }
 
 void SwiftPassInvocation::startFunctionPassRun(SILFunctionTransform *transform) {
-  assert(!this->function && !this->transform && "a pass is already running");
-  this->function = transform->getFunction();
+  assert(!this->transform && "a pass is already running");
   this->transform = transform;
+  beginTransformFunction(transform->getFunction());
 }
 
 void SwiftPassInvocation::startInstructionPassRun(SILInstruction *inst) {
@@ -1302,13 +1296,15 @@ void SwiftPassInvocation::startInstructionPassRun(SILInstruction *inst) {
 void SwiftPassInvocation::finishedModulePassRun() {
   endPassRunChecks();
   assert(!function && transform && "not running a pass");
+  assert(changeNotifications == SILAnalysis::InvalidationKind::Nothing
+         && "unhandled change notifications at end of module pass");
   transform = nullptr;
 }
 
 void SwiftPassInvocation::finishedFunctionPassRun() {
   endPassRunChecks();
-  assert(function && transform && "not running a pass");
-  function = nullptr;
+  endTransformFunction();
+  assert(allocatedSlabs.empty() && "StackList is leaking slabs");
   transform = nullptr;
 }
 
@@ -1322,6 +1318,24 @@ void SwiftPassInvocation::endPassRunChecks() {
   assert(numNodeSetsAllocated == 0 && "Not all NodeSets deallocated");
 }
 
+void SwiftPassInvocation::beginTransformFunction(SILFunction *function) {
+  assert(!this->function && transform && "not running a pass");
+  assert(changeNotifications == SILAnalysis::InvalidationKind::Nothing
+         && "change notifications not cleared");
+  this->function = function;
+}
+
+void SwiftPassInvocation::endTransformFunction() {
+  assert(function && transform && "not running a pass");
+  if (changeNotifications != SILAnalysis::InvalidationKind::Nothing) {
+    passManager->invalidateAnalysis(function, changeNotifications);
+    changeNotifications = SILAnalysis::InvalidationKind::Nothing;
+  }
+  function = nullptr;
+  assert(numBlockSetsAllocated == 0 && "Not all BasicBlockSets deallocated");
+  assert(numNodeSetsAllocated == 0 && "Not all NodeSets deallocated");
+}
+
 //===----------------------------------------------------------------------===//
 //                            Swift Bridging
 //===----------------------------------------------------------------------===//