[llvm] bccd1b9 - [StringMap] Invalidate iterators in remove() (#202520)

[via llvm-commits]

Author: Fangrui Song
Date: 2026-06-09T08:12:33Z
New Revision: bccd1b9cb744e5dd96ee59baa4bf4583457feea3

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

LOG: [StringMap] Invalidate iterators in remove() (#202520)

erase() bumps the epoch to invalidate iterators (#202237), but the
lower-level remove() — which detaches an entry without destroying it,
used
by ValueSymbolTable via Value::setName() — did not. Move the
incrementEpoch() into remove() so remove-while-iterating fails fast
under
LLVM_ENABLE_ABI_BREAKING_CHECKS too.

Aided by Claude Opus 4.8

Added: 
    

Modified: 
    llvm/include/llvm/ADT/StringMap.h
    llvm/lib/Transforms/IPO/StripSymbols.cpp
    llvm/unittests/ADT/StringMapTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ADT/StringMap.h b/llvm/include/llvm/ADT/StringMap.h
index 5ce35f5c7a6ce..77853f240fc10 100644
--- a/llvm/include/llvm/ADT/StringMap.h
+++ b/llvm/include/llvm/ADT/StringMap.h
@@ -430,11 +430,13 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
 
   /// remove - Remove the specified key/value pair from the map, but do not
   /// erase it.  This aborts if the key is not in the map.
-  void remove(MapEntryTy *KeyValue) { RemoveKey(KeyValue); }
+  void remove(MapEntryTy *KeyValue) {
+    incrementEpoch();
+    RemoveKey(KeyValue);
+  }
 
   void erase(iterator I) {
     MapEntryTy &V = *I;
-    incrementEpoch();
     remove(&V);
     V.Destroy(getAllocator());
   }

diff  --git a/llvm/lib/Transforms/IPO/StripSymbols.cpp b/llvm/lib/Transforms/IPO/StripSymbols.cpp
index ec701b6d0037b..5e00ef468f9ef 100644
--- a/llvm/lib/Transforms/IPO/StripSymbols.cpp
+++ b/llvm/lib/Transforms/IPO/StripSymbols.cpp
@@ -22,6 +22,7 @@
 #include "llvm/Transforms/IPO/StripSymbols.h"
 
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DebugInfo.h"
 #include "llvm/IR/DerivedTypes.h"
@@ -76,15 +77,17 @@ static void RemoveDeadConstant(Constant *C) {
 // Strip the symbol table of its names.
 //
 static void StripSymtab(ValueSymbolTable &ST, bool PreserveDbgInfo) {
-  for (ValueSymbolTable::iterator VI = ST.begin(), VE = ST.end(); VI != VE; ) {
-    Value *V = VI->getValue();
-    ++VI;
-    if (!isa<GlobalValue>(V) || cast<GlobalValue>(V)->hasLocalLinkage()) {
+  // Collect the values to rename first: setName("") removes the value from the
+  // symbol table, which invalidates iterators into it.
+  SmallVector<Value *, 0> ToStrip;
+  for (const ValueName &VN : ST) {
+    Value *V = VN.getValue();
+    if (!isa<GlobalValue>(V) || cast<GlobalValue>(V)->hasLocalLinkage())
       if (!PreserveDbgInfo || !V->getName().starts_with("llvm.dbg"))
-        // Set name to "", removing from symbol table!
-        V->setName("");
-    }
+        ToStrip.push_back(V);
   }
+  for (Value *V : ToStrip)
+    V->setName("");
 }
 
 // Strip any named types of their names.

diff  --git a/llvm/unittests/ADT/StringMapTest.cpp b/llvm/unittests/ADT/StringMapTest.cpp
index d350b46525ba0..6b357d20504e1 100644
--- a/llvm/unittests/ADT/StringMapTest.cpp
+++ b/llvm/unittests/ADT/StringMapTest.cpp
@@ -785,6 +785,17 @@ TEST(StringMapCustomTest, EraseInvalidatesIterators) {
   EXPECT_DEATH((void)It->second, "invalid iterator access");
 }
 
+TEST(StringMapCustomTest, RemoveInvalidatesIterators) {
+  StringMap<int> Map;
+  Map["a"] = 1;
+  Map["b"] = 2;
+  auto It = Map.find("a");
+  auto *Entry = &*Map.find("b");
+  Map.remove(Entry);
+  Entry->Destroy(Map.getAllocator());
+  EXPECT_DEATH((void)It->second, "invalid iterator access");
+}
+
 TEST(StringMapCustomTest, ClearInvalidatesIterators) {
   StringMap<int> Map;
   Map["a"] = 1;