[compiler-rt] 65b4a77 - [sanitizer] Optimize DenseMap::{find,erase}

[via llvm-commits]

Author: Fangrui Song
Date: 2024-08-07T22:02:32-07:00
New Revision: 65b4a7704beaf87c5eb12b672a837501a981a77a

URL: https://github.com/llvm/llvm-project/commit/65b4a7704beaf87c5eb12b672a837501a981a77a
DIFF: https://github.com/llvm/llvm-project/commit/65b4a7704beaf87c5eb12b672a837501a981a77a.diff

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

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

Pull Request: https://github.com/llvm/llvm-project/pull/101785

Added: 
    

Modified: 
    compiler-rt/lib/sanitizer_common/sanitizer_dense_map.h

Removed: 
    


################################################################################
diff  --git a/compiler-rt/lib/sanitizer_common/sanitizer_dense_map.h b/compiler-rt/lib/sanitizer_common/sanitizer_dense_map.h
index 046d77dddc9c1..c63788653de75 100644
--- a/compiler-rt/lib/sanitizer_common/sanitizer_dense_map.h
+++ b/compiler-rt/lib/sanitizer_common/sanitizer_dense_map.h
@@ -69,24 +69,14 @@ class DenseMapBase {
     setNumTombstones(0);
   }
 
+  /// Return true if the specified key is in the map, false otherwise.
+  bool contains(const KeyT &Key) const { return doFind(Key) != nullptr; }
+
   /// 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;
-  }
+  size_type count(const KeyT &Key) const { return contains(Key) ? 1 : 0; }
 
-  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;
-  }
+  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 
diff erent, and possibly
   /// less expensive, key type.
@@ -95,25 +85,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 +167,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 +432,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