1 files changed, 44 insertions, 3 deletions

diff --git a/lib/VMCore/PassManager.cpp b/lib/VMCore/PassManager.cpp
index 86cf10e..46f1243 100644
--- a/lib/VMCore/PassManager.cpp
+++ b/lib/VMCore/PassManager.cpp
@@ -165,11 +165,13 @@ namespace llvm {
 class FunctionPassManagerImpl : public Pass,
                                 public PMDataManager,
                                 public PMTopLevelManager {
+private:
+  bool wasRun;
 public:
   static char ID;
   explicit FunctionPassManagerImpl(int Depth) : 
     Pass(&ID), PMDataManager(Depth), 
-    PMTopLevelManager(TLM_Function) { }
+    PMTopLevelManager(TLM_Function), wasRun(false) { }
 
   /// add - Add a pass to the queue of passes to run.  This passes ownership of
   /// the Pass to the PassManager.  When the PassManager is destroyed, the pass
@@ -179,6 +181,10 @@ public:
     schedulePass(P);
   }
  
+  // Prepare for running an on the fly pass, freeing memory if needed
+  // from a previous run.
+  void releaseMemoryOnTheFly();
+
   /// run - Execute all of the passes scheduled for execution.  Keep track of
   /// whether any of the passes modifies the module, and if so, return true.
   bool run(Function &F);
@@ -272,8 +278,10 @@ public:
     for (unsigned Index = 0; Index < getNumContainedPasses(); ++Index) {
       ModulePass *MP = getContainedPass(Index);
       MP->dumpPassStructure(Offset + 1);
-      if (FunctionPassManagerImpl *FPP = OnTheFlyManagers[MP])
-        FPP->dumpPassStructure(Offset + 2);
+      std::map<Pass *, FunctionPassManagerImpl *>::const_iterator I =
+        OnTheFlyManagers.find(MP);
+      if (I != OnTheFlyManagers.end())
+        I->second->dumpPassStructure(Offset + 2);
       dumpLastUses(MP, Offset+1);
     }
   }
@@ -1290,6 +1298,18 @@ void FPPassManager::cleanup() {
  }
 }
 
+void FunctionPassManagerImpl::releaseMemoryOnTheFly() {
+  if (!wasRun)
+    return;
+  for (unsigned Index = 0; Index < getNumContainedManagers(); ++Index) {
+    FPPassManager *FPPM = getContainedManager(Index);
+    for (unsigned Index = 0; Index < FPPM->getNumContainedPasses(); ++Index) {
+      FPPM->getContainedPass(Index)->releaseMemory();
+    }
+  }
+  wasRun = false;
+}
+
 // Execute all the passes managed by this top level manager.
 // Return true if any function is modified by a pass.
 bool FunctionPassManagerImpl::run(Function &F) {
@@ -1306,6 +1326,7 @@ bool FunctionPassManagerImpl::run(Function &F) {
   for (unsigned Index = 0; Index < getNumContainedManagers(); ++Index)
     getContainedManager(Index)->cleanup();
 
+  wasRun = true;
   return Changed;
 }
 
@@ -1404,6 +1425,14 @@ bool
 MPPassManager::runOnModule(Module &M) {
   bool Changed = false;
 
+  // Initialize on-the-fly passes
+  for (std::map<Pass *, FunctionPassManagerImpl *>::iterator
+       I = OnTheFlyManagers.begin(), E = OnTheFlyManagers.end();
+       I != E; ++I) {
+    FunctionPassManagerImpl *FPP = I->second;
+    Changed |= FPP->doInitialization(M);
+  }
+
   for (unsigned Index = 0; Index < getNumContainedPasses(); ++Index) {
     ModulePass *MP = getContainedPass(Index);
 
@@ -1430,6 +1459,17 @@ MPPassManager::runOnModule(Module &M) {
     recordAvailableAnalysis(MP);
     removeDeadPasses(MP, M.getModuleIdentifier().c_str(), ON_MODULE_MSG);
   }
+
+  // Finalize on-the-fly passes
+  for (std::map<Pass *, FunctionPassManagerImpl *>::iterator
+       I = OnTheFlyManagers.begin(), E = OnTheFlyManagers.end();
+       I != E; ++I) {
+    FunctionPassManagerImpl *FPP = I->second;
+    // We don't know when is the last time an on-the-fly pass is run,
+    // so we need to releaseMemory / finalize here
+    FPP->releaseMemoryOnTheFly();
+    Changed |= FPP->doFinalization(M);
+  }
   return Changed;
 }
 
@@ -1466,6 +1506,7 @@ Pass* MPPassManager::getOnTheFlyPass(Pass *MP, const PassInfo *PI, Function &F){
   FunctionPassManagerImpl *FPP = OnTheFlyManagers[MP];
   assert(FPP && "Unable to find on the fly pass");
   
+  FPP->releaseMemoryOnTheFly();
   FPP->run(F);
   return (dynamic_cast<PMTopLevelManager *>(FPP))->findAnalysisPass(PI);
 }