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

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 different, 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