[ORC] Add JITDylib argument to ResourceManager notify-removing/transferring ops.

In some cases it's helpful to group trackers by JITDylib. E.g. Platform classes
may want to track initializer symbols with a `JITDylib -> Tracker -> [ Symbol ]`
map. This makes it easy to collect all symbols for the JITDylib, while still
allowing efficient removal of a single tracker. Passing the JITDylib as an
argument to ResourceManager::notifyRemovingResources and
ResourceManager::notifyTransferringResources supports such use-cases.

5 files changed, 23 insertions, 15 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index d058da2..4a9d0d4 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -2285,9 +2285,10 @@ Error ExecutionSession::removeResourceTracker(ResourceTracker &RT) {
 
   Error Err = Error::success();
 
+  auto &JD = RT.getJITDylib();
   for (auto *L : reverse(CurrentResourceManagers))
-    Err =
-        joinErrors(std::move(Err), L->handleRemoveResources(RT.getKeyUnsafe()));
+    Err = joinErrors(std::move(Err),
+                     L->handleRemoveResources(JD, RT.getKeyUnsafe()));
 
   for (auto &Q : QueriesToFail)
     Q->handleFailed(
@@ -2316,7 +2317,8 @@ void ExecutionSession::transferResourceTracker(ResourceTracker &DstRT,
     auto &JD = DstRT.getJITDylib();
     JD.transferTracker(DstRT, SrcRT);
     for (auto *L : reverse(ResourceManagers))
-      L->handleTransferResources(DstRT.getKeyUnsafe(), SrcRT.getKeyUnsafe());
+      L->handleTransferResources(JD, DstRT.getKeyUnsafe(),
+                                 SrcRT.getKeyUnsafe());
   });
 }
 
diff --git a/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp b/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
index 418dc39..02c3e61 100644
--- a/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/DebugObjectManagerPlugin.cpp
@@ -492,7 +492,8 @@ Error DebugObjectManagerPlugin::notifyFailed(
   return Error::success();
 }
 
-void DebugObjectManagerPlugin::notifyTransferringResources(ResourceKey DstKey,
+void DebugObjectManagerPlugin::notifyTransferringResources(JITDylib &JD,
+                                                           ResourceKey DstKey,
                                                            ResourceKey SrcKey) {
   // Debug objects are stored by ResourceKey only after registration.
   // Thus, pending objects don't need to be updated here.
@@ -507,7 +508,8 @@ void DebugObjectManagerPlugin::notifyTransferringResources(ResourceKey DstKey,
   }
 }
 
-Error DebugObjectManagerPlugin::notifyRemovingResources(ResourceKey Key) {
+Error DebugObjectManagerPlugin::notifyRemovingResources(JITDylib &JD,
+                                                        ResourceKey Key) {
   // Removing the resource for a pending object fails materialization, so they
   // get cleaned up in the notifyFailed() handler.
   std::lock_guard<std::mutex> Lock(RegisteredObjsLock);
diff --git a/llvm/lib/ExecutionEngine/Orc/DebuggerSupportPlugin.cpp b/llvm/lib/ExecutionEngine/Orc/DebuggerSupportPlugin.cpp
index 009e2d7..15e7ffb 100644
--- a/llvm/lib/ExecutionEngine/Orc/DebuggerSupportPlugin.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/DebuggerSupportPlugin.cpp
@@ -390,12 +390,12 @@ Error GDBJITDebugInfoRegistrationPlugin::notifyFailed(
 }
 
 Error GDBJITDebugInfoRegistrationPlugin::notifyRemovingResources(
-    ResourceKey K) {
+    JITDylib &JD, ResourceKey K) {
   return Error::success();
 }
 
 void GDBJITDebugInfoRegistrationPlugin::notifyTransferringResources(
-    ResourceKey DstKey, ResourceKey SrcKey) {}
+    JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) {}
 
 void GDBJITDebugInfoRegistrationPlugin::modifyPassConfig(
     MaterializationResponsibility &MR, LinkGraph &LG,
diff --git a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
index 1a4a451..2b11c47 100644
--- a/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp
@@ -680,12 +680,12 @@ Error ObjectLinkingLayer::notifyEmitted(MaterializationResponsibility &MR,
       [&](ResourceKey K) { Allocs[K].push_back(std::move(FA)); });
 }
 
-Error ObjectLinkingLayer::handleRemoveResources(ResourceKey K) {
+Error ObjectLinkingLayer::handleRemoveResources(JITDylib &JD, ResourceKey K) {
 
   {
     Error Err = Error::success();
     for (auto &P : Plugins)
-      Err = joinErrors(std::move(Err), P->notifyRemovingResources(K));
+      Err = joinErrors(std::move(Err), P->notifyRemovingResources(JD, K));
     if (Err)
       return Err;
   }
@@ -705,7 +705,8 @@ Error ObjectLinkingLayer::handleRemoveResources(ResourceKey K) {
   return MemMgr.deallocate(std::move(AllocsToRemove));
 }
 
-void ObjectLinkingLayer::handleTransferResources(ResourceKey DstKey,
+void ObjectLinkingLayer::handleTransferResources(JITDylib &JD,
+                                                 ResourceKey DstKey,
                                                  ResourceKey SrcKey) {
   auto I = Allocs.find(SrcKey);
   if (I != Allocs.end()) {
@@ -721,7 +722,7 @@ void ObjectLinkingLayer::handleTransferResources(ResourceKey DstKey,
   }
 
   for (auto &P : Plugins)
-    P->notifyTransferringResources(DstKey, SrcKey);
+    P->notifyTransferringResources(JD, DstKey, SrcKey);
 }
 
 EHFrameRegistrationPlugin::EHFrameRegistrationPlugin(
@@ -773,7 +774,8 @@ Error EHFrameRegistrationPlugin::notifyFailed(
   return Error::success();
 }
 
-Error EHFrameRegistrationPlugin::notifyRemovingResources(ResourceKey K) {
+Error EHFrameRegistrationPlugin::notifyRemovingResources(JITDylib &JD,
+                                                         ResourceKey K) {
   std::vector<ExecutorAddrRange> RangesToRemove;
 
   ES.runSessionLocked([&] {
@@ -797,7 +799,7 @@ Error EHFrameRegistrationPlugin::notifyRemovingResources(ResourceKey K) {
 }
 
 void EHFrameRegistrationPlugin::notifyTransferringResources(
-    ResourceKey DstKey, ResourceKey SrcKey) {
+    JITDylib &JD, ResourceKey DstKey, ResourceKey SrcKey) {
   auto SI = EHFrameRanges.find(SrcKey);
   if (SI == EHFrameRanges.end())
     return;
diff --git a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
index 8e9d4cd..9e781e9 100644
--- a/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp
@@ -350,7 +350,8 @@ void RTDyldObjectLinkingLayer::onObjEmit(
   }
 }
 
-Error RTDyldObjectLinkingLayer::handleRemoveResources(ResourceKey K) {
+Error RTDyldObjectLinkingLayer::handleRemoveResources(JITDylib &JD,
+                                                      ResourceKey K) {
 
   std::vector<MemoryManagerUP> MemMgrsToRemove;
 
@@ -374,7 +375,8 @@ Error RTDyldObjectLinkingLayer::handleRemoveResources(ResourceKey K) {
   return Error::success();
 }
 
-void RTDyldObjectLinkingLayer::handleTransferResources(ResourceKey DstKey,
+void RTDyldObjectLinkingLayer::handleTransferResources(JITDylib &JD,
+                                                       ResourceKey DstKey,
                                                        ResourceKey SrcKey) {
   auto I = MemMgrs.find(SrcKey);
   if (I != MemMgrs.end()) {