[llvm] c2df149 - [StringMap] Add remove_if and use it for erase-while-iterating (#202272)

[via llvm-commits]

Author: Fangrui Song
Date: 2026-06-08T02:18:28-07:00
New Revision: c2df14973f7f51abfe284641125abe02b3bd374c

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

LOG: [StringMap] Add remove_if and use it for erase-while-iterating (#202272)

Add a `remove_if` member to StringMap (and to HashKeyMap, the base of
SampleProfileMap) as a replacement for the erase-while-iterating idiom,
and convert the two in-tree users: SymbolStringPool::clearDeadEntries
and llvm-profdata's filterFunctions (a template over StringMap and
SampleProfileMap).

Extracted from #202237 - making StringMap's mutation invalidates
iterators
so that we can remove the tombstone state.

Aided by Claude Opus 4.8

Added: 
    

Modified: 
    llvm/include/llvm/ADT/StringMap.h
    llvm/include/llvm/ExecutionEngine/Orc/SymbolStringPool.h
    llvm/include/llvm/ProfileData/HashKeyMap.h
    llvm/tools/llvm-profdata/llvm-profdata.cpp
    llvm/unittests/ADT/StringMapTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ADT/StringMap.h b/llvm/include/llvm/ADT/StringMap.h
index 7901365daa462..7f051b4685751 100644
--- a/llvm/include/llvm/ADT/StringMap.h
+++ b/llvm/include/llvm/ADT/StringMap.h
@@ -430,6 +430,26 @@ class LLVM_ALLOCATORHOLDER_EMPTYBASE StringMap
     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.
+  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;
+      }
+    }
+    return Removed;
+  }
+
   bool erase(StringRef Key) {
     iterator I = find(Key);
     if (I == end())

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/include/llvm/ProfileData/HashKeyMap.h b/llvm/include/llvm/ProfileData/HashKeyMap.h
index fceb95143340f..00624b9cba18e 100644
--- a/llvm/include/llvm/ProfileData/HashKeyMap.h
+++ b/llvm/include/llvm/ProfileData/HashKeyMap.h
@@ -119,6 +119,21 @@ class HashKeyMap :
   iterator erase(const_iterator It) {
     return base_type::erase(It);
   }
+
+  /// Remove entries that match the given predicate. \p Pred is invoked with a
+  /// reference to each entry. Returns whether anything was removed.
+  template <typename Predicate> bool remove_if(Predicate Pred) {
+    bool Removed = false;
+    for (auto It = base_type::begin(), E = base_type::end(); It != E;) {
+      if (Pred(*It)) {
+        It = base_type::erase(It);
+        Removed = true;
+      } else {
+        ++It;
+      }
+    }
+    return Removed;
+  }
 };
 
 }

diff  --git a/llvm/tools/llvm-profdata/llvm-profdata.cpp b/llvm/tools/llvm-profdata/llvm-profdata.cpp
index 6bf3179696e4a..08a1ea51af8e5 100644
--- a/llvm/tools/llvm-profdata/llvm-profdata.cpp
+++ b/llvm/tools/llvm-profdata/llvm-profdata.cpp
@@ -932,17 +932,15 @@ static void filterFunctions(T &ProfileMap) {
   std::string NegativeMD5Name =
       std::to_string(llvm::MD5Hash(FuncNameNegativeFilter));
 
-  for (auto I = ProfileMap.begin(); I != ProfileMap.end();) {
-    auto Tmp = I++;
-    const auto &FuncName = getFuncName(*Tmp);
+  ProfileMap.remove_if([&](const auto &Entry) {
+    const auto &FuncName = getFuncName(Entry);
     // Negative filter has higher precedence than positive filter.
-    if ((hasNegativeFilter &&
-         (NegativePattern.match(FuncName) ||
-          (FunctionSamples::UseMD5 && NegativeMD5Name == FuncName))) ||
-        (hasFilter && !(Pattern.match(FuncName) ||
-                        (FunctionSamples::UseMD5 && MD5Name == FuncName))))
-      ProfileMap.erase(Tmp);
-  }
+    return (hasNegativeFilter &&
+            (NegativePattern.match(FuncName) ||
+             (FunctionSamples::UseMD5 && NegativeMD5Name == FuncName))) ||
+           (hasFilter && !(Pattern.match(FuncName) ||
+                           (FunctionSamples::UseMD5 && MD5Name == FuncName)));
+  });
 
   llvm::dbgs() << Count - ProfileMap.size() << " of " << Count << " functions "
                << "in the original profile are filtered.\n";

diff  --git a/llvm/unittests/ADT/StringMapTest.cpp b/llvm/unittests/ADT/StringMapTest.cpp
index 1d92de4e92325..51225a8d096ad 100644
--- a/llvm/unittests/ADT/StringMapTest.cpp
+++ b/llvm/unittests/ADT/StringMapTest.cpp
@@ -747,4 +747,24 @@ 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());
+}
+
 } // end anonymous namespace