[ORC] Fix transfer to unknown ResourceTrackers

llvm/lib/ExecutionEngine/Orc/ObjectLinkingLayer.cpp

@@ -701,16 +701,15 @@ Error ObjectLinkingLayer::handleRemoveResources(JITDylib &JD, ResourceKey K) {
 void ObjectLinkingLayer::handleTransferResources(JITDylib &JD,
                                                  ResourceKey DstKey,
                                                  ResourceKey SrcKey) {
-  auto I = Allocs.find(SrcKey);
-  if (I != Allocs.end()) {
-    auto &SrcAllocs = I->second;
+  if (Allocs.contains(SrcKey)) {
+    // DstKey may not be in the DenseMap yet, so the following line may resize
+    // the container and invalidate iterators and value references.
     auto &DstAllocs = Allocs[DstKey];
+    auto &SrcAllocs = Allocs[SrcKey];
     DstAllocs.reserve(DstAllocs.size() + SrcAllocs.size());
     for (auto &Alloc : SrcAllocs)
       DstAllocs.push_back(std::move(Alloc));
 
-    // Erase SrcKey entry using value rather than iterator I: I may have been
-    // invalidated when we looked up DstKey.
     Allocs.erase(SrcKey);
   }
 

llvm/lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp

@@ -430,16 +430,15 @@ Error RTDyldObjectLinkingLayer::handleRemoveResources(JITDylib &JD,
 void RTDyldObjectLinkingLayer::handleTransferResources(JITDylib &JD,
                                                        ResourceKey DstKey,
                                                        ResourceKey SrcKey) {
-  auto I = MemMgrs.find(SrcKey);
-  if (I != MemMgrs.end()) {
-    auto &SrcMemMgrs = I->second;
+  if (MemMgrs.contains(SrcKey)) {
+    // DstKey may not be in the DenseMap yet, so the following line may resize
+    // the container and invalidate iterators and value references.
     auto &DstMemMgrs = MemMgrs[DstKey];
+    auto &SrcMemMgrs = MemMgrs[SrcKey];
     DstMemMgrs.reserve(DstMemMgrs.size() + SrcMemMgrs.size());
     for (auto &MemMgr : SrcMemMgrs)
       DstMemMgrs.push_back(std::move(MemMgr));
 
-    // Erase SrcKey entry using value rather than iterator I: I may have been
-    // invalidated when we looked up DstKey.
     MemMgrs.erase(SrcKey);
   }
 }

llvm/unittests/ExecutionEngine/Orc/ObjectLinkingLayerTest.cpp

@@ -65,6 +65,36 @@ TEST_F(ObjectLinkingLayerTest, AddLinkGraph) {
   EXPECT_THAT_EXPECTED(ES.lookup(&JD, "_X"), Succeeded());
 }
 
+TEST_F(ObjectLinkingLayerTest, ResourceTracker) {
+  // This test transfers allocations to previously unknown ResourceTrackers,
+  // while increasing the number of trackers in the ObjectLinkingLayer, which
+  // may invalidate some iterators internally.
+  std::vector<ResourceTrackerSP> Trackers;
+  for (unsigned I = 0; I < 64; I++) {
+    auto G = std::make_unique<LinkGraph>("foo", Triple("x86_64-apple-darwin"),
+                                         8, llvm::endianness::little,
+                                         x86_64::getEdgeKindName);
+
+    auto &Sec1 = G->createSection("__data", MemProt::Read | MemProt::Write);
+    auto &B1 = G->createContentBlock(Sec1, BlockContent,
+                                     orc::ExecutorAddr(0x1000), 8, 0);
+    llvm::SmallString<0> SymbolName;
+    SymbolName += "_X";
+    SymbolName += std::to_string(I);
+    G->addDefinedSymbol(B1, 4, SymbolName, 4, Linkage::Strong, Scope::Default,
+                        false, false);
+
+    auto RT1 = JD.createResourceTracker();
+    EXPECT_THAT_ERROR(ObjLinkingLayer.add(RT1, std::move(G)), Succeeded());
+    EXPECT_THAT_EXPECTED(ES.lookup(&JD, SymbolName), Succeeded());
+
+    auto RT2 = JD.createResourceTracker();
+    RT1->transferTo(*RT2);
+
+    Trackers.push_back(RT2);
+  }
+}
+
 TEST_F(ObjectLinkingLayerTest, ClaimLateDefinedWeakSymbols) {
   // Check that claiming weak symbols works as expected.
   //