1 files changed, 51 insertions, 0 deletions
diff --git a/unittests/ADT/SmallPtrSetTest.cpp b/unittests/ADT/SmallPtrSetTest.cpp
index d4d963fdc5bd..fc14c684d67f 100644
--- a/unittests/ADT/SmallPtrSetTest.cpp
+++ b/unittests/ADT/SmallPtrSetTest.cpp
@@ -12,7 +12,9 @@
 //===----------------------------------------------------------------------===//
 
 #include "gtest/gtest.h"
+#include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/Support/PointerLikeTypeTraits.h"
 
 using namespace llvm;
 
@@ -279,3 +281,52 @@ TEST(SmallPtrSetTest, EraseTest) {
   SmallPtrSet<int *, 2> A;
   checkEraseAndIterators(A);
 }
+
+// Verify that dereferencing and iteration work.
+TEST(SmallPtrSetTest, dereferenceAndIterate) {
+  int Ints[] = {0, 1, 2, 3, 4, 5, 6, 7};
+  SmallPtrSet<const int *, 4> S;
+  for (int &I : Ints) {
+    EXPECT_EQ(&I, *S.insert(&I).first);
+    EXPECT_EQ(&I, *S.find(&I));
+  }
+
+  // Iterate from each and count how many times each element is found.
+  int Found[sizeof(Ints)/sizeof(int)] = {0};
+  for (int &I : Ints)
+    for (auto F = S.find(&I), E = S.end(); F != E; ++F)
+      ++Found[*F - Ints];
+
+  // Sort.  We should hit the first element just once and the final element N
+  // times.
+  std::sort(std::begin(Found), std::end(Found));
+  for (auto F = std::begin(Found), E = std::end(Found); F != E; ++F)
+    EXPECT_EQ(F - Found + 1, *F);
+}
+
+// Verify that const pointers work for count and find even when the underlying
+// SmallPtrSet is not for a const pointer type.
+TEST(SmallPtrSetTest, ConstTest) {
+  SmallPtrSet<int *, 8> IntSet;
+  int A;
+  int *B = &A;
+  const int *C = &A;
+  IntSet.insert(B);
+  EXPECT_EQ(IntSet.count(B), 1u);
+  EXPECT_EQ(IntSet.count(C), 1u);
+  EXPECT_NE(IntSet.find(B), IntSet.end());
+  EXPECT_NE(IntSet.find(C), IntSet.end());
+}
+
+// Verify that we automatically get the const version of PointerLikeTypeTraits
+// filled in for us, even for a non-pointer type
+using TestPair = PointerIntPair<int *, 1>;
+
+TEST(SmallPtrSetTest, ConstNonPtrTest) {
+  SmallPtrSet<TestPair, 8> IntSet;
+  int A[1];
+  TestPair Pair(&A[0], 1);
+  IntSet.insert(Pair);
+  EXPECT_EQ(IntSet.count(Pair), 1u);
+  EXPECT_NE(IntSet.find(Pair), IntSet.end());
+}