[llvm] 6252237 - [llvm][ADT] Add erase method to ScopedHashTableScope (#193516)

[via llvm-commits]

Author: lonely eagle
Date: 2026-04-29T19:56:43+08:00
New Revision: 6252237008b8a2cce7dc6da599cdfb74b057f410

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

LOG: [llvm][ADT] Add erase method to ScopedHashTableScope (#193516)

This PR introduces an `erase` method to `ScopedHashTable`, designed to
remove the most recent value associated with a given key within the
scope stack. To support efficient deletion, the internal
`ScopedHashTableVal` structure has been refactored into a doubly linked
list, allowing the predecessor of a node to be identified in O(1) time
during removal. Fix the MLIR CSE issue
https://github.com/llvm/llvm-project/pull/191135#discussion_r3100850607.
Part of https://github.com/llvm/llvm-project/pull/193778.

Added: 
    

Modified: 
    llvm/include/llvm/ADT/ScopedHashTable.h
    llvm/unittests/ADT/ScopedHashTableTest.cpp

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ADT/ScopedHashTable.h b/llvm/include/llvm/ADT/ScopedHashTable.h
index 97d15463fdfd7..cc977012b2394 100644
--- a/llvm/include/llvm/ADT/ScopedHashTable.h
+++ b/llvm/include/llvm/ADT/ScopedHashTable.h
@@ -46,6 +46,7 @@ template <typename K, typename V>
 class ScopedHashTableVal {
   ScopedHashTableVal *NextInScope;
   ScopedHashTableVal *NextForKey;
+  ScopedHashTableVal *PreInScope;
   K Key;
   V Val;
 
@@ -59,6 +60,7 @@ class ScopedHashTableVal {
   ScopedHashTableVal *getNextForKey() { return NextForKey; }
   const ScopedHashTableVal *getNextForKey() const { return NextForKey; }
   ScopedHashTableVal *getNextInScope() { return NextInScope; }
+  ScopedHashTableVal *getPreInScope() { return PreInScope; }
 
   template <typename AllocatorTy>
   static ScopedHashTableVal *Create(ScopedHashTableVal *nextInScope,
@@ -70,6 +72,9 @@ class ScopedHashTableVal {
     new (New) ScopedHashTableVal(key, val);
     New->NextInScope = nextInScope;
     New->NextForKey = nextForKey;
+    New->PreInScope = nullptr;
+    if (nextInScope)
+      nextInScope->PreInScope = New;
     return New;
   }
 
@@ -78,6 +83,20 @@ class ScopedHashTableVal {
     this->~ScopedHashTableVal();
     Allocator.Deallocate(this);
   }
+
+  template <typename AllocatorTy>
+  static void erase(ScopedHashTableVal<K, V> *&ThisEntry,
+                    AllocatorTy &Allocator) {
+    ScopedHashTableVal<K, V> *ToDestroy = ThisEntry;
+    ScopedHashTableVal<K, V> *NextInScope = ThisEntry->NextInScope;
+    ScopedHashTableVal<K, V> *PrevInScope = ThisEntry->PreInScope;
+    if (PrevInScope)
+      PrevInScope->NextInScope = NextInScope;
+    if (NextInScope)
+      NextInScope->PreInScope = PrevInScope;
+    ThisEntry = ThisEntry->NextForKey;
+    ToDestroy->Destroy(Allocator);
+  }
 };
 
 template <typename K, typename V, typename KInfo = DenseMapInfo<K>,
@@ -101,6 +120,7 @@ class ScopedHashTableScope {
 
   ScopedHashTableScope *getParentScope() { return PrevScope; }
   const ScopedHashTableScope *getParentScope() const { return PrevScope; }
+  void erase(const K &key);
 
 private:
   friend class ScopedHashTable<K, V, KInfo, AllocatorTy>;
@@ -219,6 +239,8 @@ class ScopedHashTable : detail::AllocatorHolder<AllocatorTy> {
                              getAllocator());
     S->setLastValInScope(KeyEntry);
   }
+
+  void erase(const K &key) { CurScope->erase(key); }
 };
 
 /// ScopedHashTableScope ctor - Install this as the current scope for the hash
@@ -257,6 +279,33 @@ ScopedHashTableScope<K, V, KInfo, Allocator>::~ScopedHashTableScope() {
   }
 }
 
+/// This method undoes the latest binding of the given key, effectively
+/// reverting to the previous state for that key. In the example at the
+/// beginning of this file, if we execute `HT.erase(0)` immediately after
+/// `HT.insert(0, 42);`, then the value associated with key "0" reverts to 0.
+/// This value is owned by "Scope1(HT)".
+template <typename K, typename V, typename KInfo, typename Allocator>
+void ScopedHashTableScope<K, V, KInfo, Allocator>::erase(const K &Key) {
+  auto It = HT.TopLevelMap.find(Key);
+  if (It == HT.TopLevelMap.end())
+    return;
+  ScopedHashTableVal<K, V> *&ThisEntry = It->second;
+
+  // `ThisEntry` may be the LastValInScope of a parent scope rather than the
+  // current scope. We iterate through the scope chain to find the scope
+  // that owns ThisEntry as its LastValInScope and update it accordingly.
+  auto *S = this;
+  while (S) {
+    if (ThisEntry == S->LastValInScope) {
+      S->LastValInScope = ThisEntry->getNextInScope();
+      break;
+    }
+    S = S->PrevScope;
+  }
+  if (ThisEntry->getNextForKey() == nullptr)
+    HT.TopLevelMap.erase(It);
+  ScopedHashTableVal<K, V>::erase(ThisEntry, HT.getAllocator());
+}
 } // end namespace llvm
 
 #endif // LLVM_ADT_SCOPEDHASHTABLE_H

diff  --git a/llvm/unittests/ADT/ScopedHashTableTest.cpp b/llvm/unittests/ADT/ScopedHashTableTest.cpp
index 8ce5c7cecf998..d6f32b4002dcc 100644
--- a/llvm/unittests/ADT/ScopedHashTableTest.cpp
+++ b/llvm/unittests/ADT/ScopedHashTableTest.cpp
@@ -31,6 +31,7 @@ class ScopedHashTableTest : public Test {
   static constexpr StringLiteral kLocalName = "local";
   static constexpr StringLiteral kLocalValue = "lvalue";
   static constexpr StringLiteral kLocalValue2 = "lvalue2";
+  static constexpr StringLiteral KLocalValue3 = "lvalue3";
 };
 
 TEST_F(ScopedHashTableTest, AccessWithNoActiveScope) {
@@ -151,3 +152,32 @@ TEST_F(ScopedHashTableTest, DISABLED_PopScopeOnStack) {
   // There is no way to write an assert on an assert in googletest so that we
   // mark the test case as DISABLED.
 }
+
+TEST_F(ScopedHashTableTest, EraseKeyInHashTable) {
+  using SymbolTableScopeTy = ScopedHashTable<StringRef, StringRef>::ScopeTy;
+  StringRef curTestVarA = "a";
+  StringRef curTestVarB = "b";
+  SymbolTableScopeTy funScope(symbolTable);
+  symbolTable.insert(kLocalName, kLocalValue);
+  symbolTable.insert(curTestVarA, curTestVarA);
+  symbolTable.insert(curTestVarB, curTestVarB);
+
+  SymbolTableScopeTy funScope2(symbolTable);
+  symbolTable.insert(curTestVarA, curTestVarA);
+  symbolTable.insert(kLocalName, kLocalValue2);
+  symbolTable.insert(curTestVarB, curTestVarB);
+
+  SymbolTableScopeTy funScope3(symbolTable);
+  symbolTable.insert(curTestVarA, curTestVarA);
+  symbolTable.insert(curTestVarB, curTestVarB);
+  symbolTable.insert(kLocalName, KLocalValue3);
+  SymbolTableScopeTy funScope4(symbolTable);
+
+  EXPECT_EQ(symbolTable.lookup(kLocalName), KLocalValue3);
+  symbolTable.erase(kLocalName);
+  EXPECT_EQ(symbolTable.lookup(kLocalName), kLocalValue2);
+  symbolTable.erase(kLocalName);
+  EXPECT_EQ(symbolTable.lookup(kLocalName), kLocalValue);
+  symbolTable.erase(kLocalName);
+  EXPECT_EQ(symbolTable.lookup(kLocalName), StringRef());
+}