[FoldingSet] Add typed lookup/insert/erase/getOrInsert (#219644)

FindNodeOrInsertPos hands the insertion state to InsertNode as a `void *`:
any pointer converts to it, only a comment says where it may come from,
and since the switch to linear probing the value is the node's hash, not
a position at all.

Carry it in FoldingSetInsertToken instead, the hash wrapped in a type
whose constructor and accessor are private to FoldingSetBase, with
NotAHash as the no-token state. The `void *` overloads remain,
forwarding to the typed ones, until their callers are migrated.

LLM-aided

Author: MaskRay

lld/ELF/SyntheticSections.cpp

@@ -2875,14 +2875,14 @@ void DebugNamesBaseSection::computeHdrAndAbbrevTable(
         FoldingSetNodeID id;
         abbrev.Profile(id);
         uint32_t newCode;
-        void *insertPos;
-        if (Abbrev *existing = abbrevSet.FindNodeOrInsertPos(id, insertPos)) {
+        FoldingSetInsertToken token;
+        if (Abbrev *existing = abbrevSet.lookup(id, token)) {
           // Found it; we've already seen an identical abbreviation.
           newCode = existing->code;
         } else {
           Abbrev *abbrev2 =
               new (abbrevAlloc.Allocate()) Abbrev(std::move(abbrev));
-          abbrevSet.InsertNode(abbrev2, insertPos);
+          abbrevSet.insert(abbrev2, token);
           abbrevTable.push_back(abbrev2);
           newCode = abbrevTable.size();
           abbrev2->code = newCode;

llvm/include/llvm/ADT/FoldingSet.h

@@ -74,38 +74,37 @@ namespace llvm {
 /// 1) If you have an existing node that you want add to the set but unsure
 /// that the node might already exist then call;
 ///
-///    MyNode *M = MyFoldingSet.GetOrInsertNode(N);
+///    MyNode *M = MyFoldingSet.getOrInsert(N);
 ///
 /// If The result is equal to the input then the node has been inserted.
 /// Otherwise, the result is the node existing in the folding set, and the
 /// input can be discarded (use the result instead.)
 ///
 /// 2) If you are ready to construct a node but want to check if it already
-/// exists, then call FindNodeOrInsertPos with a FoldingSetNodeID of the bits to
-/// check;
+/// exists, then call lookup with a FoldingSetNodeID of the bits to check;
 ///
 ///   FoldingSetNodeID ID;
 ///   ID.AddString(Name);
 ///   ID.AddInteger(Value);
-///   void *InsertPoint;
+///   FoldingSetInsertToken Token;
 ///
-///    MyNode *M = MyFoldingSet.FindNodeOrInsertPos(ID, InsertPoint);
+///    MyNode *M = MyFoldingSet.lookup(ID, Token);
 ///
-/// If found then M will be non-NULL, else InsertPoint will point to where it
-/// should be inserted using InsertNode.
+/// If found then M will be non-NULL, else Token holds what insert needs to
+/// place the node.
 ///
-/// 3) If you get a NULL result from FindNodeOrInsertPos then you can insert a
-/// new node with InsertNode;
+/// 3) If you get a NULL result from lookup then you can insert a new node with
+/// insert;
 ///
 ///    MyNode *N = new MyNode(Name, Value);
-///    MyFoldingSet.InsertNode(N, InsertPoint);
+///    MyFoldingSet.insert(N, Token);
 ///
-/// InsertPoint survives intervening insertions, but N must profile identically
-/// to the ID that produced it, or N becomes unfindable.
+/// Token survives intervening insertions, but N must profile identically to
+/// the ID that produced it, or N becomes unfindable.
 ///
 /// 4) Finally, if you want to remove a node from the folding set call;
 ///
-///    bool WasRemoved = MyFoldingSet.RemoveNode(M);
+///    bool WasRemoved = MyFoldingSet.erase(M);
 ///
 /// The result indicates whether the node existed in the folding set.
 
@@ -177,8 +176,8 @@ 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.
-  // Never returns NotAHash: FoldingSetBase uses it to keep the InsertPos token
-  // non-null and to mark a node belonging to no set.
+  // Never returns NotAHash: FoldingSetBase reserves it for the empty insert
+  // token and for a node belonging to no set.
   unsigned ComputeHash() const {
     unsigned Hash =
         static_cast<unsigned>(hash_combine_range(Data, Data + Size));
@@ -288,6 +287,31 @@ class FoldingSetNodeID {
   LLVM_ABI FoldingSetNodeIDRef Intern(BumpPtrAllocator &Allocator) const;
 };
 
+/// Insertion token: a failed lookup fills it in, the matching insert consumes
+/// it.
+class FoldingSetInsertToken {
+  uint32_t Hash = FoldingSetNodeIDRef::NotAHash;
+
+  explicit FoldingSetInsertToken(uint32_t Hash) : Hash(Hash) {
+    assert(Hash != FoldingSetNodeIDRef::NotAHash && "Invalid insert token");
+  }
+
+  friend class FoldingSetBase;
+
+public:
+  FoldingSetInsertToken() = default;
+  explicit operator bool() const {
+    return Hash != FoldingSetNodeIDRef::NotAHash;
+  }
+
+  friend bool operator==(FoldingSetInsertToken A, FoldingSetInsertToken B) {
+    return A.Hash == B.Hash;
+  }
+  friend bool operator!=(FoldingSetInsertToken A, FoldingSetInsertToken B) {
+    return !(A == B);
+  }
+};
+
 //===----------------------------------------------------------------------===//
 /// Non-templated base class for FoldingSet and ContextualFoldingSet, holding
 /// the memory management and probing that does not depend on the node type.
@@ -379,15 +403,19 @@ class FoldingSetBase : public DebugEpochBase {
   /// return it.  Otherwise, insert \p N and return it instead.
   LLVM_ABI Node *GetOrInsertNode(Node *N, const FoldingSetInfo &Info);
 
-  /// Look up the node specified by ID.  If it exists, return it.  If not,
-  /// return the insertion token that will make insertion faster.
+  /// Look up the node specified by ID. If it exists, return it and clear
+  /// \p Token; otherwise return null and set \p Token for a subsequent insert.
+  LLVM_ABI Node *lookup(const FoldingSetNodeID &ID,
+                        FoldingSetInsertToken &Token,
+                        const FoldingSetInfo &Info);
   LLVM_ABI Node *FindNodeOrInsertPos(const FoldingSetNodeID &ID,
                                      void *&InsertPos,
                                      const FoldingSetInfo &Info);
 
-  /// Insert the specified node into the folding set, knowing that
-  /// it is not already in the folding set.  InsertPos must be obtained from
-  /// FindNodeOrInsertPos for an ID that \p N profiles identically to.
+  /// Insert the specified node into the folding set, knowing that it is not
+  /// already in the folding set.  \p Token must come from lookup for an ID that
+  /// \p N profiles identically to.
+  LLVM_ABI void insert(Node *N, FoldingSetInsertToken Token);
   LLVM_ABI void InsertNode(Node *N, void *InsertPos);
 };
 
@@ -476,36 +504,46 @@ class FoldingSetImpl : public FoldingSetBase, public Trait::ContextStorage {
 
   /// Remove a node from the folding set, returning true if one
   /// was removed or false if the node was not in the folding set.
-  bool RemoveNode(T *N) { return FoldingSetBase::RemoveNode(N); }
+  bool erase(T *N) { return FoldingSetBase::RemoveNode(N); }
+  bool RemoveNode(T *N) { return erase(N); }
 
   /// If there is an existing node exactly equal to the specified node,
   /// return it.  Otherwise, insert 'N' and return it instead.
-  T *GetOrInsertNode(T *N) {
+  T *getOrInsert(T *N) {
     return static_cast<T *>(
         FoldingSetBase::GetOrInsertNode(N, getFoldingSetInfo()));
   }
+  T *GetOrInsertNode(T *N) { return getOrInsert(N); }
 
-  /// Look up the node specified by ID.  If it exists, return it.  If not,
-  /// return the insertion token that will make insertion faster.
+  /// Look up the node specified by ID. If it exists, return it and clear
+  /// \p Token; otherwise return null and set \p Token for a subsequent insert.
+  T *lookup(const FoldingSetNodeID &ID, FoldingSetInsertToken &Token) {
+    return static_cast<T *>(
+        FoldingSetBase::lookup(ID, Token, getFoldingSetInfo()));
+  }
   T *FindNodeOrInsertPos(const FoldingSetNodeID &ID, void *&InsertPos) {
     return static_cast<T *>(FoldingSetBase::FindNodeOrInsertPos(
         ID, InsertPos, getFoldingSetInfo()));
   }
 
-  /// Insert the specified node into the folding set, knowing that
-  /// it is not already in the folding set.  InsertPos must be obtained from
-  /// FindNodeOrInsertPos.
+  /// Insert the specified node into the folding set, knowing that it is not
+  /// already in the folding set.  \p Token must come from lookup for an ID that
+  /// \p N profiles identically to.
+  void insert(T *N, FoldingSetInsertToken Token) {
+    FoldingSetBase::insert(N, Token);
+  }
   void InsertNode(T *N, void *InsertPos) {
     FoldingSetBase::InsertNode(N, InsertPos);
   }
 
   /// Insert the specified node into the folding set, knowing that it is not
   /// already in the folding set.
-  void InsertNode(T *N) {
-    T *Inserted = GetOrInsertNode(N);
+  void insert(T *N) {
+    T *Inserted = getOrInsert(N);
     (void)Inserted;
     assert(Inserted == N && "Node already inserted!");
   }
+  void InsertNode(T *N) { insert(N); }
 };
 
 //===----------------------------------------------------------------------===//
@@ -560,35 +598,44 @@ template <class T, class VectorT = SmallVector<T *, 8>> class FoldingSetVector {
     Vector.clear();
   }
 
-  /// Look up the node specified by ID.  If it exists, return it.  If not,
-  /// return the insertion token that will make insertion faster.
+  /// Look up the node specified by ID. If it exists, return it and clear
+  /// \p Token; otherwise return null and set \p Token for a subsequent insert.
+  T *lookup(const FoldingSetNodeID &ID, FoldingSetInsertToken &Token) {
+    return Set.lookup(ID, Token);
+  }
   T *FindNodeOrInsertPos(const FoldingSetNodeID &ID, void *&InsertPos) {
     return Set.FindNodeOrInsertPos(ID, InsertPos);
   }
 
   /// If there is an existing node exactly equal to the specified node,
   /// return it.  Otherwise, insert 'N' and return it instead.
-  T *GetOrInsertNode(T *N) {
-    T *Result = Set.GetOrInsertNode(N);
+  T *getOrInsert(T *N) {
+    T *Result = Set.getOrInsert(N);
     if (Result == N)
       Vector.push_back(N);
     return Result;
   }
+  T *GetOrInsertNode(T *N) { return getOrInsert(N); }
 
-  /// Insert the specified node into the folding set, knowing that
-  /// it is not already in the folding set.  InsertPos must be obtained from
-  /// FindNodeOrInsertPos.
+  /// Insert the specified node into the folding set, knowing that it is not
+  /// already in the folding set.  \p Token must come from lookup for an ID that
+  /// \p N profiles identically to.
+  void insert(T *N, FoldingSetInsertToken Token) {
+    Set.insert(N, Token);
+    Vector.push_back(N);
+  }
   void InsertNode(T *N, void *InsertPos) {
     Set.InsertNode(N, InsertPos);
     Vector.push_back(N);
   }
 
   /// Insert the specified node into the folding set, knowing that
   /// it is not already in the folding set.
-  void InsertNode(T *N) {
-    Set.InsertNode(N);
+  void insert(T *N) {
+    Set.insert(N);
     Vector.push_back(N);
   }
+  void InsertNode(T *N) { insert(N); }
 
   /// Returns the number of nodes in the folding set.
   unsigned size() const { return Set.size(); }

llvm/lib/Support/FoldingSet.cpp

@@ -219,33 +219,46 @@ FoldingSetBase::nodeEquals(const FoldingSetInfo &Info,
   return Info.NodeEquals(Self, N, ID, TempID);
 }
 
-FoldingSetBase::Node *FoldingSetBase::FindNodeOrInsertPos(
-    const FoldingSetNodeID &ID, void *&InsertPos, const FoldingSetInfo &Info) {
+FoldingSetBase::Node *FoldingSetBase::lookup(const FoldingSetNodeID &ID,
+                                             FoldingSetInsertToken &Token,
+                                             const FoldingSetInfo &Info) {
   unsigned IDHash = ID.ComputeHash();
   unsigned Mask = NumBuckets - 1;
   for (unsigned I = IDHash & Mask; Buckets[I]; I = (I + 1) & Mask) {
     Node *N = static_cast<Node *>(Buckets[I]);
     if (N->getFoldingSetHash() == IDHash && nodeEquals(Info, this, N, ID)) {
-      InsertPos = nullptr;
+      Token = {};
       return N;
     }
   }
 
-  InsertPos = encodeHash(IDHash);
+  Token = FoldingSetInsertToken(IDHash);
   return nullptr;
 }
 
-void FoldingSetBase::InsertNode(Node *N, void *InsertPos) {
+FoldingSetBase::Node *FoldingSetBase::FindNodeOrInsertPos(
+    const FoldingSetNodeID &ID, void *&InsertPos, const FoldingSetInfo &Info) {
+  FoldingSetInsertToken Token;
+  Node *N = lookup(ID, Token, Info);
+  InsertPos = Token ? encodeHash(Token.Hash) : nullptr;
+  return N;
+}
+
+void FoldingSetBase::insert(Node *N, FoldingSetInsertToken Token) {
   assert(N && "Cannot insert a null node");
-  assert(InsertPos && "Invalid InsertPos!");
+  assert(Token && "Invalid token!");
   incrementEpoch();
   if (LLVM_UNLIKELY((NumNodes + 1) * 4 > NumBuckets * 3))
     grow(NumBuckets * 2);
-  uint32_t Hash = decodeHash(InsertPos);
+  uint32_t Hash = Token.Hash;
   placeNode(N, Hash);
   N->setFoldingSetHash(Hash);
 }
 
+void FoldingSetBase::InsertNode(Node *N, void *InsertPos) {
+  insert(N, FoldingSetInsertToken(decodeHash(InsertPos)));
+}
+
 bool FoldingSetBase::RemoveNode(Node *N) {
   uint32_t Hash = N->getFoldingSetHash();
   if (Hash == FoldingSetNodeIDRef::NotAHash)
@@ -280,9 +293,9 @@ FoldingSetBase::Node *
 FoldingSetBase::GetOrInsertNode(Node *N, const FoldingSetInfo &Info) {
   FoldingSetNodeID ID;
   Info.GetNodeProfile(this, N, ID);
-  void *IP;
-  if (Node *E = FindNodeOrInsertPos(ID, IP, Info))
+  FoldingSetInsertToken Token;
+  if (Node *E = lookup(ID, Token, Info))
     return E;
-  InsertNode(N, IP);
+  insert(N, Token);
   return N;
 }