[llvm] [ADT] Speed up FoldingSet with linear probing and Algorithm R deletion (PR #218190)

Fangrui Song via llvm-commits llvm-commits at lists.llvm.org
Sat Aug 22 21:03:54 PDT 2026


https://github.com/MaskRay created https://github.com/llvm/llvm-project/pull/218190

FoldingSet is a chained hash table: each node holds a pointer to the
next node in its bucket, so a lookup chases a pointer into a node at
every occupied bucket it probes.

Switch to the StringMap layout: an array of node pointers, null for an
empty slot, followed by a parallel array of bucket hashes that rejects
mismatches before the profile compare. Deletion uses Knuth TAOCP 6.4
Algorithm R, as DenseMap, SmallPtrSet and StringMap do, so erase
invalidates iterators while leaving pointers to nodes valid.

The next-in-bucket pointer becomes a cached 32-bit hash, halving
FoldingSetNode. FindNodeOrInsertPos returns that hash instead of a
bucket address, so an InsertPos survives intervening insertions.

FoldingSetBase derives from DebugEpochBase and its iterators from
HandleBase, so a stale iterator fails under
LLVM_ENABLE_ABI_BREAKING_CHECKS.


>From 6015ce53e6ee769a55222f9306366d6f69957a4f Mon Sep 17 00:00:00 2001
From: Fangrui Song <i at maskray.me>
Date: Sat, 22 Aug 2026 18:10:10 -0700
Subject: [PATCH] [ADT] Speed up FoldingSet with linear probing and Algorithm R
 deletion

FoldingSet is a chained hash table: each node holds a pointer to the
next node in its bucket, so a lookup chases a pointer into a node at
every occupied bucket it probes.

Switch to the StringMap layout: an array of node pointers, null for an
empty slot, followed by a parallel array of bucket hashes that rejects
mismatches before the profile compare. Deletion uses Knuth TAOCP 6.4
Algorithm R, as DenseMap, SmallPtrSet and StringMap do, so erase
invalidates iterators while leaving pointers to nodes valid.

The next-in-bucket pointer becomes a cached 32-bit hash, halving
FoldingSetNode. FindNodeOrInsertPos returns that hash instead of a
bucket address, so an InsertPos survives intervening insertions.

FoldingSetBase derives from DebugEpochBase and its iterators from
HandleBase, so a stale iterator fails under
LLVM_ENABLE_ABI_BREAKING_CHECKS.
---
 llvm/docs/ProgrammersManual.md     |   9 +-
 llvm/include/llvm/ADT/FoldingSet.h |  96 ++++++-----
 llvm/lib/Support/FoldingSet.cpp    | 259 +++++++++++------------------
 llvm/unittests/ADT/FoldingSet.cpp  |  99 +++++++++++
 4 files changed, 259 insertions(+), 204 deletions(-)

diff --git a/llvm/docs/ProgrammersManual.md b/llvm/docs/ProgrammersManual.md
index 16f24d6565fad..e78ce9a36fe4c 100644
--- a/llvm/docs/ProgrammersManual.md
+++ b/llvm/docs/ProgrammersManual.md
@@ -2058,8 +2058,8 @@ building composite data structures.
 #### llvm/ADT/FoldingSet.h
 
 `FoldingSet` is an aggregate class that is really good at uniquing
-expensive-to-create or polymorphic objects.  It is a combination of a chained
-hash table with intrusive links (uniqued objects are required to inherit from
+expensive-to-create or polymorphic objects.  It is a linear-probed hash table
+whose buckets point to the uniqued objects (which are required to inherit from
 `FoldingSetNode`) that uses {ref}`SmallVector <dss_smallvector>` as part of its ID
 process.
 
@@ -2076,11 +2076,12 @@ element that we want to query for.  The query either returns the element
 matching the ID or it returns an opaque ID that indicates where insertion should
 take place.  Construction of the ID usually does not require heap traffic.
 
-Because `FoldingSet` uses intrusive links, it can support polymorphic objects in
+Because the buckets are pointers, `FoldingSet` can support polymorphic objects in
 the set (for example, you can have `SDNode` instances mixed with `LoadSDNodes`).
 Because the elements are individually allocated, pointers to the elements are
 stable: inserting or removing elements does not invalidate any pointers to other
-elements.
+elements.  Removing an element relocates other buckets, so it invalidates
+iterators.
 
 (dss_set)=
 
diff --git a/llvm/include/llvm/ADT/FoldingSet.h b/llvm/include/llvm/ADT/FoldingSet.h
index ab4fa2712d4a5..29efe3e48ee96 100644
--- a/llvm/include/llvm/ADT/FoldingSet.h
+++ b/llvm/include/llvm/ADT/FoldingSet.h
@@ -16,6 +16,7 @@
 #ifndef LLVM_ADT_FOLDINGSET_H
 #define LLVM_ADT_FOLDINGSET_H
 
+#include "llvm/ADT/EpochTracker.h"
 #include "llvm/ADT/Hashing.h"
 #include "llvm/ADT/STLForwardCompat.h"
 #include "llvm/ADT/SmallVector.h"
@@ -37,9 +38,8 @@ namespace llvm {
 ///      it, otherwise return the bucket it should be inserted into.
 ///   2. Given a node that has already been created, remove it from the set.
 ///
-/// This class is implemented as a single-link chained hash table, where the
-/// "buckets" are actually the nodes themselves (the next pointer is in the
-/// node).  The last node points back to the bucket to simplify node removal.
+/// This class is implemented as an open-addressed hash table using linear
+/// probing, with Knuth TAOCP 6.4 Algorithm R deletion.
 ///
 /// Any node that is to be included in the folding set must be a subclass of
 /// FoldingSetNode.  The node class must also define a Profile method used to
@@ -182,7 +182,11 @@ class FoldingSetNodeIDRef {
   // Compute a strong hash value used to lookup the node in the FoldingSetBase.
   // The hash value is not guaranteed to be deterministic across processes.
   unsigned ComputeHash() const {
-    return static_cast<unsigned>(hash_combine_range(Data, Data + Size));
+    unsigned Hash =
+        static_cast<unsigned>(hash_combine_range(Data, Data + Size));
+    // FoldingSetBase hands a hash back to the caller as a non-null InsertPos
+    // token, which on a 32-bit host leaves no encoding for UINT32_MAX.
+    return Hash == UINT32_MAX ? 0 : Hash;
   }
 
   // Compute a deterministic hash value across processes that is suitable for
@@ -289,23 +293,19 @@ class FoldingSetNodeID {
 };
 
 //===----------------------------------------------------------------------===//
-/// Implements the folding set functionality. The main structure is an array of
-/// buckets.  Each bucket is indexed by the hash of the nodes it contains. The
-/// bucket itself points to the nodes contained in the bucket via a singly
-/// linked list.  The last node in the list points back to the bucket to
-/// facilitate node removal.
-///
-class FoldingSetBase {
+/// Non-templated base class for FoldingSet and ContextualFoldingSet, holding
+/// the memory management and probing that does not depend on the node type.
+class FoldingSetBase : public DebugEpochBase {
 protected:
-  /// Array of bucket chains.
-  void **Buckets;
+  /// Array of node pointers; a null entry marks an empty slot. A parallel
+  /// array of 32-bit hashes shares the same allocation, see getHashes().
+  void **Buckets = nullptr;
 
   /// Length of the Buckets array.  Always a power of 2.
-  unsigned NumBuckets;
+  unsigned NumBuckets = 0;
 
-  /// Number of nodes in the folding set. Growth occurs when NumNodes
-  /// is greater than twice the number of buckets.
-  unsigned NumNodes;
+  /// Number of nodes in the folding set.
+  unsigned NumNodes = 0;
 
   LLVM_ABI explicit FoldingSetBase(unsigned Log2InitSize);
   LLVM_ABI FoldingSetBase(FoldingSetBase &&Arg);
@@ -314,19 +314,19 @@ class FoldingSetBase {
 
 public:
   //===--------------------------------------------------------------------===//
-  /// This class is used to maintain the singly linked bucket list in
-  /// a folding set.
+  /// This class is used to maintain node state in a folding set.
   class Node {
   private:
-    // NextInFoldingSetBucket - next link in the bucket list.
-    void *NextInFoldingSetBucket = nullptr;
+    // Hash of the node's profile, cached so that growth and removal never
+    // re-run Profile().
+    uint32_t FoldingSetHash = 0;
 
   public:
     Node() = default;
 
     // Accessors
-    void *getNextInBucket() const { return NextInFoldingSetBucket; }
-    void SetNextInBucket(void *N) { NextInFoldingSetBucket = N; }
+    uint32_t getFoldingSetHash() const { return FoldingSetHash; }
+    void setFoldingSetHash(uint32_t Hash) { FoldingSetHash = Hash; }
   };
 
   /// Remove all nodes from the folding set.
@@ -340,11 +340,7 @@ class FoldingSetBase {
 
   /// Returns the number of nodes permitted in the folding set
   /// before a rebucket operation is performed.
-  unsigned capacity() const {
-    // We allow a load factor of up to 2.0,
-    // so that means our capacity is NumBuckets * 2
-    return NumBuckets * 2;
-  }
+  unsigned capacity() const { return NumBuckets - NumBuckets / 4; }
 
 protected:
   /// Functions provided by the derived class to compute folding properties.
@@ -369,6 +365,18 @@ class FoldingSetBase {
   };
 
 private:
+  /// The hashes of the nodes in Buckets, in the same order. Only entries whose
+  /// bucket is non-null are live.
+  uint32_t *getHashes() const {
+    return reinterpret_cast<uint32_t *>(Buckets + NumBuckets);
+  }
+
+  /// Insert \p N at the first empty slot following its home, without checking
+  /// capacity.
+  void insertImpl(Node *N, uint32_t Hash);
+
+  friend class FoldingSetIteratorImpl;
+
   /// Resize the hash table and rehash everything. \p NewBucketCount must be a
   /// power of two, and must be greater than the old bucket count.
   void GrowBucketCount(unsigned NewBucketCount, const FoldingSetInfo &Info);
@@ -496,13 +504,13 @@ class FoldingSetImpl : public FoldingSetBase, public Trait::ContextStorage {
 public:
   using iterator = FoldingSetIterator<T>;
 
-  iterator begin() { return iterator(Buckets); }
-  iterator end() { return iterator(Buckets + NumBuckets); }
+  iterator begin() { return iterator(this, 0); }
+  iterator end() { return iterator(this, NumBuckets); }
 
   using const_iterator = FoldingSetIterator<const T>;
 
-  const_iterator begin() const { return const_iterator(Buckets); }
-  const_iterator end() const { return const_iterator(Buckets + NumBuckets); }
+  const_iterator begin() const { return const_iterator(this, 0); }
+  const_iterator end() const { return const_iterator(this, NumBuckets); }
 
   /// Grow the number of buckets so that we can hold at least \p EltCount
   /// nodes before rebucketing. May allocate more space than requested.
@@ -636,30 +644,38 @@ template <class T, class VectorT = SmallVector<T *, 8>> class FoldingSetVector {
 //===----------------------------------------------------------------------===//
 /// This is the common iterator support shared by all folding sets, which knows
 /// how to walk the folding set hash table.
-class FoldingSetIteratorImpl {
+class FoldingSetIteratorImpl : DebugEpochBase::HandleBase {
 protected:
-  FoldingSetNode *NodePtr;
+  const FoldingSetBase *Set = nullptr;
+  unsigned Index = 0;
 
-  LLVM_ABI FoldingSetIteratorImpl(void **Bucket);
+  LLVM_ABI FoldingSetIteratorImpl(const FoldingSetBase *Set, unsigned Index);
 
   LLVM_ABI void advance();
 
+  FoldingSetNode *getNode() const {
+    assert(isHandleInSync() && "invalid iterator access!");
+    return static_cast<FoldingSetNode *>(Set->Buckets[Index]);
+  }
+
 public:
   bool operator==(const FoldingSetIteratorImpl &RHS) const {
-    return NodePtr == RHS.NodePtr;
+    assert(isHandleInSync() && RHS.isHandleInSync() && "handle not in sync!");
+    return Index == RHS.Index;
   }
   bool operator!=(const FoldingSetIteratorImpl &RHS) const {
-    return NodePtr != RHS.NodePtr;
+    return !(*this == RHS);
   }
 };
 
 template <class T> class FoldingSetIterator : public FoldingSetIteratorImpl {
 public:
-  explicit FoldingSetIterator(void **Bucket) : FoldingSetIteratorImpl(Bucket) {}
+  explicit FoldingSetIterator(const FoldingSetBase *Set, unsigned Index)
+      : FoldingSetIteratorImpl(Set, Index) {}
 
-  T &operator*() const { return *static_cast<T *>(NodePtr); }
+  T &operator*() const { return *static_cast<T *>(getNode()); }
 
-  T *operator->() const { return static_cast<T *>(NodePtr); }
+  T *operator->() const { return static_cast<T *>(getNode()); }
 
   inline FoldingSetIterator &operator++() { // Preincrement
     advance();
diff --git a/llvm/lib/Support/FoldingSet.cpp b/llvm/lib/Support/FoldingSet.cpp
index d9ae1aca5fc4a..acaca9d5fe451 100644
--- a/llvm/lib/Support/FoldingSet.cpp
+++ b/llvm/lib/Support/FoldingSet.cpp
@@ -131,66 +131,55 @@ FoldingSetNodeID::Intern(BumpPtrAllocator &Allocator) const {
 }
 
 //===----------------------------------------------------------------------===//
-/// Helper functions for FoldingSetBase.
-
-/// GetNextPtr - In order to save space, each bucket is a
-/// singly-linked-list. In order to make deletion more efficient, we make
-/// the list circular, so we can delete a node without computing its hash.
-/// The problem with this is that the start of the hash buckets are not
-/// Nodes. If NextInBucketPtr is a bucket pointer, this method returns null:
-/// use GetBucketPtr when this happens.
-static FoldingSetBase::Node *GetNextPtr(void *NextInBucketPtr) {
-  // The low bit is set if this is the pointer back to the bucket.
-  if (reinterpret_cast<intptr_t>(NextInBucketPtr) & 1)
-    return nullptr;
-
-  return static_cast<FoldingSetBase::Node *>(NextInBucketPtr);
-}
-
-/// GetBucketPtr - Provides a casting of a bucket pointer for isNode
-/// testing.
-static void **GetBucketPtr(void *NextInBucketPtr) {
-  intptr_t Ptr = reinterpret_cast<intptr_t>(NextInBucketPtr);
-  assert((Ptr & 1) && "Not a bucket pointer");
-  return reinterpret_cast<void **>(Ptr & ~intptr_t(1));
+// FoldingSetBase Implementation
+//
+// FoldingSet is an open-addressed hash set using linear probing. One
+// allocation holds the bucket array followed by a parallel array of the
+// buckets' 32-bit hashes:
+//   [ Buckets (NumBuckets * sizeof(void *)) ][ Hashes (NumBuckets * 4) ]
+// A null bucket marks an empty slot, and the hash array rejects mismatches
+// before the profile compare, so walking a probe chain touches no nodes.
+// Nodes cache their own hash as well, which is what lets RemoveNode() find a
+// node without re-running Profile().
+//
+// Removal uses Knuth TAOCP vol. 3 6.4 Algorithm R, as StringMap, DenseMap and
+// SmallPtrSet do: it closes the hole rather than leaving a tombstone, so the
+// table stays sized to the live node count under insert/erase churn.
+
+/// Encode a hash into the non-null token FindNodeOrInsertPos hands back.
+/// Unlike a bucket address the token survives intervening insertions.
+static void *encodeHash(uint32_t Hash) {
+  return reinterpret_cast<void *>(static_cast<uintptr_t>(Hash) + 1);
 }
 
-/// GetBucketFor - Hash the specified node ID and return the hash bucket for
-/// the specified ID.
-static void **GetBucketFor(unsigned Hash, void **Buckets, unsigned NumBuckets) {
-  // NumBuckets is always a power of 2.
-  unsigned BucketNum = Hash & (NumBuckets - 1);
-  return Buckets + BucketNum;
+static uint32_t decodeHash(void *InsertPos) {
+  return static_cast<uint32_t>(reinterpret_cast<uintptr_t>(InsertPos) - 1);
 }
 
-/// AllocateBuckets - Allocate initialized bucket memory.
+/// AllocateBuckets - Allocate zeroed bucket and hash arrays.
 static void **AllocateBuckets(unsigned NumBuckets) {
-  void **Buckets =
-      static_cast<void **>(safe_calloc(NumBuckets + 1, sizeof(void *)));
-  // Set the very last bucket to be a non-null "pointer".
-  Buckets[NumBuckets] = reinterpret_cast<void *>(-1);
-  return Buckets;
+  return static_cast<void **>(
+      safe_calloc(NumBuckets, sizeof(void *) + sizeof(uint32_t)));
 }
 
-//===----------------------------------------------------------------------===//
-// FoldingSetBase Implementation
-
 FoldingSetBase::FoldingSetBase(unsigned Log2InitSize) {
   assert(5 < Log2InitSize && Log2InitSize < 32 &&
          "Initial hash table size out of range");
   NumBuckets = 1 << Log2InitSize;
   Buckets = AllocateBuckets(NumBuckets);
-  NumNodes = 0;
 }
 
 FoldingSetBase::FoldingSetBase(FoldingSetBase &&Arg)
     : Buckets(Arg.Buckets), NumBuckets(Arg.NumBuckets), NumNodes(Arg.NumNodes) {
+  Arg.incrementEpoch();
   Arg.Buckets = nullptr;
   Arg.NumBuckets = 0;
   Arg.NumNodes = 0;
 }
 
 FoldingSetBase &FoldingSetBase::operator=(FoldingSetBase &&RHS) {
+  incrementEpoch();
+  RHS.incrementEpoch();
   free(Buckets); // This may be null if the set is in a moved-from state.
   Buckets = RHS.Buckets;
   NumBuckets = RHS.NumBuckets;
@@ -204,16 +193,25 @@ FoldingSetBase &FoldingSetBase::operator=(FoldingSetBase &&RHS) {
 FoldingSetBase::~FoldingSetBase() { free(Buckets); }
 
 void FoldingSetBase::clear() {
-  // Set all but the last bucket to null pointers.
-  memset(Buckets, 0, NumBuckets * sizeof(void *));
-
-  // Set the very last bucket to be a non-null "pointer".
-  Buckets[NumBuckets] = reinterpret_cast<void *>(-1);
-
-  // Reset the node count to zero.
+  incrementEpoch();
+  // Stale hashes are unreachable, so only the occupancy needs resetting.
+  if (NumBuckets)
+    memset(Buckets, 0, NumBuckets * sizeof(void *));
   NumNodes = 0;
 }
 
+void FoldingSetBase::insertImpl(Node *N, uint32_t Hash) {
+  incrementEpoch();
+  unsigned Mask = NumBuckets - 1;
+  unsigned I = Hash & Mask;
+  while (Buckets[I])
+    I = (I + 1) & Mask;
+  Buckets[I] = N;
+  getHashes()[I] = Hash;
+  N->setFoldingSetHash(Hash);
+  ++NumNodes;
+}
+
 void FoldingSetBase::GrowBucketCount(unsigned NewBucketCount,
                                      const FoldingSetInfo &Info) {
   assert((NewBucketCount > NumBuckets) &&
@@ -221,126 +219,82 @@ void FoldingSetBase::GrowBucketCount(unsigned NewBucketCount,
   assert(isPowerOf2_32(NewBucketCount) && "Bad bucket count!");
 
   FoldingSetBase Tmp(llvm::Log2_32(NewBucketCount));
-  FoldingSetNodeID TempID;
-  for (unsigned i = 0; i != NumBuckets; ++i) {
-    void *Probe = Buckets[i];
-    if (!Probe)
-      continue;
-    while (Node *NodeInBucket = GetNextPtr(Probe)) {
-      // Figure out the next link, remove NodeInBucket from the old link.
-      Probe = NodeInBucket->getNextInBucket();
-      NodeInBucket->SetNextInBucket(nullptr);
-
-      // Insert the node into the new bucket, after recomputing the hash.
-      Tmp.InsertNode(
-          NodeInBucket,
-          GetBucketFor(Info.ComputeNodeHash(this, NodeInBucket, TempID),
-                       Tmp.Buckets, Tmp.NumBuckets),
-          Info);
-      TempID.clear();
-    }
-  }
+  const uint32_t *Hashes = getHashes();
+  for (unsigned I = 0; I != NumBuckets; ++I)
+    if (void *N = Buckets[I])
+      Tmp.insertImpl(static_cast<Node *>(N), Hashes[I]);
 
   *this = std::move(Tmp);
 }
 
 void FoldingSetBase::reserve(unsigned EltCount, const FoldingSetInfo &Info) {
-  // This will give us somewhere between EltCount / 2 and
-  // EltCount buckets.  This puts us in the load factor
-  // range of 1.0 - 2.0.
   if (EltCount <= capacity())
     return;
-  GrowBucketCount(llvm::bit_floor(EltCount), Info);
+  uint64_t Required = divideCeil(uint64_t(EltCount) * 4, 3);
+  GrowBucketCount(
+      static_cast<unsigned>(llvm::bit_ceil(std::max<uint64_t>(Required, 64))),
+      Info);
 }
 
 FoldingSetBase::Node *FoldingSetBase::FindNodeOrInsertPos(
     const FoldingSetNodeID &ID, void *&InsertPos, const FoldingSetInfo &Info) {
   unsigned IDHash = ID.ComputeHash();
-  void **Bucket = GetBucketFor(IDHash, Buckets, NumBuckets);
-  void *Probe = *Bucket;
-
-  InsertPos = nullptr;
+  const uint32_t *Hashes = getHashes();
+  unsigned Mask = NumBuckets - 1;
 
   FoldingSetNodeID TempID;
-  while (Node *NodeInBucket = GetNextPtr(Probe)) {
-    if (Info.NodeEquals(this, NodeInBucket, ID, IDHash, TempID))
-      return NodeInBucket;
+  for (unsigned I = IDHash & Mask; Buckets[I]; I = (I + 1) & Mask) {
+    // Reject on the hash first: the common case only reads the bucket and hash
+    // arrays, which matters for cache locality.
+    if (Hashes[I] != IDHash)
+      continue;
+    Node *N = static_cast<Node *>(Buckets[I]);
+    if (Info.NodeEquals(this, N, ID, IDHash, TempID)) {
+      InsertPos = nullptr;
+      return N;
+    }
     TempID.clear();
-
-    Probe = NodeInBucket->getNextInBucket();
   }
 
-  // Didn't find the node, return null with the bucket as the InsertPos.
-  InsertPos = Bucket;
+  // Didn't find the node, hand back the hash so that InsertNode can place it.
+  InsertPos = encodeHash(IDHash);
   return nullptr;
 }
 
 void FoldingSetBase::InsertNode(Node *N, void *InsertPos,
                                 const FoldingSetInfo &Info) {
-  assert(!N->getNextInBucket());
-  // Do we need to grow the hashtable?
-  if (NumNodes + 1 > capacity()) {
+  assert(InsertPos && "Invalid InsertPos!");
+  if (NumNodes + 1 > capacity())
     GrowBucketCount(NumBuckets * 2, Info);
-    FoldingSetNodeID TempID;
-    InsertPos = GetBucketFor(Info.ComputeNodeHash(this, N, TempID), Buckets,
-                             NumBuckets);
-  }
-
-  ++NumNodes;
-
-  /// The insert position is actually a bucket pointer.
-  void **Bucket = static_cast<void **>(InsertPos);
-
-  void *Next = *Bucket;
-
-  // If this is the first insertion into this bucket, its next pointer will be
-  // null.  Pretend as if it pointed to itself, setting the low bit to indicate
-  // that it is a pointer to the bucket.
-  if (!Next)
-    Next = reinterpret_cast<void *>(reinterpret_cast<intptr_t>(Bucket) | 1);
-
-  // Set the node's next pointer, and make the bucket point to the node.
-  N->SetNextInBucket(Next);
-  *Bucket = N;
+  insertImpl(N, decodeHash(InsertPos));
 }
 
 bool FoldingSetBase::RemoveNode(Node *N) {
-  // Because each bucket is a circular list, we don't need to compute N's hash
-  // to remove it.
-  void *Ptr = N->getNextInBucket();
-  if (!Ptr)
-    return false; // Not in folding set.
+  uint32_t *Hashes = getHashes();
+  unsigned Mask = NumBuckets - 1;
+
+  unsigned I = N->getFoldingSetHash() & Mask;
+  while (Buckets[I] != N) {
+    if (!Buckets[I])
+      return false; // Not in folding set.
+    I = (I + 1) & Mask;
+  }
 
-  --NumNodes;
-  N->SetNextInBucket(nullptr);
-
-  // Remember what N originally pointed to, either a bucket or another node.
-  void *NodeNextPtr = Ptr;
-
-  // Chase around the list until we find the node (or bucket) which points to N.
-  while (true) {
-    if (Node *NodeInBucket = GetNextPtr(Ptr)) {
-      // Advance pointer.
-      Ptr = NodeInBucket->getNextInBucket();
-
-      // We found a node that points to N, change it to point to N's next node,
-      // removing N from the list.
-      if (Ptr == N) {
-        NodeInBucket->SetNextInBucket(NodeNextPtr);
-        return true;
-      }
-    } else {
-      void **Bucket = GetBucketPtr(Ptr);
-      Ptr = *Bucket;
-
-      // If we found that the bucket points to N, update the bucket to point to
-      // whatever is next.
-      if (Ptr == N) {
-        *Bucket = NodeNextPtr;
-        return true;
-      }
+  incrementEpoch();
+
+  // Knuth TAOCP 6.4 Algorithm R: walk forward sliding each following entry
+  // whose probe path crosses the hole.
+  for (unsigned J = (I + 1) & Mask; Buckets[J]; J = (J + 1) & Mask) {
+    unsigned Ideal = Hashes[J];
+    if (((I - Ideal) & Mask) < ((J - Ideal) & Mask)) {
+      Buckets[I] = Buckets[J];
+      Hashes[I] = Hashes[J];
+      I = J;
     }
   }
+  Buckets[I] = nullptr;
+  --NumNodes;
+  return true;
 }
 
 FoldingSetBase::Node *
@@ -357,31 +311,16 @@ FoldingSetBase::GetOrInsertNode(Node *N, const FoldingSetInfo &Info) {
 //===----------------------------------------------------------------------===//
 // FoldingSetIteratorImpl Implementation
 
-FoldingSetIteratorImpl::FoldingSetIteratorImpl(void **Bucket) {
-  // Skip to the first non-null non-self-cycle bucket.
-  while (*Bucket != reinterpret_cast<void *>(-1) &&
-         (!*Bucket || !GetNextPtr(*Bucket)))
-    ++Bucket;
-
-  NodePtr = static_cast<FoldingSetNode *>(*Bucket);
+FoldingSetIteratorImpl::FoldingSetIteratorImpl(const FoldingSetBase *Set,
+                                               unsigned Index)
+    : DebugEpochBase::HandleBase(Set), Set(Set), Index(Index) {
+  while (this->Index < Set->NumBuckets && !Set->Buckets[this->Index])
+    ++this->Index;
 }
 
 void FoldingSetIteratorImpl::advance() {
-  // If there is another link within this bucket, go to it.
-  void *Probe = NodePtr->getNextInBucket();
-
-  if (FoldingSetNode *NextNodeInBucket = GetNextPtr(Probe))
-    NodePtr = NextNodeInBucket;
-  else {
-    // Otherwise, this is the last link in this bucket.
-    void **Bucket = GetBucketPtr(Probe);
-
-    // Skip to the next non-null non-self-cycle bucket.
-    do {
-      ++Bucket;
-    } while (*Bucket != reinterpret_cast<void *>(-1) &&
-             (!*Bucket || !GetNextPtr(*Bucket)));
-
-    NodePtr = static_cast<FoldingSetNode *>(*Bucket);
-  }
+  assert(isHandleInSync() && "invalid iterator access!");
+  do
+    ++Index;
+  while (Index < Set->NumBuckets && !Set->Buckets[Index]);
 }
diff --git a/llvm/unittests/ADT/FoldingSet.cpp b/llvm/unittests/ADT/FoldingSet.cpp
index f02fbdb0d459e..cc479be708245 100644
--- a/llvm/unittests/ADT/FoldingSet.cpp
+++ b/llvm/unittests/ADT/FoldingSet.cpp
@@ -13,6 +13,10 @@
 #include "llvm/ADT/FoldingSet.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
+#include <map>
+#include <memory>
+#include <random>
+#include <set>
 #include <string>
 
 using namespace llvm;
@@ -358,4 +362,99 @@ TEST(FoldingSetTest, ContextualFoldingSetBasic) {
   EXPECT_THAT(Set, SizeIs(0));
 }
 
+// Exercise growth, and the Algorithm R shifting that erase performs, against a
+// reference model. Nothing else in this file inserts enough nodes to rehash.
+TEST(FoldingSetTest, InsertEraseStress) {
+  FoldingSet<TrivialPair> Set;
+  std::map<unsigned, std::unique_ptr<TrivialPair>> Model;
+  std::mt19937 Rng(42);
+  for (unsigned Op = 0; Op != 1000; ++Op) {
+    unsigned Key = Rng() % 4096;
+    FoldingSetNodeID ID;
+    ID.AddInteger(Key);
+    ID.AddInteger(Key);
+
+    auto It = Model.find(Key);
+    if (Rng() & 1) {
+      void *InsertPos = nullptr;
+      TrivialPair *Found = Set.FindNodeOrInsertPos(ID, InsertPos);
+      if (It != Model.end()) {
+        ASSERT_EQ(It->second.get(), Found);
+        continue;
+      }
+      ASSERT_EQ(nullptr, Found);
+      auto N = std::make_unique<TrivialPair>(Key, Key);
+      Set.InsertNode(N.get(), InsertPos);
+      Model.emplace(Key, std::move(N));
+    } else if (It != Model.end()) {
+      ASSERT_TRUE(Set.RemoveNode(It->second.get()));
+      ASSERT_FALSE(Set.RemoveNode(It->second.get()));
+      Model.erase(It);
+    }
+    ASSERT_EQ(Model.size(), Set.size());
+  }
+
+  // Every surviving node must still be reachable along its probe chain, and
+  // iteration must visit each of them exactly once.
+  for (const auto &KV : Model) {
+    FoldingSetNodeID ID;
+    ID.AddInteger(KV.first);
+    ID.AddInteger(KV.first);
+    void *InsertPos = nullptr;
+    EXPECT_EQ(KV.second.get(), Set.FindNodeOrInsertPos(ID, InsertPos));
+  }
+  std::set<TrivialPair *> Visited;
+  for (TrivialPair &N : Set)
+    EXPECT_TRUE(Visited.insert(&N).second);
+  EXPECT_EQ(Model.size(), Visited.size());
+}
+
+#if LLVM_ENABLE_ABI_BREAKING_CHECKS
+TEST(FoldingSetTest, InsertInvalidatesIterators) {
+  FoldingSet<TrivialPair> Set;
+  TrivialPair T1(1, 1), T2(2, 2);
+  Set.InsertNode(&T1);
+  auto It = Set.begin();
+  Set.InsertNode(&T2);
+  EXPECT_DEATH((void)It->Value, "invalid iterator access");
+}
+
+TEST(FoldingSetTest, RemoveInvalidatesIterators) {
+  FoldingSet<TrivialPair> Set;
+  TrivialPair T1(1, 1), T2(2, 2);
+  Set.InsertNode(&T1);
+  Set.InsertNode(&T2);
+  auto It = Set.begin();
+  Set.RemoveNode(&T2);
+  EXPECT_DEATH((void)It->Value, "invalid iterator access");
+}
+
+TEST(FoldingSetTest, RemoveOfAbsentNodeKeepsIterators) {
+  FoldingSet<TrivialPair> Set;
+  TrivialPair T1(1, 1), Absent(2, 2);
+  Set.InsertNode(&T1);
+  auto It = Set.begin();
+  EXPECT_FALSE(Set.RemoveNode(&Absent));
+  EXPECT_EQ(&T1, &*It);
+}
+
+TEST(FoldingSetTest, ClearInvalidatesIterators) {
+  FoldingSet<TrivialPair> Set;
+  TrivialPair T1(1, 1);
+  Set.InsertNode(&T1);
+  auto It = Set.begin();
+  Set.clear();
+  EXPECT_DEATH((void)It->Value, "invalid iterator access");
+}
+
+TEST(FoldingSetTest, MoveInvalidatesIterators) {
+  FoldingSet<TrivialPair> Set;
+  TrivialPair T1(1, 1);
+  Set.InsertNode(&T1);
+  auto It = Set.begin();
+  FoldingSet<TrivialPair> Other(std::move(Set));
+  EXPECT_DEATH((void)It->Value, "invalid iterator access");
+}
+#endif
+
 } // namespace



More information about the llvm-commits mailing list