2 files changed, 52 insertions, 0 deletions

diff --git a/llvm/include/llvm/ADT/DenseMapInfo.h b/llvm/include/llvm/ADT/DenseMapInfo.h
index 6f17a647b63..b0a05307207 100644
--- a/llvm/include/llvm/ADT/DenseMapInfo.h
+++ b/llvm/include/llvm/ADT/DenseMapInfo.h
@@ -14,6 +14,8 @@
 #ifndef LLVM_ADT_DENSEMAPINFO_H
 #define LLVM_ADT_DENSEMAPINFO_H
 
+#include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/StringRef.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 #include "llvm/Support/type_traits.h"
 
@@ -163,6 +165,31 @@ struct DenseMapInfo<std::pair<T, U> > {
   }
 };
 
+// Provide DenseMapInfo for StringRefs.
+template <> struct DenseMapInfo<StringRef> {
+  static inline StringRef getEmptyKey() {
+    return StringRef(reinterpret_cast<const char *>(~static_cast<uintptr_t>(0)),
+                     0);
+  }
+  static inline StringRef getTombstoneKey() {
+    return StringRef(reinterpret_cast<const char *>(~static_cast<uintptr_t>(1)),
+                     0);
+  }
+  static unsigned getHashValue(StringRef Val) {
+    assert(Val.data() != getEmptyKey().data() && "Cannot hash the empty key!");
+    assert(Val.data() != getTombstoneKey().data() &&
+           "Cannot hash the tombstone key!");
+    return (unsigned)(hash_value(Val));
+  }
+  static bool isEqual(StringRef LHS, StringRef RHS) {
+    if (RHS.data() == getEmptyKey().data())
+      return LHS.data() == getEmptyKey().data();
+    if (RHS.data() == getTombstoneKey().data())
+      return LHS.data() == getTombstoneKey().data();
+    return LHS == RHS;
+  }
+};
+
 } // end namespace llvm
 
 #endif
diff --git a/llvm/unittests/ADT/DenseMapTest.cpp b/llvm/unittests/ADT/DenseMapTest.cpp
index 97807771725..2c6fe358832 100644
--- a/llvm/unittests/ADT/DenseMapTest.cpp
+++ b/llvm/unittests/ADT/DenseMapTest.cpp
@@ -323,6 +323,31 @@ TYPED_TEST(DenseMapTest, ConstIteratorTest) {
   EXPECT_TRUE(cit == cit2);
 }
 
+// Make sure DenseMap works with StringRef keys.
+TEST(DenseMapCustomTest, StringRefTest) {
+  DenseMap<StringRef, int> M;
+
+  M["a"] = 1;
+  M["b"] = 2;
+  M["c"] = 3;
+
+  EXPECT_EQ(3u, M.size());
+  EXPECT_EQ(1, M.lookup("a"));
+  EXPECT_EQ(2, M.lookup("b"));
+  EXPECT_EQ(3, M.lookup("c"));
+
+  EXPECT_EQ(0, M.lookup("q"));
+
+  // Test the empty string, spelled various ways.
+  EXPECT_EQ(0, M.lookup(""));
+  EXPECT_EQ(0, M.lookup(StringRef()));
+  EXPECT_EQ(0, M.lookup(StringRef("a", 0)));
+  M[""] = 42;
+  EXPECT_EQ(42, M.lookup(""));
+  EXPECT_EQ(42, M.lookup(StringRef()));
+  EXPECT_EQ(42, M.lookup(StringRef("a", 0)));
+}
+
 // 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 {