From babd5980d8a8b4aad9814212a269f6197ebf1f2e Mon Sep 17 00:00:00 2001
From: Talin <viridia@gmail.com>
Date: Mon, 30 Jan 2012 06:55:43 +0000
Subject: [PATCH] DenseMap::find_as() and unit tests.

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@149229 91177308-0d34-0410-b5e6-96231b3b80d8
---
 docs/ProgrammersManual.html    | 10 ++++++++-
 include/llvm/ADT/DenseMap.h    | 29 ++++++++++++++++++++++--
 unittests/ADT/DenseMapTest.cpp | 41 ++++++++++++++++++++++++++++++++++
 3 files changed, 77 insertions(+), 3 deletions(-)

diff --git a/docs/ProgrammersManual.html b/docs/ProgrammersManual.html
index 51862034b73..0b1a875da1a 100644
--- a/docs/ProgrammersManual.html
+++ b/docs/ProgrammersManual.html
@@ -1753,7 +1753,7 @@ pointers, or map other small types to each other.
 
 <p>
 There are several aspects of DenseMap that you should be aware of, however.  The
-iterators in a densemap are invalidated whenever an insertion occurs, unlike
+iterators in a DenseMap are invalidated whenever an insertion occurs, unlike
 map.  Also, because DenseMap allocates space for a large number of key/value
 pairs (it starts with 64 by default), it will waste a lot of space if your keys
 or values are large.  Finally, you must implement a partial specialization of
@@ -1761,6 +1761,14 @@ DenseMapInfo for the key that you want, if it isn't already supported.  This
 is required to tell DenseMap about two special marker values (which can never be
 inserted into the map) that it needs internally.</p>
 
+<p>
+DenseMap's find_as() method supports lookup operations using an alternate key
+type. This is useful in cases where the normal key type is expensive to
+construct, but cheap to compare against. The DenseMapInfo is responsible for
+defining the appropriate comparison and hashing methods for each alternate
+key type used.
+</p>
+
 </div>
 
 <!-- _______________________________________________________________________ -->
diff --git a/include/llvm/ADT/DenseMap.h b/include/llvm/ADT/DenseMap.h
index 2ba6db78cc4..672147dd4d1 100644
--- a/include/llvm/ADT/DenseMap.h
+++ b/include/llvm/ADT/DenseMap.h
@@ -147,6 +147,26 @@ public:
     return end();
   }
 
+  /// Alternate version of find() which allows a different, and possibly
+  /// less expensive, key type.
+  /// The DenseMapInfo is responsible for supplying methods
+  /// getHashValue(LookupKeyT) and isEqual(LookupKeyT, KeyT) for each key
+  /// type used.
+  template<class LookupKeyT>
+  iterator find_as(const LookupKeyT &Val) {
+    BucketT *TheBucket;
+    if (LookupBucketFor(Val, TheBucket))
+      return iterator(TheBucket, Buckets+NumBuckets, true);
+    return end();
+  }
+  template<class LookupKeyT>
+  const_iterator find_as(const LookupKeyT &Val) const {
+    BucketT *TheBucket;
+    if (LookupBucketFor(Val, TheBucket))
+      return const_iterator(TheBucket, Buckets+NumBuckets, true);
+    return end();
+  }
+
   /// lookup - Return the entry for the specified key, or a default
   /// constructed value if no such entry exists.
   ValueT lookup(const KeyT &Val) const {
@@ -309,6 +329,10 @@ private:
   static unsigned getHashValue(const KeyT &Val) {
     return KeyInfoT::getHashValue(Val);
   }
+  template<typename LookupKeyT>
+  static unsigned getHashValue(const LookupKeyT &Val) {
+    return KeyInfoT::getHashValue(Val);
+  }
   static const KeyT getEmptyKey() {
     return KeyInfoT::getEmptyKey();
   }
@@ -320,7 +344,8 @@ private:
   /// 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
   /// returns false.
-  bool LookupBucketFor(const KeyT &Val, BucketT *&FoundBucket) const {
+  template<typename LookupKeyT>
+  bool LookupBucketFor(const LookupKeyT &Val, BucketT *&FoundBucket) const {
     unsigned BucketNo = getHashValue(Val);
     unsigned ProbeAmt = 1;
     BucketT *BucketsPtr = Buckets;
@@ -341,7 +366,7 @@ private:
     while (1) {
       BucketT *ThisBucket = BucketsPtr + (BucketNo & (NumBuckets-1));
       // Found Val's bucket?  If so, return it.
-      if (KeyInfoT::isEqual(ThisBucket->first, Val)) {
+      if (KeyInfoT::isEqual(Val, ThisBucket->first)) {
         FoundBucket = ThisBucket;
         return true;
       }
diff --git a/unittests/ADT/DenseMapTest.cpp b/unittests/ADT/DenseMapTest.cpp
index afac651a6b2..e0ee7782ccb 100644
--- a/unittests/ADT/DenseMapTest.cpp
+++ b/unittests/ADT/DenseMapTest.cpp
@@ -176,4 +176,45 @@ TEST_F(DenseMapTest, ConstIteratorTest) {
   EXPECT_TRUE(cit == cit2);
 }
 
+// Key traits that allows lookup with either an unsigned or char* key;
+// In the latter case, "a" == 0, "b" == 1 and so on.
+struct TestDenseMapInfo {
+  static inline unsigned getEmptyKey() { return ~0; }
+  static inline unsigned getTombstoneKey() { return ~0U - 1; }
+  static unsigned getHashValue(const unsigned& Val) { return Val * 37U; }
+  static unsigned getHashValue(const char* Val) {
+    return (unsigned)(Val[0] - 'a') * 37U;
+  }
+  static bool isEqual(const unsigned& LHS, const unsigned& RHS) {
+    return LHS == RHS;
+  }
+  static bool isEqual(const char* LHS, const unsigned& RHS) {
+    return (unsigned)(LHS[0] - 'a') == RHS;
+  }
+};
+
+// find_as() tests
+TEST_F(DenseMapTest, FindAsTest) {
+  DenseMap<unsigned, unsigned, TestDenseMapInfo> map;
+  map[0] = 1;
+  map[1] = 2;
+  map[2] = 3;
+
+  // Size tests
+  EXPECT_EQ(3u, map.size());
+
+  // Normal lookup tests
+  EXPECT_EQ(1, map.count(1));
+  EXPECT_EQ(1u, map.find(0)->second);
+  EXPECT_EQ(2u, map.find(1)->second);
+  EXPECT_EQ(3u, map.find(2)->second);
+  EXPECT_TRUE(map.find(3) == map.end());
+
+  // find_as() tests
+  EXPECT_EQ(1u, map.find_as("a")->second);
+  EXPECT_EQ(2u, map.find_as("b")->second);
+  EXPECT_EQ(3u, map.find_as("c")->second);
+  EXPECT_TRUE(map.find_as("d") == map.end());
+}
+
 }