Add growth API to the new hashtables. (#4044)

This adds two different growth APIs. This is instead of the more
conventional STL `reserve` method. One allows users that aren't trying
to grow in anticipation of an *exact* count of insertions, but generally
trying to size the table to the correct ballpark with a power-of-two
estimate.

The other API allows pre-growing to allow a specific number of
insertions to be performed without further growth. This API takes the
maximum load factor and other implementation details into account.

---------

Co-authored-by: josh11b <15258583+josh11b@users.noreply.github.com>

common/map.h

@@ -323,6 +323,21 @@ class MapBase : protected RawHashtable::BaseImpl<InputKeyT, InputValueT,
              std::invocable<InsertCallbackT, LookupKeyT, void*, void*> &&
              std::invocable<UpdateCallbackT, KeyT&, ValueT&>);
 
+  // Grow the map to a specific allocation size.
+  //
+  // This will grow the map's hashtable if necessary for it to have an
+  // allocation size of `target_alloc_size` which must be a power of two. Note
+  // that this will not allow that many keys to be inserted, but a smaller
+  // number based on the maximum load factor. If a specific number of insertions
+  // need to be achieved without triggering growth, use the `GrowForInsertCount`
+  // method.
+  auto GrowToAllocSize(ssize_t target_alloc_size,
+                       KeyContextT key_context = KeyContextT()) -> void;
+
+  // Grow the map sufficiently to allow inserting the specified number of keys.
+  auto GrowForInsertCount(ssize_t count,
+                          KeyContextT key_context = KeyContextT()) -> void;
+
   // Erase a key from the map.
   template <typename LookupKeyT>
   auto Erase(LookupKeyT lookup_key, KeyContextT key_context = KeyContextT())
@@ -533,6 +548,18 @@ MapBase<InputKeyT, InputValueT, InputKeyContextT>::Update(
   return InsertKVResult(true, *entry);
 }
 
+template <typename InputKeyT, typename InputValueT, typename InputKeyContextT>
+void MapBase<InputKeyT, InputValueT, InputKeyContextT>::GrowToAllocSize(
+    ssize_t target_alloc_size, KeyContextT key_context) {
+  this->GrowToAllocSizeImpl(target_alloc_size, key_context);
+}
+
+template <typename InputKeyT, typename InputValueT, typename InputKeyContextT>
+void MapBase<InputKeyT, InputValueT, InputKeyContextT>::GrowForInsertCount(
+    ssize_t count, KeyContextT key_context) {
+  this->GrowForInsertCountImpl(count, key_context);
+}
+
 template <typename InputKeyT, typename InputValueT, typename InputKeyContextT>
 template <typename LookupKeyT>
 auto MapBase<InputKeyT, InputValueT, InputKeyContextT>::Erase(

common/map_test.cpp

@@ -222,6 +222,112 @@ TYPED_TEST(MapTest, Conversions) {
   EXPECT_EQ(104, *cmv3[4]);
 }
 
+TYPED_TEST(MapTest, GrowToAllocSize) {
+  using MapT = TypeParam;
+
+  MapT m;
+  // Grow when empty. May be a no-op for some small sizes.
+  m.GrowToAllocSize(32);
+
+  // Add some elements that will need to be propagated through subsequent
+  // growths. Also delete some.
+  ssize_t storage_bytes = m.ComputeMetrics().storage_bytes;
+  for (int i : llvm::seq(1, 24)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(m.Insert(i, i * 100).is_inserted());
+  }
+  for (int i : llvm::seq(1, 8)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(m.Erase(i));
+  }
+  // No further growth triggered.
+  EXPECT_EQ(storage_bytes, m.ComputeMetrics().storage_bytes);
+
+  // No-op.
+  m.GrowToAllocSize(16);
+  ExpectMapElementsAre(
+      m, MakeKeyValues([](int k) { return k * 100; }, llvm::seq(8, 24)));
+  // No further growth triggered.
+  EXPECT_EQ(storage_bytes, m.ComputeMetrics().storage_bytes);
+
+  // Get a few doubling based growths, and at least one beyond the largest small
+  // size.
+  m.GrowToAllocSize(64);
+  ExpectMapElementsAre(
+      m, MakeKeyValues([](int k) { return k * 100; }, llvm::seq(8, 24)));
+  m.GrowToAllocSize(128);
+  ExpectMapElementsAre(
+      m, MakeKeyValues([](int k) { return k * 100; }, llvm::seq(8, 24)));
+  // Update the storage bytes after growth.
+  EXPECT_LT(storage_bytes, m.ComputeMetrics().storage_bytes);
+  storage_bytes = m.ComputeMetrics().storage_bytes;
+
+  // Add some more, but not enough to trigger further growth, and then grow by
+  // several more multiples of two to test handling large growth.
+  for (int i : llvm::seq(24, 48)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(m.Insert(i, i * 100).is_inserted());
+  }
+  for (int i : llvm::seq(8, 16)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(m.Erase(i));
+  }
+  // No growth from insertions.
+  EXPECT_EQ(storage_bytes, m.ComputeMetrics().storage_bytes);
+
+  m.GrowToAllocSize(1024);
+  ExpectMapElementsAre(
+      m, MakeKeyValues([](int k) { return k * 100; }, llvm::seq(16, 48)));
+  // Storage should have grown.
+  EXPECT_LT(storage_bytes, m.ComputeMetrics().storage_bytes);
+}
+
+TYPED_TEST(MapTest, GrowForInsert) {
+  using MapT = TypeParam;
+
+  MapT m;
+  m.GrowForInsertCount(42);
+  ssize_t storage_bytes = m.ComputeMetrics().storage_bytes;
+  for (int i : llvm::seq(1, 42)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(m.Insert(i, i * 100).is_inserted());
+  }
+  ExpectMapElementsAre(
+      m, MakeKeyValues([](int k) { return k * 100; }, llvm::seq(1, 42)));
+  EXPECT_EQ(storage_bytes, m.ComputeMetrics().storage_bytes);
+
+  // Erase many elements and grow again for another insert.
+  for (int i : llvm::seq(1, 32)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(m.Erase(i));
+  }
+  m.GrowForInsertCount(42);
+  storage_bytes = m.ComputeMetrics().storage_bytes;
+  for (int i : llvm::seq(42, 84)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(m.Insert(i, i * 100).is_inserted());
+  }
+  ExpectMapElementsAre(
+      m, MakeKeyValues([](int k) { return k * 100; }, llvm::seq(32, 84)));
+  EXPECT_EQ(storage_bytes, m.ComputeMetrics().storage_bytes);
+
+  // Erase all the elements, then grow for a much larger insertion and insert
+  // again.
+  for (int i : llvm::seq(32, 84)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(m.Erase(i));
+  }
+  m.GrowForInsertCount(321);
+  storage_bytes = m.ComputeMetrics().storage_bytes;
+  for (int i : llvm::seq(128, 321 + 128)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(m.Insert(i, i * 100).is_inserted());
+  }
+  ExpectMapElementsAre(m, MakeKeyValues([](int k) { return k * 100; },
+                                        llvm::seq(128, 321 + 128)));
+  EXPECT_EQ(storage_bytes, m.ComputeMetrics().storage_bytes);
+}
+
 // This test is largely exercising the underlying `RawHashtable` implementation
 // with complex growth, erasure, and re-growth.
 TYPED_TEST(MapTest, ComplexOpSequence) {

common/raw_hashtable.h

@@ -458,6 +458,20 @@ class BaseImpl {
   auto InsertImpl(LookupKeyT lookup_key, KeyContextT key_context)
       -> std::pair<EntryT*, bool>;
 
+  // Grow the table to specific allocation size.
+  //
+  // This will grow the the table if necessary for it to have an allocation size
+  // of `target_alloc_size` which must be a power of two. Note that this will
+  // not allow that many keys to be inserted into the hashtable, but a smaller
+  // number based on the load factor. If a specific number of insertions need to
+  // be achieved without triggering growth, use the `GrowForInsertCountImpl`
+  // method.
+  auto GrowToAllocSizeImpl(ssize_t target_alloc_size, KeyContextT key_context)
+      -> void;
+
+  // Grow the table to allow inserting the specified number of keys.
+  auto GrowForInsertCountImpl(ssize_t count, KeyContextT key_context) -> void;
+
   // Looks up the entry in the hashtable, and if found destroys the entry and
   // returns `true`. If not found, returns `false`.
   //
@@ -513,6 +527,7 @@ class BaseImpl {
   static auto ComputeNextAllocSize(ssize_t old_alloc_size) -> ssize_t;
   static auto GrowthThresholdForAllocSize(ssize_t alloc_size) -> ssize_t;
 
+  auto GrowToNextAllocSize(KeyContextT key_context) -> void;
   template <typename LookupKeyT>
   auto GrowAndInsert(LookupKeyT lookup_key, KeyContextT key_context) -> EntryT*;
 
@@ -911,6 +926,84 @@ auto BaseImpl<InputKeyT, InputValueT, InputKeyContextT>::InsertImpl(
                     "or an empty slot.";
 }
 
+template <typename InputKeyT, typename InputValueT, typename InputKeyContextT>
+[[clang::noinline]] auto
+BaseImpl<InputKeyT, InputValueT, InputKeyContextT>::GrowToAllocSizeImpl(
+    ssize_t target_alloc_size, KeyContextT key_context) -> void {
+  CARBON_CHECK(llvm::isPowerOf2_64(target_alloc_size));
+  if (target_alloc_size <= alloc_size()) {
+    return;
+  }
+
+  // If this is the next alloc size, we can used our optimized growth strategy.
+  if (target_alloc_size == ComputeNextAllocSize(alloc_size())) {
+    GrowToNextAllocSize(key_context);
+    return;
+  }
+
+  // Create locals for the old state of the table.
+  ssize_t old_size = alloc_size();
+  CARBON_DCHECK(old_size > 0);
+  bool old_small = is_small();
+  Storage* old_storage = storage();
+  uint8_t* old_metadata = metadata();
+  EntryT* old_entries = entries();
+
+  // Configure for the new size and allocate the new storage.
+  alloc_size() = target_alloc_size;
+  storage() = Allocate(target_alloc_size);
+  std::memset(metadata(), 0, target_alloc_size);
+  growth_budget_ = GrowthThresholdForAllocSize(target_alloc_size);
+
+  // Just re-insert all the entries. As we're more than doubling the table size,
+  // we don't bother with fancy optimizations here. Even using `memcpy` for the
+  // entries seems unlikely to be a significant win given how sparse the
+  // insertions will end up being.
+  ssize_t count = 0;
+  for (ssize_t group_index = 0; group_index < old_size;
+       group_index += GroupSize) {
+    auto g = MetadataGroup::Load(old_metadata, group_index);
+    auto present_matched_range = g.MatchPresent();
+    for (ssize_t byte_index : present_matched_range) {
+      ++count;
+      ssize_t index = group_index + byte_index;
+      EntryT* new_entry =
+          InsertIntoEmpty(old_entries[index].key(), key_context);
+      new_entry->MoveFrom(std::move(old_entries[index]));
+    }
+  }
+  growth_budget_ -= count;
+
+  if (!old_small) {
+    // Old isn't a small buffer, so we need to deallocate it.
+    Deallocate(old_storage, old_size);
+  }
+}
+
+template <typename InputKeyT, typename InputValueT, typename InputKeyContextT>
+auto BaseImpl<InputKeyT, InputValueT, InputKeyContextT>::GrowForInsertCountImpl(
+    ssize_t count, KeyContextT key_context) -> void {
+  if (count < growth_budget_) {
+    // Already space for the needed growth.
+    return;
+  }
+
+  // Currently, we don't account for any tombstones marking deleted elements,
+  // and just conservatively ensure the growth will create adequate growth
+  // budget for insertions. We could make this more precise by instead walking
+  // the table and only counting present slots, as once we grow we'll be able to
+  // reclaim all of the deleted slots. But this adds complexity and it isn't
+  // clear this is necessary so we do the simpler conservative thing.
+  ssize_t used_budget =
+      GrowthThresholdForAllocSize(alloc_size()) - growth_budget_;
+  ssize_t budget_needed = used_budget + count;
+  ssize_t space_needed = budget_needed + (budget_needed / 7);
+  ssize_t target_alloc_size = llvm::NextPowerOf2(space_needed);
+  CARBON_CHECK(GrowthThresholdForAllocSize(target_alloc_size) >
+               (budget_needed));
+  GrowToAllocSizeImpl(target_alloc_size, key_context);
+}
+
 template <typename InputKeyT, typename InputValueT, typename InputKeyContextT>
 template <typename LookupKeyT>
 auto BaseImpl<InputKeyT, InputValueT, InputKeyContextT>::EraseImpl(
@@ -1088,16 +1181,36 @@ auto BaseImpl<InputKeyT, InputValueT,
   return alloc_size - alloc_size / 8;
 }
 
-// Grow the hashtable to create space and then insert into it. Returns the
-// selected insertion entry. Never returns null. In addition to growing and
-// selecting the insertion entry, this routine updates the metadata array so
-// that this function can be directly called and the result returned from
-// `InsertImpl`.
+// Optimized routine for growing to the next alloc size.
+//
+// A particularly common and important-to-optimize path is growing to the next
+// alloc size, which will always be a doubling of the allocated size. This
+// allows an important optimization -- we're adding exactly one more high bit to
+// the hash-computed index for each entry. This in turn means we can classify
+// every entry in the table into three cases:
+//
+// 1) The new high bit is zero, the entry is at the same index in the new
+//    table as the old.
+//
+// 2) The new high bit is one, the entry is at the old index plus the old
+//    size.
+//
+// 3) The entry's current index doesn't match the initial hash index because
+//    it required some amount of probing to find an empty slot.
+//
+// The design of the hash table tries to minimize how many entries fall into
+// case (3), so we expect the vast majority of entries to be in (1) or (2). This
+// lets us model growth notionally as copying the hashtable twice into the lower
+// and higher halves of the new allocation, clearing out the now-empty slots
+// (from both deleted entries and entries in the other half of the table after
+// growth), and inserting any probed elements. That model in turn is much more
+// efficient than re-inserting all of the elements as it avoids the unnecessary
+// parts of insertion and avoids interleaving random accesses for the probed
+// elements. But most importantly, for trivially relocatable types it allows us
+// to use `memcpy` rather than moving the elements individually.
 template <typename InputKeyT, typename InputValueT, typename InputKeyContextT>
-template <typename LookupKeyT>
-[[clang::noinline]] auto
-BaseImpl<InputKeyT, InputValueT, InputKeyContextT>::GrowAndInsert(
-    LookupKeyT lookup_key, KeyContextT key_context) -> EntryT* {
+auto BaseImpl<InputKeyT, InputValueT, InputKeyContextT>::GrowToNextAllocSize(
+    KeyContextT key_context) -> void {
   // We collect the probed elements in a small vector for re-insertion. It is
   // tempting to reuse the already allocated storage, but doing so appears to
   // be a (very slight) performance regression. These are relatively rare and
@@ -1109,12 +1222,9 @@ BaseImpl<InputKeyT, InputValueT, InputKeyContextT>::GrowAndInsert(
   // we can handle this most efficiently with temporary, additional storage.
   llvm::SmallVector<ssize_t, 128> probed_indices;
 
-  // We grow into a new `MapBase` so that both the new and old maps are
-  // fully functional until all the entries are moved over. However, we directly
-  // manipulate the internals to short circuit many aspects of the growth.
+  // Create locals for the old state of the table.
   ssize_t old_size = alloc_size();
   CARBON_DCHECK(old_size > 0);
-  CARBON_DCHECK(growth_budget_ == 0);
 
   bool old_small = is_small();
   Storage* old_storage = storage();
@@ -1137,7 +1247,7 @@ BaseImpl<InputKeyT, InputValueT, InputKeyContextT>::GrowAndInsert(
       << ", size: " << old_size;
 #endif
 
-  // Compute the new size and grow the storage in place (if possible).
+  // Configure for the new size and allocate the new storage.
   ssize_t new_size = ComputeNextAllocSize(old_size);
   alloc_size() = new_size;
   storage() = Allocate(new_size);
@@ -1147,25 +1257,11 @@ BaseImpl<InputKeyT, InputValueT, InputKeyContextT>::GrowAndInsert(
   uint8_t* new_metadata = metadata();
   EntryT* new_entries = entries();
 
-  // We always double the size when we grow. This allows an important
-  // optimization -- we're adding exactly one more high bit to the hash-computed
-  // index for each entry. This in turn means we can classify every entry in the
-  // table into three cases:
-  //
-  // 1) The new high bit is zero, the entry is at the same index in the new
-  //    table as the old.
-  //
-  // 2) The new high bit is one, the entry is at the old index plus the old
-  //    size.
-  //
-  // 3) The entry's current index doesn't match the initial hash index because
-  //    it required some amount of probing to find an empty slot.
-  //
-  // The design of the hash table tries to minimize how many entries fall into
-  // case (3), so we expect the vast majority of entries to be in (1) or (2).
-  // This lets us model growth notionally as duplicating the hash table,
-  // clearing out the empty slots, and inserting any probed elements.
-
+  // Walk the metadata groups, clearing deleted to empty, duplicating the
+  // metadata for the low and high halves, and updating it based on where each
+  // entry will go in the new table. The updated metadata group is written to
+  // the new table, and for non-trivially relocatable entry types, the entry is
+  // also moved to its new location.
   ssize_t count = 0;
   for (ssize_t group_index = 0; group_index < old_size;
        group_index += GroupSize) {
@@ -1280,9 +1376,22 @@ BaseImpl<InputKeyT, InputValueT, InputKeyContextT>::GrowAndInsert(
     // Old isn't a small buffer, so we need to deallocate it.
     Deallocate(old_storage, old_size);
   }
+}
+
+// Grow the hashtable to create space and then insert into it. Returns the
+// selected insertion entry. Never returns null. In addition to growing and
+// selecting the insertion entry, this routine updates the metadata array so
+// that this function can be directly called and the result returned from
+// `InsertImpl`.
+template <typename InputKeyT, typename InputValueT, typename InputKeyContextT>
+template <typename LookupKeyT>
+[[clang::noinline]] auto
+BaseImpl<InputKeyT, InputValueT, InputKeyContextT>::GrowAndInsert(
+    LookupKeyT lookup_key, KeyContextT key_context) -> EntryT* {
+  GrowToNextAllocSize(key_context);
 
-  // And lastly insert the lookup_key into an index in the newly grown map and
-  // return that index for use.
+  // And insert the lookup_key into an index in the newly grown map and return
+  // that index for use.
   --growth_budget_;
   return InsertIntoEmpty(lookup_key, key_context);
 }

common/set.h

@@ -220,6 +220,21 @@ class SetBase
               KeyContextT key_context = KeyContextT()) -> InsertResult
     requires std::invocable<InsertCallbackT, LookupKeyT, void*>;
 
+  // Grow the set to a specific allocation size.
+  //
+  // This will grow the set's hashtable if necessary for it to have an
+  // allocation size of `target_alloc_size` which must be a power of two. Note
+  // that this will not allow that many keys to be inserted, but a smaller
+  // number based on the maximum load factor. If a specific number of insertions
+  // need to be achieved without triggering growth, use the `GrowForInsertCount`
+  // method.
+  auto GrowToAllocSize(ssize_t target_alloc_size,
+                       KeyContextT key_context = KeyContextT()) -> void;
+
+  // Grow the set sufficiently to allow inserting the specified number of keys.
+  auto GrowForInsertCount(ssize_t count,
+                          KeyContextT key_context = KeyContextT()) -> void;
+
   // Erase a key from the set.
   template <typename LookupKeyT>
   auto Erase(LookupKeyT lookup_key, KeyContextT key_context = KeyContextT())
@@ -333,6 +348,18 @@ auto SetBase<InputKeyT, InputKeyContextT>::Insert(LookupKeyT lookup_key,
   return InsertResult(true, entry->key());
 }
 
+template <typename InputKeyT, typename InputKeyContextT>
+void SetBase<InputKeyT, InputKeyContextT>::GrowToAllocSize(
+    ssize_t target_alloc_size, KeyContextT key_context) {
+  this->GrowToAllocSizeImpl(target_alloc_size, key_context);
+}
+
+template <typename InputKeyT, typename InputKeyContextT>
+void SetBase<InputKeyT, InputKeyContextT>::GrowForInsertCount(
+    ssize_t count, KeyContextT key_context) {
+  this->GrowForInsertCountImpl(count, key_context);
+}
+
 template <typename InputKeyT, typename InputKeyContextT>
 template <typename LookupKeyT>
 auto SetBase<InputKeyT, InputKeyContextT>::Erase(LookupKeyT lookup_key,

common/set_test.cpp

@@ -176,6 +176,107 @@ TYPED_TEST(SetTest, Conversions) {
   EXPECT_TRUE(csv2.Contains(3));
 }
 
+TYPED_TEST(SetTest, GrowToAllocSize) {
+  using SetT = TypeParam;
+
+  SetT s;
+  // Grow when empty. May be a no-op for some small sizes.
+  s.GrowToAllocSize(32);
+
+  // Add some elements that will need to be propagated through subsequent
+  // growths. Also delete some.
+  ssize_t storage_bytes = s.ComputeMetrics().storage_bytes;
+  for (int i : llvm::seq(1, 24)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(s.Insert(i).is_inserted());
+  }
+  for (int i : llvm::seq(1, 8)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(s.Erase(i));
+  }
+  // No further growth triggered.
+  EXPECT_EQ(storage_bytes, s.ComputeMetrics().storage_bytes);
+
+  // No-op.
+  s.GrowToAllocSize(16);
+  ExpectSetElementsAre(s, MakeElements(llvm::seq(8, 24)));
+  // No further growth triggered.
+  EXPECT_EQ(storage_bytes, s.ComputeMetrics().storage_bytes);
+
+  // Get a few doubling based growths, and at least one beyond the largest small
+  // size.
+  s.GrowToAllocSize(64);
+  ExpectSetElementsAre(s, MakeElements(llvm::seq(8, 24)));
+  s.GrowToAllocSize(128);
+  ExpectSetElementsAre(s, MakeElements(llvm::seq(8, 24)));
+  s.GrowToAllocSize(256);
+  ExpectSetElementsAre(s, MakeElements(llvm::seq(8, 24)));
+  // Update the storage bytes after growth.
+  EXPECT_LT(storage_bytes, s.ComputeMetrics().storage_bytes);
+  storage_bytes = s.ComputeMetrics().storage_bytes;
+
+  // Add some more, but not enough to trigger further growth, and then grow by
+  // several more multiples of two to test handling large growth.
+  for (int i : llvm::seq(24, 48)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(s.Insert(i).is_inserted());
+  }
+  for (int i : llvm::seq(8, 16)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(s.Erase(i));
+  }
+  // No growth from insertions.
+  EXPECT_EQ(storage_bytes, s.ComputeMetrics().storage_bytes);
+
+  s.GrowToAllocSize(1024);
+  ExpectSetElementsAre(s, MakeElements(llvm::seq(16, 48)));
+  // Storage should have grown.
+  EXPECT_LT(storage_bytes, s.ComputeMetrics().storage_bytes);
+}
+
+TYPED_TEST(SetTest, GrowForInsert) {
+  using SetT = TypeParam;
+
+  SetT s;
+  s.GrowForInsertCount(42);
+  ssize_t storage_bytes = s.ComputeMetrics().storage_bytes;
+  for (int i : llvm::seq(1, 42)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(s.Insert(i).is_inserted());
+  }
+  ExpectSetElementsAre(s, MakeElements(llvm::seq(1, 42)));
+  EXPECT_EQ(storage_bytes, s.ComputeMetrics().storage_bytes);
+
+  // Erase many elements and grow again for another insert.
+  for (int i : llvm::seq(1, 32)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(s.Erase(i));
+  }
+  s.GrowForInsertCount(42);
+  storage_bytes = s.ComputeMetrics().storage_bytes;
+  for (int i : llvm::seq(42, 84)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(s.Insert(i).is_inserted());
+  }
+  ExpectSetElementsAre(s, MakeElements(llvm::seq(32, 84)));
+  EXPECT_EQ(storage_bytes, s.ComputeMetrics().storage_bytes);
+
+  // Erase all the elements, then grow for a much larger insertion and insert
+  // again.
+  for (int i : llvm::seq(32, 84)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(s.Erase(i));
+  }
+  s.GrowForInsertCount(321);
+  storage_bytes = s.ComputeMetrics().storage_bytes;
+  for (int i : llvm::seq(128, 321 + 128)) {
+    SCOPED_TRACE(llvm::formatv("Key: {0}", i).str());
+    ASSERT_TRUE(s.Insert(i).is_inserted());
+  }
+  ExpectSetElementsAre(s, MakeElements(llvm::seq(128, 321 + 128)));
+  EXPECT_EQ(storage_bytes, s.ComputeMetrics().storage_bytes);
+}
+
 TEST(SetContextTest, Basic) {
   llvm::SmallVector<TestData> keys;
   for (int i : llvm::seq(0, 513)) {