[ORC] Fix LazyReexports resource key management.

Multiple reentry points may be associated with a single key.

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h

@@ -230,7 +230,7 @@ class LazyReexportsManager : public ResourceManager {
   EmitTrampolinesFn EmitTrampolines;
   RedirectableSymbolManager &RSMgr;
 
-  DenseMap<ResourceKey, ExecutorAddr> KeyToReentryAddr;
+  DenseMap<ResourceKey, std::vector<ExecutorAddr>> KeyToReentryAddr;
   DenseMap<ExecutorAddr, CallThroughInfo> CallThroughs;
 };
 

llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp

@@ -296,8 +296,11 @@ Error LazyReexportsManager::handleRemoveResources(JITDylib &JD, ResourceKey K) {
   JD.getExecutionSession().runSessionLocked([&]() {
     auto I = KeyToReentryAddr.find(K);
     if (I != KeyToReentryAddr.end()) {
-      auto ReentryAddr = I->second;
-      CallThroughs.erase(ReentryAddr);
+      auto &ReentryAddrs = I->second;
+      for (auto &ReentryAddr : ReentryAddrs) {
+        assert(CallThroughs.count(ReentryAddr) && "CallTrhough missing");
+        CallThroughs.erase(ReentryAddr);
+      }
       KeyToReentryAddr.erase(I);
     }
   });
@@ -309,9 +312,18 @@ void LazyReexportsManager::handleTransferResources(JITDylib &JD,
                                                    ResourceKey SrcK) {
   auto I = KeyToReentryAddr.find(SrcK);
   if (I != KeyToReentryAddr.end()) {
-    auto ReentryAddr = I->second;
-    KeyToReentryAddr.erase(I);
-    KeyToReentryAddr[DstK] = ReentryAddr;
+    auto J = KeyToReentryAddr.find(DstK);
+    if (J == KeyToReentryAddr.end()) {
+      auto Tmp = std::move(I->second);
+      KeyToReentryAddr.erase(I);
+      KeyToReentryAddr[DstK] = std::move(Tmp);
+    } else {
+      auto &SrcReentryAddrs = I->second;
+      auto &DstReentryAddrs = J->second;
+      for (auto &ReentryAddr : SrcReentryAddrs)
+        DstReentryAddrs.push_back(std::move(ReentryAddr));
+      KeyToReentryAddr.erase(I);
+    }
   }
 }
 
@@ -378,7 +390,7 @@ void LazyReexportsManager::emitRedirectableSymbols(
           const auto &ReentryPoint = (*ReentryPoints)[I++];
           CallThroughs[ReentryPoint.getAddress()] = {Name, AI.Aliasee,
                                                      &MR->getTargetJITDylib()};
-          KeyToReentryAddr[K] = ReentryPoint.getAddress();
+          KeyToReentryAddr[K].push_back(ReentryPoint.getAddress());
         }
       })) {
     MR->getExecutionSession().reportError(std::move(Err));