Fix http://llvm.org/PR4822: allow module deletion after a function has been

compiled.

When functions are compiled, they accumulate references in the JITResolver's
stub maps. This patch removes those references when the functions are
destroyed.  It's illegal to destroy a Function when any thread may still try to
call its machine code.

This patch also updates r83987 to use ValueMap instead of explicit CallbackVHs
and fixes a couple "do stuff inside assert()" bugs from r84522.


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

5 files changed, 106 insertions, 56 deletions

diff --git a/include/llvm/ExecutionEngine/ExecutionEngine.h b/include/llvm/ExecutionEngine/ExecutionEngine.h
index b9da0fc..92f552d 100644
--- a/include/llvm/ExecutionEngine/ExecutionEngine.h
+++ b/include/llvm/ExecutionEngine/ExecutionEngine.h
@@ -19,6 +19,7 @@
 #include <map>
 #include <string>
 #include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/ValueMap.h"
 #include "llvm/Support/ValueHandle.h"
 #include "llvm/System/Mutex.h"
 #include "llvm/Target/TargetMachine.h"
@@ -42,26 +43,23 @@ class Type;
 
 class ExecutionEngineState {
 public:
-  class MapUpdatingCVH : public CallbackVH {
-    ExecutionEngineState &EES;
-
-  public:
-    MapUpdatingCVH(ExecutionEngineState &EES, const GlobalValue *GV);
-
-    operator const GlobalValue*() const {
-      return cast<GlobalValue>(getValPtr());
-    }
-
-    virtual void deleted();
-    virtual void allUsesReplacedWith(Value *new_value);
+  struct AddressMapConfig : public ValueMapConfig<const GlobalValue*> {
+    typedef ExecutionEngineState *ExtraData;
+    static sys::Mutex *getMutex(ExecutionEngineState *EES);
+    static void onDelete(ExecutionEngineState *EES, const GlobalValue *Old);
+    static void onRAUW(ExecutionEngineState *, const GlobalValue *,
+                       const GlobalValue *);
   };
 
+  typedef ValueMap<const GlobalValue *, void *, AddressMapConfig>
+      GlobalAddressMapTy;
+
 private:
   ExecutionEngine &EE;
 
   /// GlobalAddressMap - A mapping between LLVM global values and their
   /// actualized version...
-  std::map<MapUpdatingCVH, void *> GlobalAddressMap;
+  GlobalAddressMapTy GlobalAddressMap;
 
   /// GlobalAddressReverseMap - This is the reverse mapping of GlobalAddressMap,
   /// used to convert raw addresses into the LLVM global value that is emitted
@@ -70,13 +68,9 @@ private:
   std::map<void *, AssertingVH<const GlobalValue> > GlobalAddressReverseMap;
 
 public:
-  ExecutionEngineState(ExecutionEngine &EE) : EE(EE) {}
+  ExecutionEngineState(ExecutionEngine &EE);
 
-  MapUpdatingCVH getVH(const GlobalValue *GV) {
-    return MapUpdatingCVH(*this, GV);
-  }
-
-  std::map<MapUpdatingCVH, void *> &
+  GlobalAddressMapTy &
   getGlobalAddressMap(const MutexGuard &) {
     return GlobalAddressMap;
   }
@@ -485,15 +479,8 @@ class EngineBuilder {
   }
 
   ExecutionEngine *create();
-
 };
 
-inline bool operator<(const ExecutionEngineState::MapUpdatingCVH& lhs,
-                      const ExecutionEngineState::MapUpdatingCVH& rhs) {
-    return static_cast<const GlobalValue*>(lhs) <
-        static_cast<const GlobalValue*>(rhs);
-}
-
 } // End llvm namespace
 
 #endif
diff --git a/lib/ExecutionEngine/ExecutionEngine.cpp b/lib/ExecutionEngine/ExecutionEngine.cpp
index 053d960..16808a7 100644
--- a/lib/ExecutionEngine/ExecutionEngine.cpp
+++ b/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -117,8 +117,7 @@ Function *ExecutionEngine::FindFunctionNamed(const char *FnName) {
 
 void *ExecutionEngineState::RemoveMapping(
   const MutexGuard &, const GlobalValue *ToUnmap) {
-  std::map<MapUpdatingCVH, void *>::iterator I =
-    GlobalAddressMap.find(getVH(ToUnmap));
+  GlobalAddressMapTy::iterator I = GlobalAddressMap.find(ToUnmap);
   void *OldVal;
   if (I == GlobalAddressMap.end())
     OldVal = 0;
@@ -141,7 +140,7 @@ void ExecutionEngine::addGlobalMapping(const GlobalValue *GV, void *Addr) {
 
   DEBUG(errs() << "JIT: Map \'" << GV->getName() 
         << "\' to [" << Addr << "]\n";);
-  void *&CurVal = EEState.getGlobalAddressMap(locked)[EEState.getVH(GV)];
+  void *&CurVal = EEState.getGlobalAddressMap(locked)[GV];
   assert((CurVal == 0 || Addr == 0) && "GlobalMapping already established!");
   CurVal = Addr;
   
@@ -183,7 +182,7 @@ void ExecutionEngine::clearGlobalMappingsFromModule(Module *M) {
 void *ExecutionEngine::updateGlobalMapping(const GlobalValue *GV, void *Addr) {
   MutexGuard locked(lock);
 
-  std::map<ExecutionEngineState::MapUpdatingCVH, void *> &Map =
+  ExecutionEngineState::GlobalAddressMapTy &Map =
     EEState.getGlobalAddressMap(locked);
 
   // Deleting from the mapping?
@@ -191,7 +190,7 @@ void *ExecutionEngine::updateGlobalMapping(const GlobalValue *GV, void *Addr) {
     return EEState.RemoveMapping(locked, GV);
   }
   
-  void *&CurVal = Map[EEState.getVH(GV)];
+  void *&CurVal = Map[GV];
   void *OldVal = CurVal;
 
   if (CurVal && !EEState.getGlobalAddressReverseMap(locked).empty())
@@ -214,8 +213,8 @@ void *ExecutionEngine::updateGlobalMapping(const GlobalValue *GV, void *Addr) {
 void *ExecutionEngine::getPointerToGlobalIfAvailable(const GlobalValue *GV) {
   MutexGuard locked(lock);
   
-  std::map<ExecutionEngineState::MapUpdatingCVH, void*>::iterator I =
-    EEState.getGlobalAddressMap(locked).find(EEState.getVH(GV));
+  ExecutionEngineState::GlobalAddressMapTy::iterator I =
+    EEState.getGlobalAddressMap(locked).find(GV);
   return I != EEState.getGlobalAddressMap(locked).end() ? I->second : 0;
 }
 
@@ -227,7 +226,7 @@ const GlobalValue *ExecutionEngine::getGlobalValueAtAddress(void *Addr) {
 
   // If we haven't computed the reverse mapping yet, do so first.
   if (EEState.getGlobalAddressReverseMap(locked).empty()) {
-    for (std::map<ExecutionEngineState::MapUpdatingCVH, void *>::iterator
+    for (ExecutionEngineState::GlobalAddressMapTy::iterator
          I = EEState.getGlobalAddressMap(locked).begin(),
          E = EEState.getGlobalAddressMap(locked).end(); I != E; ++I)
       EEState.getGlobalAddressReverseMap(locked).insert(std::make_pair(I->second,
@@ -476,7 +475,7 @@ void *ExecutionEngine::getPointerToGlobal(const GlobalValue *GV) {
     return getPointerToFunction(F);
 
   MutexGuard locked(lock);
-  void *p = EEState.getGlobalAddressMap(locked)[EEState.getVH(GV)];
+  void *p = EEState.getGlobalAddressMap(locked)[GV];
   if (p)
     return p;
 
@@ -486,7 +485,7 @@ void *ExecutionEngine::getPointerToGlobal(const GlobalValue *GV) {
     EmitGlobalVariable(GVar);
   else
     llvm_unreachable("Global hasn't had an address allocated yet!");
-  return EEState.getGlobalAddressMap(locked)[EEState.getVH(GV)];
+  return EEState.getGlobalAddressMap(locked)[GV];
 }
 
 /// This function converts a Constant* into a GenericValue. The interesting 
@@ -1072,17 +1071,22 @@ void ExecutionEngine::EmitGlobalVariable(const GlobalVariable *GV) {
   ++NumGlobals;
 }
 
-ExecutionEngineState::MapUpdatingCVH::MapUpdatingCVH(
-  ExecutionEngineState &EES, const GlobalValue *GV)
-  : CallbackVH(const_cast<GlobalValue*>(GV)), EES(EES) {}
+ExecutionEngineState::ExecutionEngineState(ExecutionEngine &EE)
+  : EE(EE), GlobalAddressMap(this) {
+}
 
-void ExecutionEngineState::MapUpdatingCVH::deleted() {
-  MutexGuard locked(EES.EE.lock);
-  EES.RemoveMapping(locked, *this);  // Destroys *this.
+sys::Mutex *ExecutionEngineState::AddressMapConfig::getMutex(
+  ExecutionEngineState *EES) {
+  return &EES->EE.lock;
+}
+void ExecutionEngineState::AddressMapConfig::onDelete(
+  ExecutionEngineState *EES, const GlobalValue *Old) {
+  void *OldVal = EES->GlobalAddressMap.lookup(Old);
+  EES->GlobalAddressReverseMap.erase(OldVal);
 }
 
-void ExecutionEngineState::MapUpdatingCVH::allUsesReplacedWith(
-  Value *new_value) {
+void ExecutionEngineState::AddressMapConfig::onRAUW(
+  ExecutionEngineState *, const GlobalValue *, const GlobalValue *) {
   assert(false && "The ExecutionEngine doesn't know how to handle a"
          " RAUW on a value it has a global mapping for.");
 }
diff --git a/lib/ExecutionEngine/JIT/JITEmitter.cpp b/lib/ExecutionEngine/JIT/JITEmitter.cpp
index 073d6fb..2915f49 100644
--- a/lib/ExecutionEngine/JIT/JITEmitter.cpp
+++ b/lib/ExecutionEngine/JIT/JITEmitter.cpp
@@ -46,6 +46,7 @@
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
+#include "llvm/ADT/ValueMap.h"
 #include <algorithm>
 #ifndef NDEBUG
 #include <iomanip>
@@ -62,12 +63,29 @@ static JIT *TheJIT = 0;
 // JIT lazy compilation code.
 //
 namespace {
+  class JITResolverState;
+
+  template<typename ValueTy>
+  struct NoRAUWValueMapConfig : public ValueMapConfig<ValueTy> {
+    typedef JITResolverState *ExtraData;
+    static void onRAUW(JITResolverState *, Value *Old, Value *New) {
+      assert(false && "The JIT doesn't know how to handle a"
+             " RAUW on a value it has emitted.");
+    }
+  };
+
+  struct CallSiteValueMapConfig : public NoRAUWValueMapConfig<Function*> {
+    typedef JITResolverState *ExtraData;
+    static void onDelete(JITResolverState *JRS, Function *F);
+  };
+
   class JITResolverState {
   public:
-    typedef DenseMap<AssertingVH<Function>, void*> FunctionToStubMapTy;
+    typedef ValueMap<Function*, void*, NoRAUWValueMapConfig<Function*> >
+      FunctionToStubMapTy;
     typedef std::map<void*, AssertingVH<Function> > CallSiteToFunctionMapTy;
-    typedef DenseMap<AssertingVH<Function>, SmallPtrSet<void*, 1> >
-            FunctionToCallSitesMapTy;
+    typedef ValueMap<Function *, SmallPtrSet<void*, 1>,
+                     CallSiteValueMapConfig> FunctionToCallSitesMapTy;
     typedef std::map<AssertingVH<GlobalValue>, void*> GlobalToIndirectSymMapTy;
   private:
     /// FunctionToStubMap - Keep track of the stub created for a particular
@@ -84,6 +102,9 @@ namespace {
     GlobalToIndirectSymMapTy GlobalToIndirectSymMap;
 
   public:
+    JITResolverState() : FunctionToStubMap(this),
+                         FunctionToCallSitesMap(this) {}
+
     FunctionToStubMapTy& getFunctionToStubMap(const MutexGuard& locked) {
       assert(locked.holds(TheJIT->lock));
       return FunctionToStubMap;
@@ -111,8 +132,10 @@ namespace {
     void AddCallSite(const MutexGuard &locked, void *CallSite, Function *F) {
       assert(locked.holds(TheJIT->lock));
 
-      assert(CallSiteToFunctionMap.insert(std::make_pair(CallSite, F)).second &&
-             "Pair was already in CallSiteToFunctionMap");
+      bool Inserted = CallSiteToFunctionMap.insert(
+          std::make_pair(CallSite, F)).second;
+      (void)Inserted;
+      assert(Inserted && "Pair was already in CallSiteToFunctionMap");
       FunctionToCallSitesMap[F].insert(CallSite);
     }
 
@@ -142,8 +165,9 @@ namespace {
       FunctionToCallSitesMapTy::iterator F2C_I = FunctionToCallSitesMap.find(F);
       assert(F2C_I != FunctionToCallSitesMap.end() &&
              "FunctionToCallSitesMap broken");
-      assert(F2C_I->second.erase(Stub) &&
-             "FunctionToCallSitesMap broken");
+      bool Erased = F2C_I->second.erase(Stub);
+      (void)Erased;
+      assert(Erased && "FunctionToCallSitesMap broken");
       if (F2C_I->second.empty())
         FunctionToCallSitesMap.erase(F2C_I);
 
@@ -152,13 +176,17 @@ namespace {
 
     void EraseAllCallSites(const MutexGuard &locked, Function *F) {
       assert(locked.holds(TheJIT->lock));
+      EraseAllCallSitesPrelocked(F);
+    }
+    void EraseAllCallSitesPrelocked(Function *F) {
       FunctionToCallSitesMapTy::iterator F2C = FunctionToCallSitesMap.find(F);
       if (F2C == FunctionToCallSitesMap.end())
         return;
       for (SmallPtrSet<void*, 1>::const_iterator I = F2C->second.begin(),
              E = F2C->second.end(); I != E; ++I) {
-        assert(CallSiteToFunctionMap.erase(*I) == 1 &&
-               "Missing call site->function mapping");
+        bool Erased = CallSiteToFunctionMap.erase(*I);
+        (void)Erased;
+        assert(Erased && "Missing call site->function mapping");
       }
       FunctionToCallSitesMap.erase(F2C);
     }
@@ -245,6 +273,10 @@ namespace {
 
 JITResolver *JITResolver::TheJITResolver = 0;
 
+void CallSiteValueMapConfig::onDelete(JITResolverState *JRS, Function *F) {
+  JRS->EraseAllCallSitesPrelocked(F);
+}
+
 /// getFunctionStubIfAvailable - This returns a pointer to a function stub
 /// if it has already been created.
 void *JITResolver::getFunctionStubIfAvailable(Function *F) {
diff --git a/unittests/ExecutionEngine/JIT/JITTest.cpp b/unittests/ExecutionEngine/JIT/JITTest.cpp
index 8f9b65a..e47a437 100644
--- a/unittests/ExecutionEngine/JIT/JITTest.cpp
+++ b/unittests/ExecutionEngine/JIT/JITTest.cpp
@@ -9,6 +9,7 @@
 
 #include "gtest/gtest.h"
 #include "llvm/ADT/OwningPtr.h"
+#include "llvm/Assembly/Parser.h"
 #include "llvm/BasicBlock.h"
 #include "llvm/Constant.h"
 #include "llvm/Constants.h"
@@ -22,6 +23,7 @@
 #include "llvm/Module.h"
 #include "llvm/ModuleProvider.h"
 #include "llvm/Support/IRBuilder.h"
+#include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/TypeBuilder.h"
 #include "llvm/Target/TargetSelect.h"
 #include "llvm/Type.h"
@@ -49,14 +51,25 @@ class JITTest : public testing::Test {
  protected:
   virtual void SetUp() {
     M = new Module("<main>", Context);
+    MP = new ExistingModuleProvider(M);
     std::string Error;
-    TheJIT.reset(EngineBuilder(M).setEngineKind(EngineKind::JIT)
+    TheJIT.reset(EngineBuilder(MP).setEngineKind(EngineKind::JIT)
                  .setErrorStr(&Error).create());
     ASSERT_TRUE(TheJIT.get() != NULL) << Error;
   }
 
+  void LoadAssembly(const char *assembly) {
+    SMDiagnostic Error;
+    bool success = NULL != ParseAssemblyString(assembly, M, Error, Context);
+    std::string errMsg;
+    raw_string_ostream os(errMsg);
+    Error.Print("", os);
+    ASSERT_TRUE(success) << os.str();
+  }
+
   LLVMContext Context;
-  Module *M;  // Owned by ExecutionEngine.
+  Module *M;  // Owned by MP.
+  ModuleProvider *MP;  // Owned by ExecutionEngine.
   OwningPtr<ExecutionEngine> TheJIT;
 };
 
@@ -264,6 +277,20 @@ TEST_F(JITTest, NonLazyLeaksNoStubs) {
 }
 #endif
 
+TEST_F(JITTest, ModuleDeletion) {
+  LoadAssembly("define void @main() { "
+               "  call i32 @computeVal() "
+               "  ret void "
+               "} "
+               " "
+               "define internal i32 @computeVal()  { "
+               "  ret i32 0 "
+               "} ");
+  Function *func = M->getFunction("main");
+  TheJIT->getPointerToFunction(func);
+  TheJIT->deleteModuleProvider(MP);
+}
+
 // This code is copied from JITEventListenerTest, but it only runs once for all
 // the tests in this directory.  Everything seems fine, but that's strange
 // behavior.
diff --git a/unittests/ExecutionEngine/JIT/Makefile b/unittests/ExecutionEngine/JIT/Makefile
index 0069c76..048924a 100644
--- a/unittests/ExecutionEngine/JIT/Makefile
+++ b/unittests/ExecutionEngine/JIT/Makefile
@@ -9,7 +9,7 @@
 
 LEVEL = ../../..
 TESTNAME = JIT
-LINK_COMPONENTS := core support jit native
+LINK_COMPONENTS := asmparser core support jit native
 
 include $(LEVEL)/Makefile.config
 include $(LLVM_SRC_ROOT)/unittests/Makefile.unittest