Pass capacity to CommonFields::slot_array when it is already in register and/or known at compile time.

Change is saving ~400 bytes of binary size in raw_hash_set.cc. Some of the functions there quite hot.

PiperOrigin-RevId: 945647274
Change-Id: I5522884c1da91b83cd4bdf5f346f99b2bdd3085d
diff --git a/absl/container/internal/raw_hash_set.cc b/absl/container/internal/raw_hash_set.cc
index 3fda249..d37c781 100644
--- a/absl/container/internal/raw_hash_set.cc
+++ b/absl/container/internal/raw_hash_set.cc
@@ -237,9 +237,9 @@
 template <class Fn>
 void IterateOverFullSlotsImpl(const CommonFields& c, size_t slot_size, Fn cb) {
   const size_t cap = c.capacity();
-  ABSL_SWISSTABLE_ASSERT(!IsSmallCapacity(cap));
+  ABSL_ASSUME(cap > kMaxSmallCapacity);
   const ctrl_t* ctrl = c.control();
-  void* slot = c.slot_array();
+  void* slot = c.slot_array(cap);
   if (is_half_group(cap)) {
     // Mirrored/cloned control bytes in half-group table are also located in the
     // first group (starting from position 0). We are taking group from position
@@ -506,8 +506,10 @@
 // Sets sanitizer poisoning for slot corresponding to control byte being set.
 inline void DoSanitizeOnSetCtrl(const CommonFields& c, size_t i, ctrl_t h,
                                 size_t slot_size) {
-  ABSL_SWISSTABLE_ASSERT(i < c.capacity());
-  auto* slot_i = static_cast<const char*>(c.slot_array()) + i * slot_size;
+  const size_t cap = c.capacity();
+  ABSL_ASSUME(cap > kMaxSmallCapacity);
+  ABSL_SWISSTABLE_ASSERT(i < cap);
+  auto* slot_i = static_cast<const char*>(c.slot_array(cap)) + i * slot_size;
   if (IsFull(h)) {
     SanitizerUnpoisonMemoryRegion(slot_i, slot_size);
   } else {
@@ -593,10 +595,13 @@
     CommonFields& common, const PolicyFunctions& __restrict policy,
     size_t new_hash) {
   void* set = &common;
-  void* slot_array = common.slot_array();
   const size_t capacity = common.capacity();
   ABSL_SWISSTABLE_ASSERT(IsValidCapacity(capacity));
   ABSL_SWISSTABLE_ASSERT(!is_single_group(capacity));
+  ABSL_ASSUME(capacity > kMaxSmallCapacity);
+
+  ctrl_t* ctrl = common.control();
+  void* slot_array = common.slot_array(capacity);
   // Algorithm:
   // - mark all DELETED slots as EMPTY
   // - mark all FULL slots as DELETED
@@ -613,7 +618,6 @@
   //       swap current element with target element
   //       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);
@@ -744,7 +748,7 @@
                 capacity + 1 + NumClonedBytes());
   }
   ctrl[capacity] = ctrl_t::kSentinel;
-  SanitizerPoisonMemoryRegion(common.slot_array(),
+  SanitizerPoisonMemoryRegion(common.slot_array(capacity),
                               slot_size * (capacity - blocked_element_count));
   BlockControlBytes(common, blocked_element_count);
 }
@@ -836,7 +840,7 @@
 
 template <bool kSooEnabled>
 void* SingleSlotAddress(CommonFields& c) {
-  return kSooEnabled ? c.soo_data() : c.slot_array();
+  return kSooEnabled ? c.soo_data() : c.slot_array(/*capacity=*/1);
 }
 
 template <bool kSooEnabled>
@@ -858,7 +862,8 @@
   }
   c.decrement_size();
   c.infoz().RecordErase();
-  SanitizerPoisonMemoryRegion(c.slot_array(), slot_size);
+  SanitizerPoisonMemoryRegion(SingleSlotAddress</*kSooEnabled=*/false>(c),
+                              slot_size);
 }
 
 void EraseMetaOnlyLarge(CommonFields& c, size_t index, size_t slot_size) {
@@ -881,7 +886,7 @@
 void ClearBackingArray(CommonFields& c,
                        const PolicyFunctions& __restrict policy, void* alloc,
                        bool reuse) {
-  ABSL_SWISSTABLE_ASSERT(c.capacity() > MaxSmallCapacity());
+  ABSL_SWISSTABLE_ASSERT(c.capacity() > kMaxSmallCapacity);
   if (reuse) {
     const size_t blocked_element_count = c.blocked_element_count();
     c.set_size_to_zero();
@@ -978,7 +983,8 @@
   if (destroy_slot != nullptr) {
     if (c.is_small()) {
       if (!c.empty()) {
-        destroy_slot(&c, c.slot_array());
+        static_assert(kMaxSmallCapacity == 1);
+        destroy_slot(&c, c.slot_array(/*capacity=*/1));
       }
     } else {
       DestroySlots(c, slot_size, destroy_slot);
@@ -997,7 +1003,7 @@
 size_t FindNewPositionsAndTransferSlots(
     CommonFields& common, const PolicyFunctions& __restrict policy,
     ctrl_t* old_ctrl, void* old_slots, size_t old_capacity) {
-  void* new_slots = common.slot_array();
+  void* new_slots = common.slot_array(common.capacity());
   const void* hash_fn = policy.hash_fn(common);
   const size_t slot_size = policy.slot_size;
   const size_t seed = common.seed().seed();
@@ -1378,7 +1384,7 @@
     const ProbedItem* start, const ProbedItem* end, void* old_slots) {
   const HashtableCapacity new_capacity = c.capacity_impl();
 
-  void* new_slots = c.slot_array();
+  void* new_slots = c.slot_array(new_capacity.capacity());
   ctrl_t* new_ctrl = c.control();
   size_t total_probe_length = 0;
 
@@ -1416,7 +1422,7 @@
     void* old_slots, size_t start) {
   size_t old_capacity = PreviousCapacity(c.capacity());
   const size_t slot_size = policy.slot_size;
-  void* new_slots = c.slot_array();
+  void* new_slots = c.slot_array(c.capacity());
   size_t total_probe_length = 0;
   const void* hash_fn = policy.hash_fn(c);
   auto hash_slot = policy.hash_slot;
@@ -1692,7 +1698,8 @@
                               const PolicyFunctions& __restrict policy) {
   ABSL_SWISSTABLE_ASSERT(common.is_small());
   common.increment_size();
-  SanitizerUnpoisonMemoryRegion(common.slot_array(), policy.slot_size);
+  SanitizerUnpoisonMemoryRegion(
+      SingleSlotAddress</*kSooEnabled=*/false>(common), policy.slot_size);
 }
 
 void IncrementSmallSize(CommonFields& common,
@@ -1776,10 +1783,11 @@
       common.growth_info().GetGrowthLeftTotalSlow(old_capacity) == 0);
   ABSL_SWISSTABLE_ASSERT(old_capacity > policy.soo_capacity());
   ABSL_SWISSTABLE_ASSERT(!IsSmallCapacity(old_capacity));
+  ABSL_ASSUME(old_capacity > kMaxSmallCapacity);
 
   const size_t new_capacity = NextCapacity(old_capacity);
   ctrl_t* old_ctrl = common.control();
-  void* old_slots = common.slot_array();
+  void* old_slots = common.slot_array(old_capacity);
   size_t old_blocked_element_count = common.blocked_element_count();
 
   HashtablezInfoHandle infoz = common.infoz();
@@ -1854,7 +1862,7 @@
         common.infoz().RecordInsertMiss(get_hash(common.seed().seed()),
                                         /*distance_from_desired=*/0);
       }
-      return common.slot_array();
+      return common.slot_array(/*capacity=*/1);
     } else {
       return Grow1To3AndPrepareInsert(common, policy, get_hash);
     }
@@ -1962,6 +1970,7 @@
                              size_t hash) {
   GrowthInfoAccessor growth_info = common.growth_info();
   const size_t cap = common.capacity();
+  ABSL_ASSUME(cap > kMaxSmallCapacity);
   GrowthInfoLowerBound growth_info_lower_bound =
       growth_info.RebalanceGrowthLeftLowerBound(cap);
   if (ABSL_PREDICT_TRUE(
@@ -1983,7 +1992,7 @@
   growth_info.OverwriteControlAsFull(common.control()[target.offset]);
   SetCtrlInLargeTable(common, target.offset, H2(hash), policy.slot_size);
   common.infoz().RecordInsertMiss(hash, target.probe_length);
-  return SlotAddress(common.slot_array(), target.offset, policy.slot_size);
+  return SlotAddress(common.slot_array(cap), target.offset, policy.slot_size);
 }
 
 // Resizes empty non-allocated SOO table to NextCapacity(SooCapacity()),
@@ -1996,7 +2005,8 @@
 GrowEmptySooTableToNextCapacityForceSamplingAndPrepareInsert(
     CommonFields& common, const PolicyFunctions& __restrict policy,
     absl::FunctionRef<size_t(size_t)> get_hash) {
-  ResizeEmptyNonAllocatedTableImpl(common, policy, NextCapacity(SooCapacity()),
+  const size_t kNewCapacity = NextCapacity(SooCapacity());
+  ResizeEmptyNonAllocatedTableImpl(common, policy, kNewCapacity,
                                    /*blocked_element_count=*/0,
                                    /*force_infoz=*/true);
   PrepareInsertCommon(common);
@@ -2005,7 +2015,8 @@
   SetCtrlInSingleGroupTable(common, SooSlotIndex(), H2(new_hash),
                             policy.slot_size);
   common.infoz().RecordInsertMiss(new_hash, /*distance_from_desired=*/0);
-  return SlotAddress(common.slot_array(), SooSlotIndex(), policy.slot_size);
+  return SlotAddress(common.slot_array(kNewCapacity), SooSlotIndex(),
+                     policy.slot_size);
 }
 
 // Returns the number of elements to block for the given capacity and reserved
@@ -2098,7 +2109,7 @@
 
   const size_t old_capacity = common.capacity();
   ctrl_t* const old_ctrl = common.control();
-  void* const old_slots = common.slot_array();
+  void* const old_slots = common.slot_array(old_capacity);
   const size_t old_blocked_element_count = common.blocked_element_count();
 
   const size_t slot_size = policy.slot_size;
@@ -2236,7 +2247,7 @@
       size_t begin_offset = FindFirstFullSlot(0, cap, common.control());
       policy.transfer_n(
           &common, &tmp_slot,
-          SlotAddress(common.slot_array(), begin_offset, slot_size), 1);
+          SlotAddress(common.slot_array(cap), begin_offset, slot_size), 1);
       clear_backing_array();
       policy.transfer_n(&common, common.soo_data(), &tmp_slot, 1);
       common.set_full_soo();
@@ -2283,12 +2294,14 @@
     const size_t other_capacity = other.capacity();
     const void* other_slot =
         other_capacity <= soo_capacity ? other.soo_data()
-        : other.is_small()
-            ? other.slot_array()
-            : SlotAddress(other.slot_array(),
+        : IsSmallCapacity(other_capacity)
+            ? other.slot_array(other_capacity)
+            : SlotAddress(other.slot_array(other_capacity),
                           FindFirstFullSlot(0, other_capacity, other.control()),
                           slot_size);
-    copy_fn(soo_enabled ? common.soo_data() : common.slot_array(), other_slot);
+    copy_fn(soo_enabled ? common.soo_data()
+                        : SingleSlotAddress</*kSooEnabled=*/false>(common),
+            other_slot);
 
     if (soo_enabled && policy.is_hashtablez_eligible &&
         ShouldSampleNextTable()) {
@@ -2303,10 +2316,12 @@
   ABSL_SWISSTABLE_ASSERT(other.capacity() > soo_capacity);
   const size_t cap = common.capacity();
   ABSL_SWISSTABLE_ASSERT(cap > soo_capacity);
+  ABSL_ASSUME(cap > kMaxSmallCapacity);
   size_t offset = cap;
   const void* hash_fn = policy.hash_fn(common);
   auto hasher = policy.hash_slot;
   const size_t seed = common.seed().seed();
+  void* target_slot_array = common.slot_array(cap);
   IterateOverFullSlotsImpl(
       other, slot_size, [&](const ctrl_t*, void* that_slot) {
         // The table is guaranteed to be empty, so we can do faster than
@@ -2316,7 +2331,7 @@
         infoz.RecordInsertMiss(hash, target.probe_length);
         offset = target.offset;
         SetCtrl(common, offset, H2(hash), slot_size);
-        copy_fn(SlotAddress(common.slot_array(), offset, slot_size), that_slot);
+        copy_fn(SlotAddress(target_slot_array, offset, slot_size), that_slot);
         common.maybe_increment_generation_on_insert();
       });
   common.increment_size(size);
@@ -2365,11 +2380,13 @@
   }
   PrepareInsertCommon(common);
   growth_info.OverwriteEmptyAsFull();
+  const size_t cap = common.capacity();
+  ABSL_ASSUME(cap > kMaxSmallCapacity);
   target_group.offset += mask_empty.LowestBitSet();
-  target_group.offset &= common.capacity();
+  target_group.offset &= cap;
   SetCtrl(common, target_group.offset, H2(hash), policy.slot_size);
   common.infoz().RecordInsertMiss(hash, target_group.probe_length);
-  return SlotAddress(common.slot_array(), target_group.offset,
+  return SlotAddress(common.slot_array(cap), target_group.offset,
                      policy.slot_size);
 }
 }  // namespace
diff --git a/absl/container/internal/raw_hash_set.h b/absl/container/internal/raw_hash_set.h
index dbed05f..3c25b08 100644
--- a/absl/container/internal/raw_hash_set.h
+++ b/absl/container/internal/raw_hash_set.h
@@ -377,7 +377,7 @@
 //   tables, we would need to randomize the iteration order somehow.
 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; }
+inline constexpr size_t kMaxSmallCapacity = 1;
 // Sentinel type to indicate SOO CommonFields construction.
 struct soo_tag_t {};
 // Sentinel type to indicate SOO CommonFields construction with full size.
@@ -396,7 +396,7 @@
 
 // Whether a table is small enough that we don't need to hash any keys.
 constexpr bool IsSmallCapacity(size_t capacity) {
-  return capacity <= MaxSmallCapacity();
+  return capacity <= kMaxSmallCapacity;
 }
 
 // Whether a table fits entirely into a probing group.
@@ -555,7 +555,7 @@
     // Comparing capacity_data_ directly leads to a better generated code.
     // One byte comparison is used before computing the capacity in order to
     // detect small tables faster for critical path.
-    static_assert(MaxSmallCapacity() == 1);
+    static_assert(kMaxSmallCapacity == 1);
     return capacity_data_ <= 1;
   }
 
@@ -1284,11 +1284,11 @@
   void set_control(ctrl_t* c) { heap_or_soo_.control().set(c); }
 
   // Note: we can't use slots() because Qt defines "slots" as a macro.
-  void* slot_array() const { return slot_array(capacity()); }
   // Returns pointer to the slots of a table with explicit capacity that must be
   // equal to the actual capacity of the table.
-  // Useful for capacity known at compile time and capacity in register with
-  // ABSL_ASSUME conditions.
+  // Capacity is often known at compile time or already in register with some
+  // ABSL_ASSUME conditions. We require passing it explicitly to eliminate
+  // branches inside of NumControlBytes in majority of cases.
   void* slot_array(size_t capacity) const {
     ABSL_SWISSTABLE_ASSERT(capacity == this->capacity());
     ctrl_t* ctrl = control();
@@ -2704,7 +2704,8 @@
   iterator begin() ABSL_ATTRIBUTE_LIFETIME_BOUND {
     if (ABSL_PREDICT_FALSE(empty())) return end();
     if (is_small()) return single_iterator();
-    iterator it = {control(), slot_array(), common().generation_ptr()};
+    iterator it = {control(), slot_array(capacity()),
+                   common().generation_ptr()};
     it.skip_empty_or_deleted();
     ABSL_SWISSTABLE_ASSERT(IsFull(*it.control()));
     return it;
@@ -3159,7 +3160,7 @@
     if (is_small()) return;
     auto seq = probe(common(), hash_of(key));
     PrefetchToLocalCache(control() + seq.offset());
-    PrefetchToLocalCache(slot_array() + seq.offset());
+    PrefetchToLocalCache(slot_array(capacity()) + seq.offset());
 #endif  // ABSL_HAVE_PREFETCH
   }
 
@@ -3343,7 +3344,7 @@
   iterator find_large(const key_arg<K>& key) {
     ABSL_SWISSTABLE_ASSERT(!is_small());
     const size_t cap = common().capacity();
-    ABSL_ASSUME(cap > 1);
+    ABSL_ASSUME(cap > kMaxSmallCapacity);
     const size_t hash = hash_of(key);
     auto seq = probe(ProbeCapacity{cap}, hash);
     const h2_t h2 = H2(hash);
@@ -3388,7 +3389,7 @@
   }
 
   void clear_backing_array(bool reuse) {
-    ABSL_SWISSTABLE_ASSERT(capacity() > MaxSmallCapacity());
+    ABSL_SWISSTABLE_ASSERT(capacity() > kMaxSmallCapacity);
     ClearBackingArray(common(), GetPolicyFunctions(), &char_alloc_ref(), reuse);
   }
 
@@ -3443,7 +3444,7 @@
     EraseMetaOnlyLarge(common(),
                        // `it` can be non-iterable iterator, so we can't use
                        // it.control().
-                       static_cast<size_t>(it.slot() - slot_array()),
+                       static_cast<size_t>(it.slot() - slot_array(capacity())),
                        sizeof(slot_type));
   }
 
@@ -3628,7 +3629,7 @@
     ABSL_SWISSTABLE_ASSERT(!is_soo());
     prefetch_heap_block();
     const size_t cap = capacity();
-    ABSL_ASSUME(cap > 1);
+    ABSL_ASSUME(cap > kMaxSmallCapacity);
     const size_t hash = hash_of(key);
     auto seq = probe(ProbeCapacity{cap}, hash);
     const h2_t h2 = H2(hash);
@@ -3796,9 +3797,9 @@
     ABSL_SWISSTABLE_ASSERT(!is_soo());
     return common().control();
   }
-  slot_type* slot_array() const {
+  slot_type* slot_array(size_t capacity) const {
     ABSL_SWISSTABLE_ASSERT(!is_soo());
-    return static_cast<slot_type*>(common().slot_array());
+    return static_cast<slot_type*>(common().slot_array(capacity));
   }
   slot_type* soo_slot() {
     ABSL_SWISSTABLE_ASSERT(is_soo());
@@ -3884,6 +3885,7 @@
       void (*encode_probed_element)(void* probed_storage, h2_t h2,
                                     size_t source_offset, size_t h1)) {
     const size_t new_capacity = common.capacity();
+    ABSL_ASSUME(new_capacity > kMaxSmallCapacity);
     const size_t old_capacity = PreviousCapacity(new_capacity);
     ABSL_ASSUME(old_capacity + 1 >= Group::kWidth);
     ABSL_ASSUME((old_capacity + 1) % Group::kWidth == 0);
@@ -3891,7 +3893,7 @@
     auto* set = reinterpret_cast<raw_hash_set*>(&common);
     slot_type* old_slots_ptr = to_slot(old_slots);
     ctrl_t* new_ctrl = common.control();
-    slot_type* new_slots = set->slot_array();
+    slot_type* new_slots = set->slot_array(new_capacity);
 
     for (size_t group_index = 0; group_index < old_capacity;
          group_index += Group::kWidth) {
@@ -4074,7 +4076,7 @@
     while (true) {
       container_internal::Group g{ctrl + seq.offset()};
       for (uint32_t i : g.Match(h2)) {
-        if (set.equal_to(key, set.slot_array() + seq.offset(i)))
+        if (set.equal_to(key, set.slot_array(set.capacity()) + seq.offset(i)))
           return num_probes;
         ++num_probes;
       }
diff --git a/absl/container/internal/raw_hash_set_test.cc b/absl/container/internal/raw_hash_set_test.cc
index fe987cd..e97648d 100644
--- a/absl/container/internal/raw_hash_set_test.cc
+++ b/absl/container/internal/raw_hash_set_test.cc
@@ -82,8 +82,8 @@
     return std::forward<C>(c).common();
   }
   template <typename C>
-  static auto GetSlots(const C& c) -> decltype(c.slot_array()) {
-    return c.slot_array();
+  static auto GetSlots(const C& c) -> decltype(c.slot_array(c.capacity())) {
+    return c.slot_array(c.capacity());
   }
   template <typename C>
   static size_t CountTombstones(const C& c) {