3 files changed, 58 insertions, 4 deletions

diff --git a/llvm/include/llvm/ADT/StringMap.h b/llvm/include/llvm/ADT/StringMap.h
index 47a87c2c51e..7562b840c46 100644
--- a/llvm/include/llvm/ADT/StringMap.h
+++ b/llvm/include/llvm/ADT/StringMap.h
@@ -15,6 +15,7 @@
 #define LLVM_ADT_STRINGMAP_H
 
 #include "llvm/ADT/StringRef.h"
+#include "llvm/ADT/Twine.h"
 #include "llvm/Support/Allocator.h"
 #include <cstring>
 #include <utility>
diff --git a/llvm/lib/Support/StringMap.cpp b/llvm/lib/Support/StringMap.cpp
index 7be946642d9..7da9ccbd40c 100644
--- a/llvm/lib/Support/StringMap.cpp
+++ b/llvm/lib/Support/StringMap.cpp
@@ -17,12 +17,26 @@
 #include <cassert>
 using namespace llvm;
 
+/// Returns the number of buckets to allocate to ensure that the DenseMap can
+/// accommodate \p NumEntries without need to grow().
+static unsigned getMinBucketToReserveForEntries(unsigned NumEntries) {
+  // Ensure that "NumEntries * 4 < NumBuckets * 3"
+  if (NumEntries == 0)
+    return 0;
+  // +1 is required because of the strict equality.
+  // For example if NumEntries is 48, we need to return 401.
+  return NextPowerOf2(NumEntries * 4 / 3 + 1);
+}
+
 StringMapImpl::StringMapImpl(unsigned InitSize, unsigned itemSize) {
   ItemSize = itemSize;
   
   // If a size is specified, initialize the table with that many buckets.
   if (InitSize) {
-    init(InitSize);
+    // The table will grow when the number of entries reach 3/4 of the number of
+    // buckets. To guarantee that "InitSize" number of entries can be inserted
+    // in the table without growing, we allocate just what is needed here.
+    init(getMinBucketToReserveForEntries(InitSize));
     return;
   }
   
diff --git a/llvm/unittests/ADT/StringMapTest.cpp b/llvm/unittests/ADT/StringMapTest.cpp
index 4ed0b76f0f4..f027f0d5a95 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