[llvm] [StringMap] Invalidate iterators on mutation (PR #202237)

Fangrui Song via llvm-commits llvm-commits at lists.llvm.org
Mon Jun 8 09:14:20 PDT 2026


https://github.com/MaskRay updated https://github.com/llvm/llvm-project/pull/202237

>From 5331ffb837663685a30f797b653d64fc462d273b Mon Sep 17 00:00:00 2001
From: Fangrui Song <i at maskray.me>
Date: Sun, 7 Jun 2026 14:51:52 -0700
Subject: [PATCH 1/3] [StringMap] Invalidate iterators on erase

Tighten StringMap's `erase` contract so that, like DenseMap (#199369),
it invalidates iterators obtained before the call. StringMap now derives
from DebugEpochBase and its iterators from HandleBase, so an
erase-while-iterating bug fails under LLVM_ENABLE_ABI_BREAKING_CHECKS.

Add a `remove_if` member as the safe single-pass replacement for
erase-while-iterating, and convert SymbolStringPool::clearDeadEntries, the
one in-tree violator, to use it.

In release builds DebugEpochBase/HandleBase are empty, so iterator size and
codegen are unchanged.
---
 llvm/docs/ProgrammersManual.rst               |  4 +
 llvm/docs/ReleaseNotes.md                     |  8 +-
 llvm/include/llvm/ADT/StringMap.h             | 78 ++++++++++++++-----
 .../ExecutionEngine/Orc/SymbolStringPool.h    |  6 +-
 llvm/unittests/ADT/StringMapTest.cpp          | 31 ++++++++
 5 files changed, 100 insertions(+), 27 deletions(-)

diff --git a/llvm/docs/ProgrammersManual.rst b/llvm/docs/ProgrammersManual.rst
index ddb9c57ad4e2a..6ee4badf6286b 100644
--- a/llvm/docs/ProgrammersManual.rst
+++ b/llvm/docs/ProgrammersManual.rst
@@ -2397,6 +2397,10 @@ copies a string if a value is inserted into the table.
 ``StringMap`` iteration order, however, is not guaranteed to be deterministic, so
 any uses which require that should instead use a ``std::map``.
 
+Like ``DenseMap``, ``StringMap`` iterators are invalidated whenever an insertion
+or erasure occurs.  To erase matching elements in a single pass, use the
+``remove_if`` member instead of erasing while iterating.
+
 .. _dss_indexmap:
 
 llvm/ADT/IndexedMap.h
diff --git a/llvm/docs/ReleaseNotes.md b/llvm/docs/ReleaseNotes.md
index 5cd88f4260011..fbab845af3dde 100644
--- a/llvm/docs/ReleaseNotes.md
+++ b/llvm/docs/ReleaseNotes.md
@@ -137,10 +137,10 @@ Makes programs 10x faster by doing Special New Thing.
 * ``ConstantFP`` now supports vector types and is the canonical form returned by
   ``ConstantVector::getSplat(C)`` when ``C`` is a scalar ``ConstantFP``.
 
-* ``DenseMap`` and ``DenseSet`` ``erase`` now invalidates all iterators and
-  references into the container, not just the iterator for the erased element.
-  Use the new ``remove_if`` member to erase matching elements in a single pass
-  instead of erasing while iterating.
+* ``DenseMap``, ``DenseSet``, ``StringMap``, and ``StringSet`` ``erase`` now
+  invalidates all iterators and references into the container, not just the
+  iterator for the erased element. Use the new ``remove_if`` member to erase
+  matching elements in a single pass instead of erasing while iterating.
 
 * ``TargetRegisterInfo::getMinimalPhysRegClass`` and related APIs have been
   refactored and no longer take a type. This API is also now precomputed in
diff --git a/llvm/include/llvm/ADT/StringMap.h b/llvm/include/llvm/ADT/StringMap.h
index 7901365daa462..41199b89a7f7b 100644
--- a/llvm/include/llvm/ADT/StringMap.h
+++ b/llvm/include/llvm/ADT/StringMap.h
@@ -14,6 +14,7 @@
 #ifndef LLVM_ADT_STRINGMAP_H
 #define LLVM_ADT_STRINGMAP_H
 
+#include "llvm/ADT/EpochTracker.h"
 #include "llvm/ADT/StringMapEntry.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/Support/AllocatorBase.h"
@@ -30,7 +31,7 @@ template <typename ValueTy> class StringMapKeyIterator;
 
 /// StringMapImpl - This is the base class of StringMap that is shared among
 /// all of its instantiations.
-class StringMapImpl {
+class StringMapImpl : public DebugEpochBase {
 protected:
   // Array of NumBuckets pointers to entries, null pointers are holes.
   // TheTable[NumBuckets] contains a sentinel value for easy iteration. Followed
@@ -220,13 +221,15 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
   using const_iterator = StringMapIterBase<ValueTy, true>;
   using iterator = StringMapIterBase<ValueTy, false>;
 
-  [[nodiscard]] iterator begin() { return iterator(TheTable, NumBuckets != 0); }
-  [[nodiscard]] iterator end() { return iterator(TheTable + NumBuckets); }
+  [[nodiscard]] iterator begin() {
+    return iterator(this, TheTable, NumBuckets != 0);
+  }
+  [[nodiscard]] iterator end() { return iterator(this, TheTable + NumBuckets); }
   [[nodiscard]] const_iterator begin() const {
-    return const_iterator(TheTable, NumBuckets != 0);
+    return const_iterator(this, TheTable, NumBuckets != 0);
   }
   [[nodiscard]] const_iterator end() const {
-    return const_iterator(TheTable + NumBuckets);
+    return const_iterator(this, TheTable + NumBuckets);
   }
 
   [[nodiscard]] iterator_range<StringMapKeyIterator<ValueTy>> keys() const {
@@ -240,7 +243,7 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
     int Bucket = FindKey(Key, FullHashValue);
     if (Bucket == -1)
       return end();
-    return iterator(TheTable + Bucket);
+    return iterator(this, TheTable + Bucket);
   }
 
   [[nodiscard]] const_iterator find(StringRef Key) const {
@@ -252,7 +255,7 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
     int Bucket = FindKey(Key, FullHashValue);
     if (Bucket == -1)
       return end();
-    return const_iterator(TheTable + Bucket);
+    return const_iterator(this, TheTable + Bucket);
   }
 
   /// lookup - Return the entry for the specified key, or a default
@@ -389,7 +392,7 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
     unsigned BucketNo = LookupBucketFor(Key, FullHashValue);
     StringMapEntryBase *&Bucket = TheTable[BucketNo];
     if (Bucket && Bucket != getTombstoneVal())
-      return {iterator(TheTable + BucketNo), false}; // Already exists in map.
+      return {iterator(this, TheTable + BucketNo), false}; // Already in map.
 
     if (Bucket == getTombstoneVal())
       --NumTombstones;
@@ -399,7 +402,7 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
     assert(NumItems + NumTombstones <= NumBuckets);
 
     BucketNo = RehashTable(BucketNo);
-    return {iterator(TheTable + BucketNo), true};
+    return {iterator(this, TheTable + BucketNo), true};
   }
 
   // clear - Empties out the StringMap
@@ -426,10 +429,36 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
 
   void erase(iterator I) {
     MapEntryTy &V = *I;
+    incrementEpoch();
     remove(&V);
     V.Destroy(getAllocator());
   }
 
+  /// Remove entries that match the given predicate. \p Pred is invoked with a
+  /// reference to each live entry and must not access the map being modified.
+  /// This is the safe replacement for erase-while-iterating.
+  ///
+  /// Returns whether anything was removed. If so, all iterators and references
+  /// into the map are invalidated.
+  template <typename Predicate> bool remove_if(Predicate Pred) {
+    bool Removed = false;
+    for (StringMapEntryBase *&Bucket : buckets()) {
+      if (!Bucket || Bucket == getTombstoneVal())
+        continue;
+      auto *Entry = static_cast<MapEntryTy *>(Bucket);
+      if (Pred(*Entry)) {
+        Entry->Destroy(getAllocator());
+        Bucket = getTombstoneVal();
+        --NumItems;
+        ++NumTombstones;
+        Removed = true;
+      }
+    }
+    if (Removed)
+      incrementEpoch();
+    return Removed;
+  }
+
   bool erase(StringRef Key) {
     iterator I = find(Key);
     if (I == end())
@@ -439,7 +468,11 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
   }
 };
 
-template <typename ValueTy, bool IsConst> class StringMapIterBase {
+template <typename ValueTy, bool IsConst>
+class StringMapIterBase : DebugEpochBase::HandleBase {
+  friend class StringMapIterBase<ValueTy, true>;
+  friend class StringMapIterBase<ValueTy, false>;
+
   StringMapEntryBase **Ptr = nullptr;
 
 public:
@@ -452,20 +485,31 @@ template <typename ValueTy, bool IsConst> class StringMapIterBase {
 
   StringMapIterBase() = default;
 
-  explicit StringMapIterBase(StringMapEntryBase **Bucket, bool Advance = false)
-      : Ptr(Bucket) {
+  explicit StringMapIterBase(const DebugEpochBase *Epoch,
+                             StringMapEntryBase **Bucket, bool Advance = false)
+      : DebugEpochBase::HandleBase(Epoch), Ptr(Bucket) {
     if (Advance)
       AdvancePastEmptyBuckets();
   }
 
+  // Converting ctor from non-const to const iterators. SFINAE'd out for const
+  // sources so it doesn't shadow the implicit copy constructor.
+  template <bool IsConstSrc,
+            typename = std::enable_if_t<!IsConstSrc && IsConst>>
+  StringMapIterBase(const StringMapIterBase<ValueTy, IsConstSrc> &I)
+      : DebugEpochBase::HandleBase(I), Ptr(I.Ptr) {}
+
   [[nodiscard]] reference operator*() const {
+    assert(isHandleInSync() && "invalid iterator access!");
     return *static_cast<value_type *>(*Ptr);
   }
   [[nodiscard]] pointer operator->() const {
+    assert(isHandleInSync() && "invalid iterator access!");
     return static_cast<value_type *>(*Ptr);
   }
 
   StringMapIterBase &operator++() { // Preincrement
+    assert(isHandleInSync() && "invalid iterator access!");
     ++Ptr;
     AdvancePastEmptyBuckets();
     return *this;
@@ -477,14 +521,12 @@ template <typename ValueTy, bool IsConst> class StringMapIterBase {
     return Tmp;
   }
 
-  template <bool ToConst,
-            typename = typename std::enable_if<!IsConst && ToConst>::type>
-  operator StringMapIterBase<ValueTy, ToConst>() const {
-    return StringMapIterBase<ValueTy, ToConst>(Ptr);
-  }
-
   friend bool operator==(const StringMapIterBase &LHS,
                          const StringMapIterBase &RHS) {
+    assert((!LHS.getEpochAddress() || LHS.isHandleInSync()) &&
+           "handle not in sync!");
+    assert((!RHS.getEpochAddress() || RHS.isHandleInSync()) &&
+           "handle not in sync!");
     return LHS.Ptr == RHS.Ptr;
   }
 
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/SymbolStringPool.h b/llvm/include/llvm/ExecutionEngine/Orc/SymbolStringPool.h
index 4b53de6af77bb..8f19ca7072ece 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/SymbolStringPool.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/SymbolStringPool.h
@@ -274,11 +274,7 @@ inline SymbolStringPtr SymbolStringPool::intern(StringRef S) {
 
 inline void SymbolStringPool::clearDeadEntries() {
   std::lock_guard<std::mutex> Lock(PoolMutex);
-  for (auto I = Pool.begin(), E = Pool.end(); I != E;) {
-    auto Tmp = I++;
-    if (Tmp->second == 0)
-      Pool.erase(Tmp);
-  }
+  Pool.remove_if([](PoolMapEntry &E) { return E.getValue() == 0; });
 }
 
 inline bool SymbolStringPool::empty() const {
diff --git a/llvm/unittests/ADT/StringMapTest.cpp b/llvm/unittests/ADT/StringMapTest.cpp
index 1d92de4e92325..4802b70102917 100644
--- a/llvm/unittests/ADT/StringMapTest.cpp
+++ b/llvm/unittests/ADT/StringMapTest.cpp
@@ -747,4 +747,35 @@ TEST(StringMapCustomTest, ConstIterator) {
                                          StringMap<int>::const_iterator>);
 }
 
+TEST(StringMapCustomTest, RemoveIf) {
+  StringMap<int> Map;
+  Map["a"] = 1;
+  Map["b"] = 2;
+  Map["c"] = 3;
+  Map["d"] = 4;
+
+  EXPECT_TRUE(Map.remove_if(
+      [](const StringMapEntry<int> &E) { return E.getValue() % 2 == 0; }));
+  EXPECT_EQ(2u, Map.size());
+  EXPECT_TRUE(Map.contains("a"));
+  EXPECT_FALSE(Map.contains("b"));
+  EXPECT_TRUE(Map.contains("c"));
+  EXPECT_FALSE(Map.contains("d"));
+
+  EXPECT_FALSE(Map.remove_if(
+      [](const StringMapEntry<int> &E) { return E.getValue() > 100; }));
+  EXPECT_EQ(2u, Map.size());
+}
+
+#if LLVM_ENABLE_ABI_BREAKING_CHECKS
+TEST(StringMapCustomTest, EraseInvalidatesIterators) {
+  StringMap<int> Map;
+  Map["a"] = 1;
+  Map["b"] = 2;
+  auto It = Map.find("a");
+  Map.erase(Map.find("b"));
+  EXPECT_DEATH((void)It->second, "invalid iterator access");
+}
+#endif
+
 } // end anonymous namespace

>From 69646524ab4e11ae3add3a61acfead931be0446e Mon Sep 17 00:00:00 2001
From: Fangrui Song <i at maskray.me>
Date: Sun, 7 Jun 2026 15:49:39 -0700
Subject: [PATCH 2/3] .

---
 llvm/include/llvm/ADT/StringMap.h | 16 ++++++++--------
 1 file changed, 8 insertions(+), 8 deletions(-)

diff --git a/llvm/include/llvm/ADT/StringMap.h b/llvm/include/llvm/ADT/StringMap.h
index 41199b89a7f7b..a219ac534c564 100644
--- a/llvm/include/llvm/ADT/StringMap.h
+++ b/llvm/include/llvm/ADT/StringMap.h
@@ -434,6 +434,14 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
     V.Destroy(getAllocator());
   }
 
+  bool erase(StringRef Key) {
+    iterator I = find(Key);
+    if (I == end())
+      return false;
+    erase(I);
+    return true;
+  }
+
   /// Remove entries that match the given predicate. \p Pred is invoked with a
   /// reference to each live entry and must not access the map being modified.
   /// This is the safe replacement for erase-while-iterating.
@@ -458,14 +466,6 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
       incrementEpoch();
     return Removed;
   }
-
-  bool erase(StringRef Key) {
-    iterator I = find(Key);
-    if (I == end())
-      return false;
-    erase(I);
-    return true;
-  }
 };
 
 template <typename ValueTy, bool IsConst>

>From b4c217f7c1114c312d6ac87364232d78d66f3d8e Mon Sep 17 00:00:00 2001
From: Fangrui Song <i at maskray.me>
Date: Sun, 7 Jun 2026 23:02:41 -0700
Subject: [PATCH 3/3] invalidate in more cases

---
 llvm/include/llvm/ADT/StringMap.h    |  5 +++++
 llvm/unittests/ADT/StringMapTest.cpp | 25 +++++++++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/llvm/include/llvm/ADT/StringMap.h b/llvm/include/llvm/ADT/StringMap.h
index a219ac534c564..5ce35f5c7a6ce 100644
--- a/llvm/include/llvm/ADT/StringMap.h
+++ b/llvm/include/llvm/ADT/StringMap.h
@@ -117,6 +117,8 @@ class StringMapImpl : public DebugEpochBase {
   [[nodiscard]] LLVM_ABI static uint32_t hash(StringRef Key);
 
   void swap(StringMapImpl &Other) {
+    incrementEpoch();
+    Other.incrementEpoch();
     std::swap(TheTable, Other.TheTable);
     std::swap(NumBuckets, Other.NumBuckets);
     std::swap(NumItems, Other.NumItems);
@@ -327,6 +329,7 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
     if (Bucket && Bucket != getTombstoneVal())
       return false; // Already exists in map.
 
+    incrementEpoch();
     if (Bucket == getTombstoneVal())
       --NumTombstones;
     Bucket = KeyValue;
@@ -394,6 +397,7 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
     if (Bucket && Bucket != getTombstoneVal())
       return {iterator(this, TheTable + BucketNo), false}; // Already in map.
 
+    incrementEpoch();
     if (Bucket == getTombstoneVal())
       --NumTombstones;
     Bucket =
@@ -407,6 +411,7 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
 
   // clear - Empties out the StringMap
   void clear() {
+    incrementEpoch();
     if (empty())
       return;
 
diff --git a/llvm/unittests/ADT/StringMapTest.cpp b/llvm/unittests/ADT/StringMapTest.cpp
index 4802b70102917..d350b46525ba0 100644
--- a/llvm/unittests/ADT/StringMapTest.cpp
+++ b/llvm/unittests/ADT/StringMapTest.cpp
@@ -768,6 +768,14 @@ TEST(StringMapCustomTest, RemoveIf) {
 }
 
 #if LLVM_ENABLE_ABI_BREAKING_CHECKS
+TEST(StringMapCustomTest, InsertInvalidatesIterators) {
+  StringMap<int> Map;
+  Map["a"] = 1;
+  auto It = Map.find("a");
+  Map.try_emplace("b", 2);
+  EXPECT_DEATH((void)It->second, "invalid iterator access");
+}
+
 TEST(StringMapCustomTest, EraseInvalidatesIterators) {
   StringMap<int> Map;
   Map["a"] = 1;
@@ -776,6 +784,23 @@ TEST(StringMapCustomTest, EraseInvalidatesIterators) {
   Map.erase(Map.find("b"));
   EXPECT_DEATH((void)It->second, "invalid iterator access");
 }
+
+TEST(StringMapCustomTest, ClearInvalidatesIterators) {
+  StringMap<int> Map;
+  Map["a"] = 1;
+  auto It = Map.find("a");
+  Map.clear();
+  EXPECT_DEATH((void)It->second, "invalid iterator access");
+}
+
+TEST(StringMapCustomTest, SwapInvalidatesIterators) {
+  StringMap<int> Map;
+  Map["a"] = 1;
+  auto It = Map.find("a");
+  StringMap<int> Other;
+  Map.swap(Other);
+  EXPECT_DEATH((void)It->second, "invalid iterator access");
+}
 #endif
 
 } // end anonymous namespace



More information about the llvm-commits mailing list