1 files changed, 42 insertions, 3 deletions
diff --git a/llvm/unittests/ADT/StringMapTest.cpp b/llvm/unittests/ADT/StringMapTest.cpp
index 4ed0b76..f027f0d 100644
--- a/llvm/unittests/ADT/StringMapTest.cpp
+++ b/llvm/unittests/ADT/StringMapTest.cpp
@@ -231,12 +231,12 @@ TEST_F(StringMapTest, InsertRehashingPairTest) {
   // moved to a different bucket during internal rehashing. This depends on
   // the particular key, and the implementation of StringMap and HashString.
   // Changes to those might result in this test not actually checking that.
-  StringMap<uint32_t> t(1);
-  EXPECT_EQ(1u, t.getNumBuckets());
+  StringMap<uint32_t> t(0);
+  EXPECT_EQ(0u, t.getNumBuckets());
 
   StringMap<uint32_t>::iterator It =
     t.insert(std::make_pair("abcdef", 42)).first;
-  EXPECT_EQ(2u, t.getNumBuckets());
+  EXPECT_EQ(16u, t.getNumBuckets());
   EXPECT_EQ("abcdef", It->first());
   EXPECT_EQ(42u, It->second);
 }
@@ -356,4 +356,43 @@ TEST_F(StringMapTest, MoveDtor) {
   ASSERT_TRUE(B.empty());
 }
 
+namespace {
+// Simple class that counts how many moves and copy happens when growing a map
+struct CountCopyAndMove {
+  static unsigned Move;
+  static unsigned Copy;
+  CountCopyAndMove() {}
+
+  CountCopyAndMove(const CountCopyAndMove &) { Copy++; }
+  CountCopyAndMove &operator=(const CountCopyAndMove &) {
+    Copy++;
+    return *this;
+  }
+  CountCopyAndMove(CountCopyAndMove &&) { Move++; }
+  CountCopyAndMove &operator=(const CountCopyAndMove &&) {
+    Move++;
+    return *this;
+  }
+};
+unsigned CountCopyAndMove::Copy = 0;
+unsigned CountCopyAndMove::Move = 0;
+
+} // anonymous namespace
+
+// Make sure creating the map with an initial size of N actually gives us enough
+// buckets to insert N items without increasing allocation size.
+TEST(StringMapCustomTest, InitialSizeTest) {
+  // 1 is an "edge value", 32 is an arbitrary power of two, and 67 is an
+  // arbitrary prime, picked without any good reason.
+  for (auto Size : {1, 32, 67}) {
+    StringMap<CountCopyAndMove> Map(Size);
+    CountCopyAndMove::Copy = 0;
+    CountCopyAndMove::Move = 0;
+    for (int i = 0; i < Size; ++i)
+      Map.insert(std::make_pair(Twine(i).str(), CountCopyAndMove()));
+    EXPECT_EQ((unsigned)Size * 3, CountCopyAndMove::Move);
+    EXPECT_EQ(0u, CountCopyAndMove::Copy);
+  }
+}
+
 } // end anonymous namespace