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: 938434294
Change-Id: Ica63d6bf63fe7a1c4b864db7163ca5f94166540c
diff --git a/absl/container/internal/raw_hash_set.cc b/absl/container/internal/raw_hash_set.cc
index 45b3fa6..ca3ddec 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);
-  ctrl_t* ctrl = c.control();
-  ctrl[i] = h;
-  ctrl[((i - NumClonedBytes()) & c.capacity()) +
-       (NumClonedBytes() & c.capacity())] = h;
+  SetCtrlNoSanitizeImpl(c, i, 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);
 }
 
-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.
+// Sets `ctrl[i]` to `ctrl_t::kSentinel`.
 //
 // Unlike setting it directly, this function will perform bounds checks and
 // mirror the value to the cloned tail if necessary.
-inline void BlockCtrlInSingleGroupTable(const CommonFields& c, size_t i) {
-  SetCtrlInSingleGroupTableNoSanitizeImpl(c, i, ctrl_t::kSentinel);
+inline void BlockCtrl(const CommonFields& c, size_t i) {
+  ABSL_SWISSTABLE_ASSERT(!c.is_small());
+  SetCtrlNoSanitizeImpl(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(c.capacity()));
+  ABSL_SWISSTABLE_ASSERT(is_single_group(cap));
   DoSanitizeOnSetCtrl(c, i, h, slot_size);
-  SetCtrlInSingleGroupTableNoSanitizeImpl(c, i, h);
+  ctrl_t* ctrl = c.control();
+  ctrl[i] = h;
+  ctrl[i + cap + 1] = 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,6 +560,14 @@
   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) {
@@ -585,7 +593,9 @@
   //       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;
@@ -662,8 +672,8 @@
   }
   // Prepare insert for the new element.
   PrepareInsertCommon(common);
-  ABSL_SWISSTABLE_ASSERT(common.blocked_element_count() == 0);
-  ResetGrowthLeft(common.growth_info(), capacity, common.size());
+  ResetGrowthLeft(common.growth_info(), capacity,
+                  common.size() + blocked_element_count);
   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);
@@ -715,10 +725,7 @@
   ctrl[capacity] = ctrl_t::kSentinel;
   SanitizerPoisonMemoryRegion(common.slot_array(),
                               slot_size * (capacity - blocked_element_count));
-  while (blocked_element_count > 0) {
-    BlockCtrlInSingleGroupTable(common, capacity - blocked_element_count);
-    --blocked_element_count;
-  }
+  BlockControlBytes(common, blocked_element_count);
 }
 
 // Initializes control bytes for growing from capacity 1 to 3.
@@ -1874,11 +1881,13 @@
 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);
-  if (cap > Group::kWidth &&
-      // Do these calculations in 64-bit to avoid overflow.
-      common.size() * uint64_t{32} <= cap * uint64_t{25}) {
+  // Do these calculations in 64-bit to avoid overflow.
+  if (common.size() * uint64_t{32} <=
+      (cap - kMaxBlockedElementsForLargeTables) * uint64_t{25}) {
     // Squash DELETED without growing if there is enough capacity.
     //
     // Rehash in place if the current size is <= 25/32 of capacity.
@@ -1983,15 +1992,17 @@
 
 // Returns the number of elements to block for the given capacity and reserved
 // size.
-size_t BlockedElementCount(size_t capacity, size_t reserved_size) {
+size_t BlockedElementCountForReservedTable(size_t capacity,
+                                           size_t reserved_size) {
   if (!IsCapacityValidForBlockedElements(capacity)) {
     return 0;
   }
-  ABSL_SWISSTABLE_ASSERT(is_single_group(capacity));
-  const size_t result = CapacityToGrowth(capacity) - reserved_size;
-  ABSL_SWISSTABLE_ASSERT(result <=
-                         HashtableInlineData::kMaxBlockedElementCount);
-  return result;
+  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);
 }
 
 // Resizes empty non-allocated table to the capacity to fit new_size elements.
@@ -2006,9 +2017,10 @@
   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,
-                                   BlockedElementCount(new_capacity, new_size),
-                                   /*force_infoz=*/false);
+  ResizeEmptyNonAllocatedTableImpl(
+      common, policy, new_capacity,
+      BlockedElementCountForReservedTable(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 62ead86..c6fb30e 100644
--- a/absl/container/internal/raw_hash_set.h
+++ b/absl/container/internal/raw_hash_set.h
@@ -367,10 +367,6 @@
 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.
@@ -401,7 +397,7 @@
 // Whether `cap` is a valid capacity for a table that can store blocked
 // elements.
 constexpr bool IsCapacityValidForBlockedElements(size_t cap) {
-  return !IsSmallCapacity(cap) && cap <= MaxCapacityWithBlockedElements();
+  return !IsSmallCapacity(cap);
 }
 
 // Converts `n` into the next valid capacity, per `IsValidCapacity`.
@@ -779,6 +775,12 @@
 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 e1dbac7..cd87dfc 100644
--- a/absl/container/internal/raw_hash_set_test.cc
+++ b/absl/container/internal/raw_hash_set_test.cc
@@ -1569,6 +1569,65 @@
   }
 }
 
+// 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;