1 files changed, 14 insertions, 20 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h b/llvm/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h
index 0a067501170..da5fb6ef5fc 100644
--- a/llvm/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h
+++ b/llvm/lib/ExecutionEngine/Orc/OrcMCJITReplacement.h
@@ -54,6 +54,7 @@ class ObjectCache;
 namespace orc {
 
 class OrcMCJITReplacement : public ExecutionEngine {
+
   // OrcMCJITReplacement needs to do a little extra book-keeping to ensure that
   // Orc's automatic finalization doesn't kick in earlier than MCJIT clients are
   // expecting - see finalizeMemory.
@@ -235,7 +236,10 @@ public:
               return ObjectLayerT::Resources{this->MemMgr, this->Resolver};
             },
             NotifyObjectLoaded, NotifyFinalized),
-        CompileLayer(ObjectLayer, SimpleCompiler(*this->TM)),
+        CompileLayer(ObjectLayer, SimpleCompiler(*this->TM),
+                     [this](VModuleKey K, std::unique_ptr<Module> M) {
+                       Modules.push_back(std::move(M));
+                     }),
         LazyEmitLayer(CompileLayer) {}
 
   static void Register() {
@@ -250,16 +254,7 @@ public:
     } else {
       assert(M->getDataLayout() == getDataLayout() && "DataLayout Mismatch");
     }
-    auto *MPtr = M.release();
-    ShouldDelete[MPtr] = true;
-    auto Deleter = [this](Module *Mod) {
-      auto I = ShouldDelete.find(Mod);
-      if (I != ShouldDelete.end() && I->second)
-        delete Mod;
-    };
-    LocalModules.push_back(std::shared_ptr<Module>(MPtr, std::move(Deleter)));
-    cantFail(
-        LazyEmitLayer.addModule(ES.allocateVModule(), LocalModules.back()));
+    cantFail(LazyEmitLayer.addModule(ES.allocateVModule(), std::move(M)));
   }
 
   void addObjectFile(std::unique_ptr<object::ObjectFile> O) override {
@@ -279,14 +274,14 @@ public:
   }
   
   bool removeModule(Module *M) override {
-    for (auto I = LocalModules.begin(), E = LocalModules.end(); I != E; ++I) {
-      if (I->get() == M) {
-        ShouldDelete[M] = false;
-        LocalModules.erase(I);
-        return true;
-      }
-    }
-    return false;
+    auto I = Modules.begin();
+    for (auto E = Modules.end(); I != E; ++I)
+      if (I->get() == M)
+        break;
+    if (I == Modules.end())
+      return false;
+    Modules.erase(I);
+    return true;
   }
 
   uint64_t getSymbolAddress(StringRef Name) {
@@ -438,7 +433,6 @@ private:
   // delete blocks in LocalModules refer to the ShouldDelete map, so
   // LocalModules needs to be destructed before ShouldDelete.
   std::map<Module*, bool> ShouldDelete;
-  std::vector<std::shared_ptr<Module>> LocalModules;
 
   NotifyObjectLoadedT NotifyObjectLoaded;
   NotifyFinalizedT NotifyFinalized;