[llvm] 380355c - [ORC] Fix an iterator invalidation issue in JITDylib::defineMaterializing.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2023-01-31T18:11:21-08:00
New Revision: 380355cb4c6ba28cb8b9a573591dcce694a047ee

URL: https://github.com/llvm/llvm-project/commit/380355cb4c6ba28cb8b9a573591dcce694a047ee
DIFF: https://github.com/llvm/llvm-project/commit/380355cb4c6ba28cb8b9a573591dcce694a047ee.diff

LOG: [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.

Added: 
    

Modified: 
    llvm/lib/ExecutionEngine/Orc/Core.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index 4a9d0d470a8ec..f85719195532d 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();
     }