ExecutionEngine::clearGlobalMappingsFromModule failed to remove reverse

mappings, which could cause errors and assert-failures.  This patch fixes that,
adds a test, and refactors the global-mapping-removal code into a single place.


git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@83678 91177308-0d34-0410-b5e6-96231b3b80d8

3 files changed, 42 insertions, 16 deletions

diff --git a/include/llvm/ExecutionEngine/ExecutionEngine.h b/include/llvm/ExecutionEngine/ExecutionEngine.h
index 6a3a914..cf2df6b 100644
--- a/include/llvm/ExecutionEngine/ExecutionEngine.h
+++ b/include/llvm/ExecutionEngine/ExecutionEngine.h
@@ -61,6 +61,9 @@ public:
   getGlobalAddressReverseMap(const MutexGuard &) {
     return GlobalAddressReverseMap;
   }
+
+  // Returns the address ToUnmap was mapped to.
+  void *RemoveMapping(const MutexGuard &, const GlobalValue *ToUnmap);
 };
 
 
diff --git a/lib/ExecutionEngine/ExecutionEngine.cpp b/lib/ExecutionEngine/ExecutionEngine.cpp
index 52718eb..700ae5e 100644
--- a/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -113,6 +113,22 @@ Function *ExecutionEngine::FindFunctionNamed(const char *FnName) {
 }
 
 
+void *ExecutionEngineState::RemoveMapping(
+  const MutexGuard &, const GlobalValue *ToUnmap) {
+  std::map<AssertingVH<const GlobalValue>, void *>::iterator I =
+    GlobalAddressMap.find(ToUnmap);
+  void *OldVal;
+  if (I == GlobalAddressMap.end())
+    OldVal = 0;
+  else {
+    OldVal = I->second;
+    GlobalAddressMap.erase(I);
+  }
+
+  GlobalAddressReverseMap.erase(OldVal);
+  return OldVal;
+}
+
 /// addGlobalMapping - Tell the execution engine that the specified global is
 /// at the specified location.  This is used internally as functions are JIT'd
 /// and as global variables are laid out in memory.  It can and should also be
@@ -151,13 +167,11 @@ void ExecutionEngine::clearGlobalMappingsFromModule(Module *M) {
   MutexGuard locked(lock);
   
   for (Module::iterator FI = M->begin(), FE = M->end(); FI != FE; ++FI) {
-    state.getGlobalAddressMap(locked).erase(&*FI);
-    state.getGlobalAddressReverseMap(locked).erase(&*FI);
+    state.RemoveMapping(locked, FI);
   }
   for (Module::global_iterator GI = M->global_begin(), GE = M->global_end(); 
        GI != GE; ++GI) {
-    state.getGlobalAddressMap(locked).erase(&*GI);
-    state.getGlobalAddressReverseMap(locked).erase(&*GI);
+    state.RemoveMapping(locked, GI);
   }
 }
 
@@ -172,18 +186,7 @@ void *ExecutionEngine::updateGlobalMapping(const GlobalValue *GV, void *Addr) {
 
   // Deleting from the mapping?
   if (Addr == 0) {
-    std::map<AssertingVH<const GlobalValue>, void *>::iterator I = Map.find(GV);
-    void *OldVal;
-    if (I == Map.end())
-      OldVal = 0;
-    else {
-      OldVal = I->second;
-      Map.erase(I); 
-    }
-    
-    if (!state.getGlobalAddressReverseMap(locked).empty())
-      state.getGlobalAddressReverseMap(locked).erase(OldVal);
-    return OldVal;
+    return state.RemoveMapping(locked, GV);
   }
   
   void *&CurVal = Map[GV];
diff --git a/unittests/ExecutionEngine/ExecutionEngineTest.cpp b/unittests/ExecutionEngine/ExecutionEngineTest.cpp
index 2106e86..97a8478 100644
--- a/unittests/ExecutionEngine/ExecutionEngineTest.cpp
+++ b/unittests/ExecutionEngine/ExecutionEngineTest.cpp
@@ -93,4 +93,24 @@ TEST_F(ExecutionEngineTest, ReverseGlobalMapping) {
     << " now-free address.";
 }
 
+TEST_F(ExecutionEngineTest, ClearModuleMappings) {
+  GlobalVariable *G1 =
+      NewExtGlobal(Type::getInt32Ty(getGlobalContext()), "Global1");
+
+  int32_t Mem1 = 3;
+  Engine->addGlobalMapping(G1, &Mem1);
+  EXPECT_EQ(G1, Engine->getGlobalValueAtAddress(&Mem1));
+
+  Engine->clearGlobalMappingsFromModule(M);
+
+  EXPECT_EQ(NULL, Engine->getGlobalValueAtAddress(&Mem1));
+
+  GlobalVariable *G2 =
+      NewExtGlobal(Type::getInt32Ty(getGlobalContext()), "Global2");
+  // After clearing the module mappings, we can assign a new GV to the
+  // same address.
+  Engine->addGlobalMapping(G2, &Mem1);
+  EXPECT_EQ(G2, Engine->getGlobalValueAtAddress(&Mem1));
+}
+
 }