From f5d2a02700f9668a5f8bd5be3e976985e9a4df0e Mon Sep 17 00:00:00 2001
From: Fangrui Song <i@maskray.me>
Date: Sun, 18 Aug 2024 23:52:13 -0700
Subject: [PATCH] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20initia?=
 =?UTF-8?q?l=20version?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Created using spr 1.3.5-bogner
---
 llvm/include/llvm/ADT/SmallPtrSet.h |  8 ++++----
 llvm/lib/Support/SmallPtrSet.cpp    | 20 +++++++++++++++++++-
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/llvm/include/llvm/ADT/SmallPtrSet.h b/llvm/include/llvm/ADT/SmallPtrSet.h
index e668df08c3e636..1fc2318342ae78 100644
--- a/llvm/include/llvm/ADT/SmallPtrSet.h
+++ b/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();
 
diff --git a/llvm/lib/Support/SmallPtrSet.cpp b/llvm/lib/Support/SmallPtrSet.cpp
index cbb87ea8717cfc..1bec911f39b13e 100644
--- a/llvm/lib/Support/SmallPtrSet.cpp
+++ b/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;