diff --git a/llvm/include/llvm/ADT/DenseMap.h b/llvm/include/llvm/ADT/DenseMap.h index f076049c55a26..d788d8e17e144 100644 --- a/llvm/include/llvm/ADT/DenseMap.h +++ b/llvm/include/llvm/ADT/DenseMap.h @@ -887,10 +887,8 @@ class SmallDenseMap AlignedCharArrayUnion storage; public: - explicit SmallDenseMap(unsigned NumInitBuckets = 0) { - if (NumInitBuckets > InlineBuckets) - NumInitBuckets = llvm::bit_ceil(NumInitBuckets); - init(NumInitBuckets); + explicit SmallDenseMap(unsigned NumElementsToReservre = 0) { + init(NumElementsToReservre); } SmallDenseMap(const SmallDenseMap &other) : BaseT() { @@ -905,7 +903,7 @@ class SmallDenseMap template SmallDenseMap(const InputIt &I, const InputIt &E) { - init(NextPowerOf2(std::distance(I, E))); + init(std::distance(I, E)); this->insert(I, E); } @@ -1017,7 +1015,8 @@ class SmallDenseMap this->BaseT::copyFrom(other); } - void init(unsigned InitBuckets) { + void init(unsigned InitNumEntries) { + auto InitBuckets = BaseT::getMinBucketToReserveForEntries(InitNumEntries); Small = true; if (InitBuckets > InlineBuckets) { Small = false; diff --git a/llvm/unittests/ADT/DenseMapTest.cpp b/llvm/unittests/ADT/DenseMapTest.cpp index 785ab16271d93..50e9c6e138ef1 100644 --- a/llvm/unittests/ADT/DenseMapTest.cpp +++ b/llvm/unittests/ADT/DenseMapTest.cpp @@ -962,4 +962,73 @@ TEST(DenseMapCustomTest, PairPrinting) { EXPECT_EQ(R"({ (1, "one"), (2, "two") })", ::testing::PrintToString(Map)); } +TEST(DenseMapCustomTest, InitSize) { + constexpr unsigned ElemSize = sizeof(std::pair); + + { + DenseMap Map; + EXPECT_EQ(ElemSize * 0U, Map.getMemorySize()); + } + { + DenseMap Map(0); + EXPECT_EQ(ElemSize * 0U, Map.getMemorySize()); + } + { + DenseMap Map(1); + EXPECT_EQ(ElemSize * 4U, Map.getMemorySize()); + } + { + DenseMap Map(2); + EXPECT_EQ(ElemSize * 4U, Map.getMemorySize()); + } + { + DenseMap Map(3); + EXPECT_EQ(ElemSize * 8U, Map.getMemorySize()); + } + { + int A, B; + DenseMap Map = {{&A, 1}, {&B, 2}}; + EXPECT_EQ(ElemSize * 4U, Map.getMemorySize()); + } + { + int A, B, C; + DenseMap Map = {{&A, 1}, {&B, 2}, {&C, 3}}; + EXPECT_EQ(ElemSize * 8U, Map.getMemorySize()); + } +} + +TEST(SmallDenseMapCustomTest, InitSize) { + constexpr unsigned ElemSize = sizeof(std::pair); + { + SmallDenseMap Map; + EXPECT_EQ(ElemSize * 4U, Map.getMemorySize()); + } + { + SmallDenseMap Map(0); + EXPECT_EQ(ElemSize * 4U, Map.getMemorySize()); + } + { + SmallDenseMap Map(1); + EXPECT_EQ(ElemSize * 4U, Map.getMemorySize()); + } + { + SmallDenseMap Map(2); + EXPECT_EQ(ElemSize * 4U, Map.getMemorySize()); + } + { + SmallDenseMap Map(3); + EXPECT_EQ(ElemSize * 8U, Map.getMemorySize()); + } + { + int A, B; + SmallDenseMap Map = {{&A, 1}, {&B, 2}}; + EXPECT_EQ(ElemSize * 4U, Map.getMemorySize()); + } + { + int A, B, C; + SmallDenseMap Map = {{&A, 1}, {&B, 2}, {&C, 3}}; + EXPECT_EQ(ElemSize * 8U, Map.getMemorySize()); + } +} + } // namespace