Extend support of "blocked elements" (sentinels at the end of the control array) to all capacities, not just single-group tables.

For large capacities there is a limit of 5 blocked elements.

PiperOrigin-RevId: 938536642
Change-Id: I8b5d6c4120a4be334f03895b570ab939ad6633b1
diff --git a/absl/container/internal/raw_hash_set.cc b/absl/container/internal/raw_hash_set.cc
index ca3ddec..45b3fa6 100644
--- a/absl/container/internal/raw_hash_set.cc
+++ b/absl/container/internal/raw_hash_set.cc
@@ -498,45 +498,45 @@
 //
 // Unlike setting it directly, this function will perform bounds checks and
 // mirror the value to the cloned tail if necessary.
-inline void SetCtrlNoSanitizeImpl(const CommonFields& c, size_t i, ctrl_t h) {
-  ABSL_SWISSTABLE_ASSERT(i < c.capacity());
-  ctrl_t* ctrl = c.control();
-  const size_t cap = c.capacity();
-  ctrl[i] = h;
-  ctrl[((i - NumClonedBytes()) & cap) + (NumClonedBytes() & cap)] = h;
-}
-
 inline void SetCtrl(const CommonFields& c, size_t i, ctrl_t h,
                     size_t slot_size) {
   ABSL_SWISSTABLE_ASSERT(!c.is_small());
   DoSanitizeOnSetCtrl(c, i, h, slot_size);
-  SetCtrlNoSanitizeImpl(c, i, h);
+  ctrl_t* ctrl = c.control();
+  ctrl[i] = h;
+  ctrl[((i - NumClonedBytes()) & c.capacity()) +
+       (NumClonedBytes() & c.capacity())] = h;
 }
 // Overload for setting to an occupied `h2_t` rather than a special `ctrl_t`.
 inline void SetCtrl(const CommonFields& c, size_t i, h2_t h, size_t slot_size) {
   SetCtrl(c, i, static_cast<ctrl_t>(h), slot_size);
 }
 
-// Sets `ctrl[i]` to `ctrl_t::kSentinel`.
+inline void SetCtrlInSingleGroupTableNoSanitizeImpl(const CommonFields& c,
+                                                    size_t i, ctrl_t h) {
+  ABSL_SWISSTABLE_ASSERT(!c.is_small());
+  ABSL_SWISSTABLE_ASSERT(is_single_group(c.capacity()));
+  ctrl_t* ctrl = c.control();
+  ctrl[i] = h;
+  ctrl[i + c.capacity() + 1] = h;
+}
+
+// Sets `ctrl[i]` to `ctrl_t::kSentinel` in single group table.
 //
 // Unlike setting it directly, this function will perform bounds checks and
 // mirror the value to the cloned tail if necessary.
-inline void BlockCtrl(const CommonFields& c, size_t i) {
-  ABSL_SWISSTABLE_ASSERT(!c.is_small());
-  SetCtrlNoSanitizeImpl(c, i, ctrl_t::kSentinel);
+inline void BlockCtrlInSingleGroupTable(const CommonFields& c, size_t i) {
+  SetCtrlInSingleGroupTableNoSanitizeImpl(c, i, ctrl_t::kSentinel);
 }
 
 // Like SetCtrl, but in a single group table, we can save some operations when
 // setting the cloned control byte.
 inline void SetCtrlInSingleGroupTable(const CommonFields& c, size_t i, ctrl_t h,
                                       size_t slot_size) {
-  const size_t cap = c.capacity();
   ABSL_SWISSTABLE_ASSERT(!c.is_small());
-  ABSL_SWISSTABLE_ASSERT(is_single_group(cap));
+  ABSL_SWISSTABLE_ASSERT(is_single_group(c.capacity()));
   DoSanitizeOnSetCtrl(c, i, h, slot_size);
-  ctrl_t* ctrl = c.control();
-  ctrl[i] = h;
-  ctrl[i + cap + 1] = h;
+  SetCtrlInSingleGroupTableNoSanitizeImpl(c, i, h);
 }
 // Overload for setting to an occupied `h2_t` rather than a special `ctrl_t`.
 inline void SetCtrlInSingleGroupTable(const CommonFields& c, size_t i, h2_t h,
@@ -560,14 +560,6 @@
   SetCtrlInLargeTable(c, i, static_cast<ctrl_t>(h), slot_size);
 }
 
-void BlockControlBytes(CommonFields& common, size_t blocked_element_count) {
-  const size_t capacity = common.capacity();
-  while (blocked_element_count > 0) {
-    BlockCtrl(common, capacity - blocked_element_count);
-    --blocked_element_count;
-  }
-}
-
 size_t DropDeletesWithoutResizeAndPrepareInsert(
     CommonFields& common, const PolicyFunctions& __restrict policy,
     size_t new_hash) {
@@ -593,9 +585,7 @@
   //       mark target as FULL
   //       repeat procedure for current slot with moved from element (target)
   ctrl_t* ctrl = common.control();
-  const size_t blocked_element_count = common.blocked_element_count();
   ConvertDeletedToEmptyAndFullToDeleted(ctrl, capacity);
-  BlockControlBytes(common, blocked_element_count);
   const void* hash_fn = policy.hash_fn(common);
   auto hasher = policy.hash_slot;
   auto transfer_n = policy.transfer_n;
@@ -672,8 +662,8 @@
   }
   // Prepare insert for the new element.
   PrepareInsertCommon(common);
-  ResetGrowthLeft(common.growth_info(), capacity,
-                  common.size() + blocked_element_count);
+  ABSL_SWISSTABLE_ASSERT(common.blocked_element_count() == 0);
+  ResetGrowthLeft(common.growth_info(), capacity, common.size());
   FindInfo find_info = find_first_non_full(common, new_hash);
   SetCtrlInLargeTable(common, find_info.offset, H2(new_hash), slot_size);
   common.infoz().RecordInsertMiss(new_hash, find_info.probe_length);
@@ -725,7 +715,10 @@
   ctrl[capacity] = ctrl_t::kSentinel;
   SanitizerPoisonMemoryRegion(common.slot_array(),
                               slot_size * (capacity - blocked_element_count));
-  BlockControlBytes(common, blocked_element_count);
+  while (blocked_element_count > 0) {
+    BlockCtrlInSingleGroupTable(common, capacity - blocked_element_count);
+    --blocked_element_count;
+  }
 }
 
 // Initializes control bytes for growing from capacity 1 to 3.
@@ -1881,13 +1874,11 @@
 size_t RehashOrGrowToNextCapacityAndPrepareInsert(
     CommonFields& common, const PolicyFunctions& __restrict policy,
     size_t new_hash) {
-  ABSL_SWISSTABLE_ASSERT(
-      !common.growth_info().GetGrowthInfoLowerBound().HasNoDeleted());
   const size_t cap = common.capacity();
   ABSL_ASSUME(cap > 0);
-  // Do these calculations in 64-bit to avoid overflow.
-  if (common.size() * uint64_t{32} <=
-      (cap - kMaxBlockedElementsForLargeTables) * uint64_t{25}) {
+  if (cap > Group::kWidth &&
+      // Do these calculations in 64-bit to avoid overflow.
+      common.size() * uint64_t{32} <= cap * uint64_t{25}) {
     // Squash DELETED without growing if there is enough capacity.
     //
     // Rehash in place if the current size is <= 25/32 of capacity.
@@ -1992,17 +1983,15 @@
 
 // Returns the number of elements to block for the given capacity and reserved
 // size.
-size_t BlockedElementCountForReservedTable(size_t capacity,
-                                           size_t reserved_size) {
+size_t BlockedElementCount(size_t capacity, size_t reserved_size) {
   if (!IsCapacityValidForBlockedElements(capacity)) {
     return 0;
   }
-  const size_t blocked_elements = CapacityToGrowth(capacity) - reserved_size;
-  if (is_single_group(capacity)) {
-    // Single group tables never probes, so we can block all the slots.
-    return blocked_elements;
-  }
-  return (std::min)(blocked_elements, kMaxBlockedElementsForLargeTables);
+  ABSL_SWISSTABLE_ASSERT(is_single_group(capacity));
+  const size_t result = CapacityToGrowth(capacity) - reserved_size;
+  ABSL_SWISSTABLE_ASSERT(result <=
+                         HashtableInlineData::kMaxBlockedElementCount);
+  return result;
 }
 
 // Resizes empty non-allocated table to the capacity to fit new_size elements.
@@ -2017,10 +2006,9 @@
   ValidateMaxSize(new_size, policy.key_size, policy.slot_size);
   ABSL_ASSUME(new_size > 0);
   const size_t new_capacity = SizeToCapacity(new_size);
-  ResizeEmptyNonAllocatedTableImpl(
-      common, policy, new_capacity,
-      BlockedElementCountForReservedTable(new_capacity, new_size),
-      /*force_infoz=*/false);
+  ResizeEmptyNonAllocatedTableImpl(common, policy, new_capacity,
+                                   BlockedElementCount(new_capacity, new_size),
+                                   /*force_infoz=*/false);
   // This is after resize, to ensure that we have completed the allocation
   // and have potentially sampled the hashtable.
   common.infoz().RecordReservation(new_size);
diff --git a/absl/container/internal/raw_hash_set.h b/absl/container/internal/raw_hash_set.h
index c6fb30e..62ead86 100644
--- a/absl/container/internal/raw_hash_set.h
+++ b/absl/container/internal/raw_hash_set.h
@@ -367,6 +367,10 @@
 constexpr size_t SooCapacity() { return 1; }
 // Maximum capacity of a table where we don't need to hash any keys.
 constexpr size_t MaxSmallCapacity() { return 1; }
+// Maximum capacity of a table where we can use blocked elements.
+constexpr size_t MaxCapacityWithBlockedElements() {
+  return Group::kWidth - 1;
+}
 // Sentinel type to indicate SOO CommonFields construction.
 struct soo_tag_t {};
 // Sentinel type to indicate SOO CommonFields construction with full size.
@@ -397,7 +401,7 @@
 // Whether `cap` is a valid capacity for a table that can store blocked
 // elements.
 constexpr bool IsCapacityValidForBlockedElements(size_t cap) {
-  return !IsSmallCapacity(cap);
+  return !IsSmallCapacity(cap) && cap <= MaxCapacityWithBlockedElements();
 }
 
 // Converts `n` into the next valid capacity, per `IsValidCapacity`.
@@ -775,12 +779,6 @@
 using PerTableSeed = HashtableInlineData::PerTableSeed;
 using HashtableCapacity = HashtableInlineData::HashtableCapacity;
 
-// For large tables, we limit the number of blocked elements to maintain O(1)
-// average case lookup complexity.
-constexpr size_t kMaxBlockedElementsForLargeTables = 5;
-static_assert(kMaxBlockedElementsForLargeTables <=
-              HashtableInlineData::kMaxBlockedElementCount);
-
 // H1 is just the low bits of the hash.
 inline size_t H1(size_t hash) { return hash; }
 
diff --git a/absl/container/internal/raw_hash_set_test.cc b/absl/container/internal/raw_hash_set_test.cc
index cd87dfc..e1dbac7 100644
--- a/absl/container/internal/raw_hash_set_test.cc
+++ b/absl/container/internal/raw_hash_set_test.cc
@@ -1569,65 +1569,6 @@
   }
 }
 
-// This test verifies that we don't rehash in place when we insert an element
-// above the growth left threshold. Otherwise we may end up with zero empty
-// slots. That would cause hard to debug infinite loop in `find`.
-// This test do the following:
-// 1. Reserve a table with `kReserveSize` elements.
-// 2. Insert `Group::kWidth` elements to fill the first group (due to bad hash
-//    function all elements are inserted into the same group).
-// 3. Erase one element to create tombstone.
-// 4. Insert the same element back. But GrowthInfo still assumes that we may
-//    have a tombstone in the table.
-// 5. Insert one more element, which should cause a rehash and growth.
-TEST(Table,
-     ReservedTableResizeNotRehashInplaceIfInsertingElementAboveGrowthLeft) {
-  if (SwisstableGenerationsEnabled()) {
-    GTEST_SKIP() << "Generations enabled, so rehash happens earlier.\n"
-                 << "Note that reservation doesn't prevent rehashing since we "
-                    "are erasing one element.";
-  }
-  constexpr int64_t kCoef = 17;
-  constexpr size_t kCapacity = 31;
-  constexpr size_t kReserveSize =
-      CapacityToGrowth(kCapacity) - kMaxBlockedElementsForLargeTables;
-
-  BadTwoValuesHashTable t(0,
-                          // Negative number goes to the end of the table.
-                          BadTwoValuesHash(kReserveSize + 2));
-  // Remove seed to make table layout deterministic.
-  RawHashSetTestOnlyAccess::GetCommon(t).set_no_seed_for_testing();
-
-  t.reserve(kReserveSize);
-  for (int64_t i = 0; i < static_cast<int64_t>(Group::kWidth); ++i) {
-    ASSERT_TRUE(t.insert(i * kCoef).second);
-  }
-  EXPECT_EQ(t.erase(kCoef), 1);
-  EXPECT_EQ(RawHashSetTestOnlyAccess::CountTombstones(t), 1);
-  EXPECT_TRUE(t.insert(kCoef).second);
-  EXPECT_EQ(RawHashSetTestOnlyAccess::CountTombstones(t), 0);
-  // We want to test codepath deciding whether to rehash in place or not.
-  // For this we need to potentially have tombstone.
-  EXPECT_FALSE(RawHashSetTestOnlyAccess::GetCommon(t)
-                   .growth_info()
-                   .GetGrowthInfoLowerBound()
-                   .HasNoDeleted());
-  for (int64_t i = static_cast<int64_t>(Group::kWidth);
-       i < static_cast<int64_t>(kReserveSize); ++i) {
-    ASSERT_TRUE(t.insert(i * kCoef).second);
-  }
-  EXPECT_EQ(t.size(), kReserveSize);
-  EXPECT_EQ(t.capacity(), kCapacity);
-  EXPECT_TRUE(t.insert(-57).second);
-  EXPECT_EQ(t.size(), kReserveSize + 1);
-  EXPECT_EQ(RawHashSetTestOnlyAccess::CountTombstones(t), 0);
-  EXPECT_EQ(t.capacity(), NextCapacity(kCapacity));
-  for (int64_t i = 0; i < static_cast<int64_t>(kReserveSize); ++i) {
-    ASSERT_TRUE(t.contains(i * kCoef));
-  }
-  EXPECT_TRUE(t.contains(-57));
-}
-
 TYPED_TEST(SooTest, EraseInSmallTables) {
   for (int64_t size = 0; size < 64; ++size) {
     TypeParam t;