2 files changed, 59 insertions, 5 deletions

diff --git a/include/llvm/ADT/StringMap.h b/include/llvm/ADT/StringMap.h
index 52982589f7..7cc9ce7b6c 100644
--- a/include/llvm/ADT/StringMap.h
+++ b/include/llvm/ADT/StringMap.h
@@ -64,6 +64,11 @@ protected:
   /// of the string.
   unsigned LookupBucketFor(const char *KeyStart, const char *KeyEnd);
   
+  /// FindKey - Look up the bucket that contains the specified key. If it exists
+  /// in the map, return the bucket number of the key.  Otherwise return -1.
+  /// This does not modify the map.
+  int FindKey(const char *KeyStart, const char *KeyEnd) const;
+  
 public:
   static StringMapEntryBase *getTombstoneVal() {
     return (StringMapEntryBase*)-1;
@@ -175,11 +180,17 @@ public:
   const_iterator begin() const { return const_iterator(TheTable); }
   const_iterator end() const { return const_iterator(TheTable+NumBuckets); }
   
-  /// FindValue - Look up the specified key in the map.  If it exists, return a
-  /// pointer to the element, otherwise return null.
-  MapEntryTy *FindValue(const char *KeyStart, const char *KeyEnd) {
-    unsigned BucketNo = LookupBucketFor(KeyStart, KeyEnd);
-    return static_cast<MapEntryTy*>(TheTable[BucketNo].Item);
+  
+  iterator find(const char *KeyStart, const char *KeyEnd) {
+    int Bucket = FindKey(KeyStart, KeyEnd);
+    if (Bucket == -1) return end();
+    return iterator(TheTable+Bucket);
+  }
+
+  const_iterator find(const char *KeyStart, const char *KeyEnd) const {
+    int Bucket = FindKey(KeyStart, KeyEnd);
+    if (Bucket == -1) return end();
+    return const_iterator(TheTable+Bucket);
   }
   
   /// GetOrCreateValue - Look up the specified key in the table.  If a value
diff --git a/lib/Support/StringMap.cpp b/lib/Support/StringMap.cpp
index 5a3896424e..02d42b9b49 100644
--- a/lib/Support/StringMap.cpp
+++ b/lib/Support/StringMap.cpp
@@ -91,6 +91,49 @@ unsigned StringMapImpl::LookupBucketFor(const char *NameStart,
   }
 }
 
+
+/// FindKey - Look up the bucket that contains the specified key. If it exists
+/// in the map, return the bucket number of the key.  Otherwise return -1.
+/// This does not modify the map.
+int StringMapImpl::FindKey(const char *KeyStart, const char *KeyEnd) const {
+  unsigned HTSize = NumBuckets;
+  unsigned FullHashValue = HashString(KeyStart, KeyEnd);
+  unsigned BucketNo = FullHashValue & (HTSize-1);
+  
+  unsigned ProbeAmt = 1;
+  while (1) {
+    ItemBucket &Bucket = TheTable[BucketNo];
+    StringMapEntryBase *BucketItem = Bucket.Item;
+    // If we found an empty bucket, this key isn't in the table yet, return.
+    if (BucketItem == 0)
+      return -1;
+    
+    // If the full hash value matches, check deeply for a match.  The common
+    // case here is that we are only looking at the buckets (for item info
+    // being non-null and for the full hash value) not at the items.  This
+    // is important for cache locality.
+    if (Bucket.FullHashValue == FullHashValue) {
+      // Do the comparison like this because NameStart isn't necessarily
+      // null-terminated!
+      char *ItemStr = (char*)BucketItem+ItemSize;
+      unsigned ItemStrLen = BucketItem->getKeyLength();
+      if (unsigned(KeyEnd-KeyStart) == ItemStrLen &&
+          memcmp(ItemStr, KeyStart, ItemStrLen) == 0) {
+        // We found a match!
+        return BucketNo;
+      }
+    }
+    
+    // Okay, we didn't find the item.  Probe to the next bucket.
+    BucketNo = (BucketNo+ProbeAmt) & (HTSize-1);
+    
+    // Use quadratic probing, it has fewer clumping artifacts than linear
+    // probing and has good cache behavior in the common case.
+    ++ProbeAmt;
+  }
+}
+
+
 /// RehashTable - Grow the table, redistributing values into the buckets with
 /// the appropriate mod-of-hashtable-size.
 void StringMapImpl::RehashTable() {