[sanitizer] Optimize DenseMap::{find,erase}

[MaskRay]

Port the ADT optimization #100517. In addition, add contains
(https://reviews.llvm.org/D145895) while changing count.

[llvmbot]

@llvm/pr-subscribers-compiler-rt-sanitizer

Author: Fangrui Song (MaskRay)

Changes

Port the ADT optimization #100517. In addition, add contains
(https://reviews.llvm.org/D145895) while changing count.

---

Full diff: https://github.com/llvm/llvm-project/pull/101785.diff

1 Files Affected:

• (modified) compiler-rt/lib/sanitizer_common/sanitizer_dense_map.h (+44-28)

diff --git a/compiler-rt/lib/sanitizer_common/sanitizer_dense_map.h b/compiler-rt/lib/sanitizer_common/sanitizer_dense_map.h
index 046d77dddc9c1..5c0b2c034925a 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_dense_map.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_dense_map.h
@@ -69,25 +69,19 @@ class DenseMapBase {
     setNumTombstones(0);
   }
 
-  /// Return 1 if the specified key is in the map, 0 otherwise.
-  size_type count(const KeyT &Key) const {
-    const BucketT *TheBucket;
-    return LookupBucketFor(Key, TheBucket) ? 1 : 0;
+  /// Return true if the specified key is in the map, false otherwise.
+  bool contains(const_arg_type_t<KeyT> Val) const {
+    return doFind(Val) != nullptr;
   }
 
-  value_type *find(const KeyT &Key) {
-    BucketT *TheBucket;
-    if (LookupBucketFor(Key, TheBucket))
-      return TheBucket;
-    return nullptr;
-  }
-  const value_type *find(const KeyT &Key) const {
-    const BucketT *TheBucket;
-    if (LookupBucketFor(Key, TheBucket))
-      return TheBucket;
-    return nullptr;
+  /// Return 1 if the specified key is in the map, 0 otherwise.
+  size_type count(const_arg_type_t<KeyT> Val) const {
+    return contains(Val) ? 1 : 0;
   }
 
+  value_type *find(const KeyT &Key) { return doFind(Key); }
+  const value_type *find(const KeyT &Key) const { return doFind(Key); }
+
   /// Alternate version of find() which allows a different, and possibly
   /// less expensive, key type.
   /// The DenseMapInfo is responsible for supplying methods
@@ -95,25 +89,18 @@ class DenseMapBase {
   /// type used.
   template <class LookupKeyT>
   value_type *find_as(const LookupKeyT &Key) {
-    BucketT *TheBucket;
-    if (LookupBucketFor(Key, TheBucket))
-      return TheBucket;
-    return nullptr;
+    return doFind(Key);
   }
   template <class LookupKeyT>
   const value_type *find_as(const LookupKeyT &Key) const {
-    const BucketT *TheBucket;
-    if (LookupBucketFor(Key, TheBucket))
-      return TheBucket;
-    return nullptr;
+    return doFind(Key);
   }
 
   /// lookup - Return the entry for the specified key, or a default
   /// constructed value if no such entry exists.
   ValueT lookup(const KeyT &Key) const {
-    const BucketT *TheBucket;
-    if (LookupBucketFor(Key, TheBucket))
-      return TheBucket->getSecond();
+    if (const BucketT *Bucket = doFind(Key))
+      return Bucket->getSecond();
     return ValueT();
   }
 
@@ -184,8 +171,8 @@ class DenseMapBase {
   }
 
   bool erase(const KeyT &Val) {
-    BucketT *TheBucket;
-    if (!LookupBucketFor(Val, TheBucket))
+    BucketT *TheBucket = doFind(Val);
+    if (!TheBucket)
       return false;  // not in map.
 
     TheBucket->getSecond().~ValueT();
@@ -449,6 +436,35 @@ class DenseMapBase {
     return TheBucket;
   }
 
+  template <typename LookupKeyT>
+  BucketT *doFind(const LookupKeyT &Val) {
+    BucketT *BucketsPtr = getBuckets();
+    const unsigned NumBuckets = getNumBuckets();
+    if (NumBuckets == 0)
+      return nullptr;
+
+    const KeyT EmptyKey = getEmptyKey();
+    unsigned BucketNo = getHashValue(Val) & (NumBuckets - 1);
+    unsigned ProbeAmt = 1;
+    while (true) {
+      BucketT *Bucket = BucketsPtr + BucketNo;
+      if (LIKELY(KeyInfoT::isEqual(Val, Bucket->getFirst())))
+        return Bucket;
+      if (LIKELY(KeyInfoT::isEqual(Bucket->getFirst(), EmptyKey)))
+        return nullptr;
+
+      // Otherwise, it's a hash collision or a tombstone, continue quadratic
+      // probing.
+      BucketNo += ProbeAmt++;
+      BucketNo &= NumBuckets - 1;
+    }
+  }
+
+  template <typename LookupKeyT>
+  const BucketT *doFind(const LookupKeyT &Val) const {
+    return const_cast<DenseMapBase *>(this)->doFind(Val);
+  }
+
   /// LookupBucketFor - Lookup the appropriate bucket for Val, returning it in
   /// FoundBucket.  If the bucket contains the key and a value, this returns
   /// true, otherwise it returns a bucket with an empty marker or tombstone and