[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