[llvm] 24c2744 - [ORC] Fix LazyReexports resource key management.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2024-12-17T22:38:15+11:00
New Revision: 24c2744a189eef9dfd39789df4983e4ffd219197

URL: https://github.com/llvm/llvm-project/commit/24c2744a189eef9dfd39789df4983e4ffd219197
DIFF: https://github.com/llvm/llvm-project/commit/24c2744a189eef9dfd39789df4983e4ffd219197.diff

LOG: [ORC] Fix LazyReexports resource key management.

Multiple reentry points may be associated with a single key.

Added: 
    

Modified: 
    llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h
    llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h b/llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h
index cc9c664d0e7c09..20431729efc362 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/LazyReexports.h
+++ b/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;
 };
 

diff  --git a/llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp b/llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp
index 6e1e3746bfa24e..e82aa44dfcd05d 100644
--- a/llvm/lib/ExecutionEngine/Orc/LazyReexports.cpp
+++ b/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));