From 380355cb4c6ba28cb8b9a573591dcce694a047ee Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames@gmail.com>
Date: Tue, 31 Jan 2023 17:55:29 -0800
Subject: [PATCH] [ORC] Fix an iterator invalidation issue in
 JITDylib::defineMaterializing.

The loop body may add and remove entries in the symbol table so we can't hold
iterators to the entries. This commit updates the method to use the newly added
NonOwningSymbolStringPtr type as keys for removal instead.

Side note: This bug has been present since the introduction of the
defineMaterializing method, but the method is called rarely and DenseMap
resizes are also rare so we didn't see any fallout until a large program was
thrown at it. There's no testcase as I haven't been able to reproduce the
failure with smaller testcases.
---
 llvm/lib/ExecutionEngine/Orc/Core.cpp | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index 4a9d0d4..f857191 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -689,10 +689,11 @@ void JITDylib::removeGenerator(DefinitionGenerator &G) {
 
 Expected<SymbolFlagsMap>
 JITDylib::defineMaterializing(SymbolFlagsMap SymbolFlags) {
+  // TODO: Should we bail out early here if MR is defunct?
 
   return ES.runSessionLocked([&]() -> Expected<SymbolFlagsMap> {
-    std::vector<SymbolTable::iterator> AddedSyms;
-    std::vector<SymbolFlagsMap::iterator> RejectedWeakDefs;
+    std::vector<NonOwningSymbolStringPtr> AddedSyms;
+    std::vector<NonOwningSymbolStringPtr> RejectedWeakDefs;
 
     for (auto SFItr = SymbolFlags.begin(), SFEnd = SymbolFlags.end();
          SFItr != SFEnd; ++SFItr) {
@@ -708,27 +709,27 @@ JITDylib::defineMaterializing(SymbolFlagsMap SymbolFlags) {
         // If this is a strong definition then error out.
         if (!Flags.isWeak()) {
           // Remove any symbols already added.
-          for (auto &SI : AddedSyms)
-            Symbols.erase(SI);
+          for (auto &S : AddedSyms)
+            Symbols.erase(Symbols.find_as(S));
 
           // FIXME: Return all duplicates.
           return make_error<DuplicateDefinition>(std::string(*Name));
         }
 
         // Otherwise just make a note to discard this symbol after the loop.
-        RejectedWeakDefs.push_back(SFItr);
+        RejectedWeakDefs.push_back(NonOwningSymbolStringPtr(Name));
         continue;
       } else
         EntryItr =
           Symbols.insert(std::make_pair(Name, SymbolTableEntry(Flags))).first;
 
-      AddedSyms.push_back(EntryItr);
+      AddedSyms.push_back(NonOwningSymbolStringPtr(Name));
       EntryItr->second.setState(SymbolState::Materializing);
     }
 
     // Remove any rejected weak definitions from the SymbolFlags map.
     while (!RejectedWeakDefs.empty()) {
-      SymbolFlags.erase(RejectedWeakDefs.back());
+      SymbolFlags.erase(SymbolFlags.find_as(RejectedWeakDefs.back()));
       RejectedWeakDefs.pop_back();
     }
 
-- 
2.7.4