2 files changed, 79 insertions, 0 deletions

diff --git a/llvm/include/llvm/ADT/DenseMapInfo.h b/llvm/include/llvm/ADT/DenseMapInfo.h
index a844ebcccf5..18c692e0cbc 100644
--- a/llvm/include/llvm/ADT/DenseMapInfo.h
+++ b/llvm/include/llvm/ADT/DenseMapInfo.h
@@ -30,6 +30,36 @@ struct DenseMapInfo {
   //static bool isEqual(const T &LHS, const T &RHS);
 };
 
+template <typename T> struct CachedHash {
+  CachedHash(T Val) : Val(std::move(Val)) {
+    Hash = DenseMapInfo<T>::getHashValue(Val);
+  }
+  CachedHash(T Val, unsigned Hash) : Val(std::move(Val)), Hash(Hash) {}
+  T Val;
+  unsigned Hash;
+};
+
+// Provide DenseMapInfo for all CachedHash<T>.
+template <typename T> struct DenseMapInfo<CachedHash<T>> {
+  static CachedHash<T> getEmptyKey() {
+    T N = DenseMapInfo<T>::getEmptyKey();
+    return {N, 0};
+  }
+  static CachedHash<T> getTombstoneKey() {
+    T N = DenseMapInfo<T>::getTombstoneKey();
+    return {N, 0};
+  }
+  static unsigned getHashValue(CachedHash<T> Val) {
+    assert(!isEqual(Val, getEmptyKey()) && "Cannot hash the empty key!");
+    assert(!isEqual(Val, getTombstoneKey()) &&
+           "Cannot hash the tombstone key!");
+    return Val.Hash;
+  }
+  static bool isEqual(CachedHash<T> A, CachedHash<T> B) {
+    return DenseMapInfo<T>::isEqual(A.Val, B.Val);
+  }
+};
+
 // Provide DenseMapInfo for all pointers.
 template<typename T>
 struct DenseMapInfo<T*> {
diff --git a/llvm/unittests/ADT/DenseMapTest.cpp b/llvm/unittests/ADT/DenseMapTest.cpp
index caeba457918..db00f8cf8e5 100644
--- a/llvm/unittests/ADT/DenseMapTest.cpp
+++ b/llvm/unittests/ADT/DenseMapTest.cpp
@@ -496,6 +496,55 @@ TEST(DenseMapCustomTest, StringRefTest) {
   EXPECT_EQ(42, M.lookup(StringRef("a", 0)));
 }
 
+struct CachedHashTest {
+  unsigned Val;
+  unsigned *Counter = nullptr;
+  CachedHashTest(unsigned Val) : Val(Val) {}
+  CachedHashTest(unsigned Val, unsigned *Counter)
+      : Val(Val), Counter(Counter) {}
+};
+}
+namespace llvm {
+template <> struct DenseMapInfo<CachedHashTest> {
+  static CachedHashTest getEmptyKey() { return ~0; }
+  static CachedHashTest getTombstoneKey() { return ~0U - 1; }
+  static unsigned getHashValue(const CachedHashTest &X) {
+    ++*X.Counter;
+    return X.Val;
+  }
+  static bool isEqual(const CachedHashTest &LHS, const CachedHashTest &RHS) {
+    return LHS.Val == RHS.Val;
+  }
+};
+}
+namespace {
+
+TEST(DenseMapCustomTest, CachedHashTest) {
+  unsigned Counter = 0;
+  CachedHashTest Val(0, &Counter);
+  DenseMap<CachedHashTest, int> Map;
+
+  Map[Val] = 0;
+  ASSERT_EQ(1u, Counter);
+
+  Map.reserve(64);
+  ASSERT_EQ(2u, Counter);
+}
+
+// Like above, but now cache the hash.
+TEST(DenseMapCustomTest, CachedHashTest2) {
+  unsigned Counter = 0;
+  CachedHashTest Val(0, &Counter);
+  typedef CachedHash<CachedHashTest> Cached;
+  DenseMap<Cached, int> Map;
+
+  Map[Val] = 0;
+  ASSERT_EQ(1u, Counter);
+
+  Map.reserve(64);
+  ASSERT_EQ(1u, Counter);
+}
+
 // 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 {