[SmallPtrSet] Optimize find/erase

Port #100517 for DenseMap.

Pull Request: #104740

Author: MaskRay

llvm/include/llvm/ADT/SmallPtrSet.h

@@ -185,8 +185,8 @@ class SmallPtrSetImplBase : public DebugEpochBase {
       return false;
     }
 
-    auto *Bucket = FindBucketFor(Ptr);
-    if (*Bucket != Ptr)
+    auto *Bucket = doFind(Ptr);
+    if (!Bucket)
       return false;
 
     *const_cast<const void **>(Bucket) = getTombstoneMarker();
@@ -211,8 +211,7 @@ class SmallPtrSetImplBase : public DebugEpochBase {
     }
 
     // Big set case.
-    auto *Bucket = FindBucketFor(Ptr);
-    if (*Bucket == Ptr)
+    if (auto *Bucket = doFind(Ptr))
       return Bucket;
     return EndPointer();
   }
@@ -222,6 +221,7 @@ class SmallPtrSetImplBase : public DebugEpochBase {
 private:
   std::pair<const void *const *, bool> insert_imp_big(const void *Ptr);
 
+  const void *const *doFind(const void *Ptr) const;
   const void * const *FindBucketFor(const void *Ptr) const;
   void shrink_and_clear();
 

llvm/lib/Support/SmallPtrSet.cpp

@@ -62,7 +62,25 @@ SmallPtrSetImplBase::insert_imp_big(const void *Ptr) {
   return std::make_pair(Bucket, true);
 }
 
-const void * const *SmallPtrSetImplBase::FindBucketFor(const void *Ptr) const {
+const void *const *SmallPtrSetImplBase::doFind(const void *Ptr) const {
+  unsigned BucketNo =
+      DenseMapInfo<void *>::getHashValue(Ptr) & (CurArraySize - 1);
+  unsigned ProbeAmt = 1;
+  while (true) {
+    const void *const *Bucket = CurArray + BucketNo;
+    if (LLVM_LIKELY(*Bucket == Ptr))
+      return Bucket;
+    if (LLVM_LIKELY(*Bucket == getEmptyMarker()))
+      return nullptr;
+
+    // Otherwise, it's a hash collision or a tombstone, continue quadratic
+    // probing.
+    BucketNo += ProbeAmt++;
+    BucketNo &= CurArraySize - 1;
+  }
+}
+
+const void *const *SmallPtrSetImplBase::FindBucketFor(const void *Ptr) const {
   unsigned Bucket = DenseMapInfo<void *>::getHashValue(Ptr) & (CurArraySize-1);
   unsigned ArraySize = CurArraySize;
   unsigned ProbeAmt = 1;