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

[Fangrui Song via llvm-commits]

https://github.com/MaskRay created https://github.com/llvm/llvm-project/pull/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

>From b339d247c74e7f2ed11152d47727d7cbe9724703 Mon Sep 17 00:00:00 2001
From: Fangrui Song <i at maskray.me>
Date: Mon, 8 Jun 2026 23:38:00 -0700
Subject: [PATCH] [StringMap] Invalidate iterators in remove()
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

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
---
 llvm/include/llvm/ADT/StringMap.h        |  6 ++++--
 llvm/lib/Transforms/IPO/StripSymbols.cpp | 17 ++++++++++-------
 llvm/unittests/ADT/StringMapTest.cpp     | 11 +++++++++++
 3 files changed, 25 insertions(+), 9 deletions(-)

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;