Add some assertions and tests regarding bucket count and load factor …

…in Map.

PiperOrigin-RevId: 705896121

src/google/protobuf/map.h

@@ -360,7 +360,7 @@ class PROTOBUF_EXPORT UntypedMapBase {
 
  protected:
   // 16 bytes is the minimum useful size for the array cache in the arena.
-  enum : map_index_t { kMinTableSize = 16 / sizeof(void*) };
+  static constexpr map_index_t kMinTableSize = 16 / sizeof(void*);
 
  public:
   Arena* arena() const { return arena_; }
@@ -436,12 +436,6 @@ class PROTOBUF_EXPORT UntypedMapBase {
 #endif
   }
 
-  // Return a power of two no less than max(kMinTableSize, n).
-  // Assumes either n < kMinTableSize or n is a power of two.
-  map_index_t TableSize(map_index_t n) {
-    return n < kMinTableSize ? kMinTableSize : n;
-  }
-
   // Alignment of the nodes is the same as alignment of NodeBase.
   NodeBase* AllocNode() { return AllocNode(type_info_.node_size); }
 
@@ -726,6 +720,7 @@ class KeyMapBase : public UntypedMapBase {
   }
 
   NodeAndBucket FindHelper(typename TS::ViewType k) const {
+    AssertLoadFactor();
     map_index_t b = BucketNumber(k);
     for (auto* node = table_[b]; node != nullptr; node = node->next) {
       if (TS::ToView(static_cast<KeyNode*>(node)->key()) == k) {
@@ -764,6 +759,7 @@ class KeyMapBase : public UntypedMapBase {
     // or whatever.  But it's probably cheap enough to recompute that here;
     // it's likely that we're inserting into an empty or short list.
     ABSL_DCHECK(FindHelper(TS::ToView(node->key())).node == nullptr);
+    AssertLoadFactor();
     auto*& head = table_[b];
     if (head == nullptr) {
       head = node;
@@ -790,6 +786,10 @@ class KeyMapBase : public UntypedMapBase {
     return num_buckets - num_buckets / 16 * 4 - num_buckets % 2;
   }
 
+  void AssertLoadFactor() const {
+    ABSL_DCHECK_LE(num_elements_, CalculateHiCutoff(num_buckets_));
+  }
+
   // Returns whether it did resize.  Currently this is only used when
   // num_elements_ increases, though it could be used in other situations.
   // It checks for load too low as well as load too high: because any number
@@ -834,6 +834,7 @@ class KeyMapBase : public UntypedMapBase {
     if (num_buckets_ == kGlobalEmptyTableSize) {
       // This is the global empty array.
       // Just overwrite with a new one. No need to transfer or free anything.
+      ABSL_DCHECK_GE(kMinTableSize, new_num_buckets);
       num_buckets_ = index_of_first_non_null_ = kMinTableSize;
       table_ = CreateEmptyTable(num_buckets_);
       return;
@@ -981,6 +982,7 @@ class Map : private internal::KeyMapBase<internal::KeyForBase<Key>> {
     // won't trigger for leaked maps that never get destructed.
     StaticValidityCheck();
 
+    this->AssertLoadFactor();
     this->ClearTable(false, this->template GetDestroyNode<Node>());
   }
 

src/google/protobuf/map_test.cc

@@ -69,6 +69,7 @@ using ::testing::FieldsAre;
 using ::testing::Ge;
 using ::testing::Le;
 using ::testing::UnorderedElementsAre;
+using ::testing::UnorderedElementsAreArray;
 
 
 TEST(MapTest, CopyConstructIntegers) {
@@ -173,6 +174,18 @@ TEST(MapTest, CopyConstructMessagesWithArena) {
   EXPECT_EQ(map1["2"].GetArena(), &arena);
 }
 
+TEST(MapTest, CopyConstructionMaintainsProperLoadFactor) {
+  Map<int, int> original;
+  for (size_t size = 1; size < 50; ++size) {
+    // Add one element.
+    original[size];
+    Map<int, int> copy = original;
+    ASSERT_THAT(copy, UnorderedElementsAreArray(original));
+    EXPECT_LE(copy.size(),
+              MapTestPeer::CalculateHiCutoff(MapTestPeer::NumBuckets(copy)));
+  }
+}
+
 TEST(MapTest, AlwaysSerializesBothEntries) {
   for (const Message* prototype :
        {static_cast<const Message*>(