Move second-chance lookup in TryFindNewIndexWithoutProbing into type-erased encode_probed_element to reduce binary size. Previously, `TryFindNewIndexWithoutProbing` (inlined into every template instantiation of `raw_hash_set<...>::transfer_unprobed_elements_to_next_capacity_fn`) contained both the unprobed element fast path and a second-chance SIMD group lookup for elements that were probed in the old table. Because `transfer_unprobed_elements_to_next_capacity_fn` is instantiated per container policy, this second-chance SIMD lookup was duplicated across every hash table instantiation. Move the second-chance lookup into `ProbedItemEncoder::EncodeProbedElement` (invoked from the type-erased `encode_probed_element` callback): - Simplify `TryFindNewIndexWithoutProbing`. PiperOrigin-RevId: 1002102000
diff --git a/absl/container/internal/raw_hash_set.cc b/absl/container/internal/raw_hash_set.cc index 374ab9b..285d269 100644 --- a/absl/container/internal/raw_hash_set.cc +++ b/absl/container/internal/raw_hash_set.cc
@@ -1170,10 +1170,8 @@ BackingArrayPtrs AllocBackingArray(CommonFields& common, const PolicyFunctions& __restrict policy, - size_t new_capacity, bool has_infoz, - void* alloc, size_t blocked_element_count) { - RawHashSetLayout layout(new_capacity, policy.slot_size, policy.slot_align, - has_infoz, blocked_element_count); + const RawHashSetLayout& layout, + void* alloc) { // Perform a direct call in the common case to allow for profile-guided // heap optimization (PGHO) to understand which allocation function is used. constexpr size_t kDefaultAlignment = BackingArrayAlignment(alignof(size_t)); @@ -1193,6 +1191,17 @@ mem + layout.slot_offset()}; } +BackingArrayPtrs AllocBackingArray(CommonFields& common, + const PolicyFunctions& __restrict policy, + size_t new_capacity, bool has_infoz, + void* alloc, size_t blocked_element_count) { + return AllocBackingArray( + common, policy, + RawHashSetLayout(new_capacity, policy.slot_size, policy.slot_align, + has_infoz, blocked_element_count), + alloc); +} + void ResizeEmptyNonAllocatedTableImpl(CommonFields& common, const PolicyFunctions& __restrict policy, size_t new_capacity, @@ -1567,8 +1576,73 @@ class ProbedItemEncoder { public: using ProbedItem = ProbedItemType; - explicit ProbedItemEncoder(ctrl_t* control) : control_(control) {} + ProbedItemEncoder(CommonFields& common, + const PolicyFunctions& __restrict policy, ctrl_t* control, + void* old_slots, size_t new_capacity) + : common_(common), + policy_(policy), + control_(control), + old_slots_(old_slots), + new_capacity_(new_capacity), + new_ctrl_(common.control()), + new_slots_(common.slot_array(new_capacity)) {} + // Transfers item to the first probe group if possible, otherwise encodes it + // into the best available location. + void EncodeProbedElement(size_t hash, size_t source_offset) { + size_t h1 = H1(hash); + const size_t new_capacity = new_capacity_; + const size_t old_capacity = PreviousCapacity(new_capacity); + ABSL_SWISSTABLE_ASSERT(((source_offset - h1) & old_capacity) >= + Group::kWidth); + // Try to insert element into the first probe group. + // new_ctrl is not yet fully initialized so we can't use regular search via + // find_first_non_full. + // + // We can search in the first probe group only if it is located in already + // initialized part of the table. + if (ABSL_PREDICT_TRUE((h1 & old_capacity) < source_offset)) { + ctrl_t* new_ctrl = new_ctrl_; + size_t offset = h1 & new_capacity; + Group new_g(new_ctrl + offset); + if (auto mask = new_g.MaskNonFull(); ABSL_PREDICT_TRUE(mask)) { + size_t new_index = offset + mask.LowestBitSet(); + ABSL_SWISSTABLE_ASSERT((new_index & old_capacity) <= source_offset); + ABSL_SWISSTABLE_ASSERT(IsEmpty(new_ctrl[new_index])); + new_ctrl[new_index] = static_cast<ctrl_t>(H2(hash)); + const size_t slot_size = policy_.slot_size; + void* new_slot = SlotAddress(new_slots_, new_index, slot_size); + void* old_slot = SlotAddress(old_slots_, source_offset, slot_size); + SanitizerUnpoisonMemoryRegion(new_slot, slot_size); + policy_.transfer_n(&common_, new_slot, old_slot, 1); + SanitizerPoisonMemoryRegion(old_slot, slot_size); + return; + } + } + EncodeItem(ProbedItem(H2(hash), source_offset, h1)); + } + + // Decodes information about probed elements from all available sources. + // Finds new position for each element and transfers it to the new slots. + // Returns the total probe length. + size_t DecodeAndInsertToTable() const { + if (pos_ == buffer_) { + return 0; + } + if constexpr (kGuaranteedFitToBuffer) { + return DecodeAndInsertImpl(common_, policy_, buffer_, pos_, old_slots_); + } + size_t total_probe_length = DecodeAndInsertImpl( + common_, policy_, buffer_, + local_buffer_full_ ? buffer_ + kBufferSize : pos_, old_slots_); + if (!local_buffer_full_) { + return total_probe_length; + } + total_probe_length += DecodeAndInsertToTableOverflow(); + return total_probe_length; + } + + private: // Encode item into the best available location. void EncodeItem(ProbedItem item) { if (ABSL_PREDICT_FALSE(!kGuaranteedFitToBuffer && pos_ >= end_)) { @@ -1579,30 +1653,6 @@ ++pos_; } - // Decodes information about probed elements from all available sources. - // Finds new position for each element and transfers it to the new slots. - // Returns the total probe length. - size_t DecodeAndInsertToTable(CommonFields& common, - const PolicyFunctions& __restrict policy, - void* old_slots) const { - if (pos_ == buffer_) { - return 0; - } - if constexpr (kGuaranteedFitToBuffer) { - return DecodeAndInsertImpl(common, policy, buffer_, pos_, old_slots); - } - size_t total_probe_length = DecodeAndInsertImpl( - common, policy, buffer_, - local_buffer_full_ ? buffer_ + kBufferSize : pos_, old_slots); - if (!local_buffer_full_) { - return total_probe_length; - } - total_probe_length += - DecodeAndInsertToTableOverflow(common, policy, old_slots); - return total_probe_length; - } - - private: static ProbedItem* AlignToNextItem(void* ptr) { return reinterpret_cast<ProbedItem*>(AlignUpTo( reinterpret_cast<uintptr_t>(ptr), alignof(ProbedItem))); @@ -1663,33 +1713,37 @@ // marked elements. // Finds new position for each element and transfers it to the new slots. // Returns the total probe length. - ABSL_ATTRIBUTE_NOINLINE size_t DecodeAndInsertToTableOverflow( - CommonFields& common, const PolicyFunctions& __restrict policy, - void* old_slots) const { + ABSL_ATTRIBUTE_NOINLINE size_t DecodeAndInsertToTableOverflow() const { ABSL_SWISSTABLE_ASSERT(local_buffer_full_ && "must not be called when local buffer is not full"); size_t total_probe_length = DecodeAndInsertImpl( - common, policy, OverflowBufferStart(), pos_, old_slots); + common_, policy_, OverflowBufferStart(), pos_, old_slots_); if (ABSL_PREDICT_TRUE(marked_elements_starting_position_ == kNoMarkedElementsSentinel)) { return total_probe_length; } total_probe_length += - ProcessProbedMarkedElements(common, policy, control_, old_slots, + ProcessProbedMarkedElements(common_, policy_, control_, old_slots_, marked_elements_starting_position_); return total_probe_length; } static constexpr size_t kBufferSize = kProbedElementsBufferSize / sizeof(ProbedItem); - ProbedItem buffer_[kBufferSize]; + CommonFields& common_; + const PolicyFunctions& policy_; + ctrl_t* const control_; + void* const old_slots_; + const size_t new_capacity_; + ctrl_t* const new_ctrl_; + void* const new_slots_; // If local_buffer_full_ is false, then pos_/end_ are in the local buffer, // otherwise, they're in the overflow buffer. ProbedItem* pos_ = buffer_; const void* end_ = buffer_ + kBufferSize; - ctrl_t* const control_; size_t marked_elements_starting_position_ = kNoMarkedElementsSentinel; bool local_buffer_full_ = false; + ProbedItem buffer_[kBufferSize]; }; // Grows to next capacity with specified encoder type. @@ -1701,18 +1755,17 @@ const PolicyFunctions& __restrict policy, ctrl_t* old_ctrl, void* old_slots, size_t new_capacity) { - using ProbedItem = typename Encoder::ProbedItem; ABSL_SWISSTABLE_ASSERT(common.capacity() == new_capacity); - ABSL_SWISSTABLE_ASSERT(new_capacity <= ProbedItem::kMaxNewCapacity); - Encoder encoder(old_ctrl); + ABSL_SWISSTABLE_ASSERT(new_capacity <= Encoder::ProbedItem::kMaxNewCapacity); + Encoder encoder(common, policy, old_ctrl, old_slots, new_capacity); policy.transfer_unprobed_elements_to_next_capacity( - common, old_ctrl, old_slots, &encoder, - [](void* probed_storage, h2_t h2, size_t source_offset, size_t h1) { + common, old_ctrl, &encoder, old_slots, PreviousCapacity(new_capacity), + [](void* probed_storage, size_t hash, size_t source_offset) { auto encoder_ptr = static_cast<Encoder*>(probed_storage); - encoder_ptr->EncodeItem(ProbedItem(h2, source_offset, h1)); + encoder_ptr->EncodeProbedElement(hash, source_offset); }); InitializeMirroredControlBytes(common.control(), new_capacity); - return encoder.DecodeAndInsertToTable(common, policy, old_slots); + return encoder.DecodeAndInsertToTable(); } // Grows to next capacity for relatively small tables so that even if all @@ -1807,7 +1860,8 @@ void* Grow1To3AndPrepareInsert(CommonFields& common, const PolicyFunctions& __restrict policy, - void* hasher, HashForKeyFn get_hash) { + void* hasher, HashForKeyFn get_hash, void* alloc, + bool has_infoz, void* new_slots) { // TODO(b/413062340): Refactor to reuse more code with // GrowSooTableToNextCapacityAndPrepareInsert. ABSL_SWISSTABLE_ASSERT(common.capacity() == 1); @@ -1817,21 +1871,8 @@ ABSL_SWISSTABLE_ASSERT(common.blocked_element_count() == 0); constexpr size_t kOldCapacity = 1; constexpr size_t kNewCapacity = NextCapacity(kOldCapacity); - void* old_slots = common.single_non_soo_slot(); - // old_slots == old_ctrl in case of capacity == 1. - ctrl_t* old_ctrl = static_cast<ctrl_t*>(old_slots); - - const size_t slot_size = policy.slot_size; - const size_t slot_align = policy.slot_align; - void* alloc = policy.get_char_alloc(common); - const bool has_infoz = common.has_infoz(); - common.set_capacity(kNewCapacity); - - const auto [new_ctrl, new_slots] = - AllocBackingArray(common, policy, kNewCapacity, has_infoz, alloc, - /*blocked_element_count=*/0); - common.set_control(new_ctrl); - SanitizerPoisonMemoryRegion(new_slots, kNewCapacity * slot_size); + ctrl_t* new_ctrl = + static_cast<ctrl_t*>(new_slots) - NumControlBytes(kNewCapacity); if (ABSL_PREDICT_TRUE(!has_infoz)) { // When we're sampled, we already have a seed. @@ -1839,29 +1880,45 @@ } const size_t new_hash = get_hash(hasher, common.seed().seed()); h2_t new_h2 = H2(new_hash); - size_t orig_hash = - policy.hash_slot(policy.hash_fn(common), old_slots, common.seed().seed()); + // `hash_fn` is in separate statement, so the hash_slot function pointer is + // not loaded before the hash_fn call. + const void* hash_fn = policy.hash_fn(common); + size_t orig_hash = policy.hash_slot( + hash_fn, common.single_non_soo_slot(), common.seed().seed()); size_t offset = Resize1To3NewOffset(new_hash); InitializeThreeElementsControlBytes(H2(orig_hash), new_h2, offset, new_ctrl); + // Read `old_slots` and `slot_size` after the hash calls to avoid keeping + // them live across indirect calls. + void* old_slots = common.single_non_soo_slot(); + // old_slots == old_ctrl in case of capacity == 1. + ctrl_t* old_ctrl = static_cast<ctrl_t*>(old_slots); + common.set_control(new_ctrl); + const size_t slot_size = policy.slot_size; + SanitizerPoisonMemoryRegion(new_slots, kNewCapacity * slot_size); + void* old_element_target = NextSlot(new_slots, slot_size); + // Compute `new_element_target_slot` before `transfer_n` to free up registers + // of `new_slots` and `offset`. + void* new_element_target_slot = SlotAddress(new_slots, offset, slot_size); SanitizerUnpoisonMemoryRegion(old_element_target, slot_size); policy.transfer_n(&common, old_element_target, old_slots, 1); - - void* new_element_target_slot = SlotAddress(new_slots, offset, slot_size); SanitizerUnpoisonMemoryRegion(new_element_target_slot, slot_size); + // Group `common` updates together so `inline_data_` is modified in a single + // load-store sequence. + common.set_capacity(kNewCapacity); + common.InitGrowthLeftNoDeleted(kNewCapacity - 2, kNewCapacity); PrepareInsertCommon(common); ABSL_SWISSTABLE_ASSERT(common.size() == 2); - common.InitGrowthLeftNoDeleted(kNewCapacity - 2, kNewCapacity); if (ABSL_PREDICT_FALSE(has_infoz)) { ReportSingleGroupTableGrowthToInfoz(common, old_ctrl, new_hash); } // Deallocate after `ReportSingleGroupTableGrowthToInfoz`, because it reads // infoz from old_ctrl. - DeallocBackingArrayImpl(alloc, kOldCapacity, old_ctrl, slot_size, slot_align, - has_infoz, + DeallocBackingArrayImpl(alloc, kOldCapacity, old_ctrl, slot_size, + policy.slot_align, has_infoz, /*blocked_element_count=*/0, policy.dealloc); return new_element_target_slot; } @@ -1951,49 +2008,86 @@ } // namespace +RawHashSetLayout RawHashSetLayout::CreateForSmallNonSooGrowth(size_t size, + size_t slot_size, + size_t slot_align, + bool has_infoz) { + ABSL_SWISSTABLE_ASSERT(size <= 1); + static_assert(NumControlBytes(1) == 0); + const size_t num_control_bytes = size * NumControlBytes(3); + static_assert(NextCapacity(0) == 1); + static_assert(NextCapacity(1) == 3); + const size_t new_capacity = NextCapacity(size); + const size_t slots_size = new_capacity * slot_size; + + // Since `GrowthInfoSizeForCapacity` is 0 for capacities 1 and 3, compute + // `metadata_size` directly without `MetadataBeforeControlSize` to avoid its + // `ABSL_PREDICT_FALSE(has_infoz)` branch and allow a branchless shift. + static_assert(GrowthInfoSizeForCapacity(1) == 0); + static_assert(GrowthInfoSizeForCapacity(3) == 0); + const size_t metadata_size = + (has_infoz ? sizeof(HashtablezInfoHandle) + sizeof(uint64_t) : 0) + + NumGenerationBytes(); + ABSL_SWISSTABLE_ASSERT(metadata_size == + MetadataBeforeControlSize(has_infoz, new_capacity)); + + RawHashSetLayout layout(uninitialized_tag_t{}); + layout.slot_offset_ = + AlignUpTo(metadata_size + num_control_bytes, slot_align); + layout.control_offset_ = layout.slot_offset_ - num_control_bytes; + layout.generation_offset_ = layout.control_offset_ - NumGenerationBytes(); + layout.alloc_size_ = layout.slot_offset_ + slots_size; + return layout; +} + void* PrepareInsertSmallNonSoo(CommonFields& common, const PolicyFunctions& __restrict policy, void* hasher, HashForKeyFn get_hash) { ABSL_SWISSTABLE_ASSERT(common.is_small()); ABSL_SWISSTABLE_ASSERT(!policy.soo_enabled); - if (common.capacity() == 1) { - if (common.empty()) { - IncrementSmallSizeNonSoo(common, policy); - // Compute before `RecordInsertMissCold` to avoid reloading `control_`. - void* res = common.single_non_soo_slot(); - // Call NOINLINE function to move infoz instructions out of line. - if (common.has_infoz()) { - ReportInsertMissToInfozAndComputeHash(common, /*probe_length=*/0, - hasher, get_hash); - } - return res; - } else { - return Grow1To3AndPrepareInsert(common, policy, hasher, get_hash); + const bool is_empty = common.empty(); + if (common.capacity() == 1 && is_empty) { + IncrementSmallSizeNonSoo(common, policy); + // Compute before `RecordInsertMissCold` to avoid reloading `control_`. + void* res = common.single_non_soo_slot(); + // Call NOINLINE function to move infoz instructions out of line. + if (common.has_infoz()) { + ReportInsertMissToInfozAndComputeHash(common, /*probe_length=*/0, hasher, + get_hash); } + return res; } - // Growing from 0 to 1 capacity. - ABSL_SWISSTABLE_ASSERT(common.capacity() == 0); - constexpr size_t kNewCapacity = 1; - - common.set_capacity(kNewCapacity); + const size_t size = common.size(); const bool has_infoz = - policy.is_hashtablez_enabled && ShouldSampleNextTable(); + is_empty ? (policy.is_hashtablez_enabled && ShouldSampleNextTable()) + : common.has_infoz(); void* alloc = policy.get_char_alloc(common); + void* new_slots = + AllocBackingArray( + common, policy, + RawHashSetLayout::CreateForSmallNonSooGrowth( + size, policy.slot_size, policy.slot_align, has_infoz), + alloc) + .slots; + if (is_empty) { + // Growing from 0 to 1 capacity. + ABSL_SWISSTABLE_ASSERT(common.capacity() == 0); + constexpr size_t kNewCapacity = 1; + common.set_capacity(kNewCapacity); + common.set_control(static_cast<ctrl_t*>(new_slots)); - const auto [new_ctrl, new_slots] = - AllocBackingArray(common, policy, kNewCapacity, has_infoz, alloc, - /*blocked_element_count=*/0); - common.set_control(new_ctrl); + static_assert(NextCapacity(0) == 1); + PrepareInsertCommon(common); - static_assert(NextCapacity(0) == 1); - PrepareInsertCommon(common); - - if (ABSL_PREDICT_FALSE(has_infoz)) { - ReportSingleElementTableGrowthToInfozAndForceSample(common, policy, hasher, - get_hash); + if (ABSL_PREDICT_FALSE(has_infoz)) { + ReportSingleElementTableGrowthToInfozAndForceSample( + common, policy, hasher, get_hash); + } + return new_slots; } - return new_slots; + return Grow1To3AndPrepareInsert(common, policy, hasher, get_hash, alloc, + has_infoz, new_slots); } namespace { @@ -2264,10 +2358,16 @@ } ABSL_SWISSTABLE_ASSERT(common.size() == policy.soo_capacity()); static constexpr size_t kNewCapacity = NextCapacity(SooCapacity()); - const size_t slot_size = policy.slot_size; - void* alloc = policy.get_char_alloc(common); - common.set_capacity(kNewCapacity); + // Generate the seed and compute hashes before `AllocBackingArray` so + // `hasher`, `get_hash`, and `new_ctrl`/`new_slots` do not need to be spilled + // across calls. + common.generate_new_seed(/*has_infoz=*/false); + const size_t new_hash = get_hash(hasher, common.seed().seed()); + const void* hash_fn = policy.hash_fn(common); + const h2_t soo_slot_h2 = + H2(policy.hash_slot(hash_fn, common.soo_data(), common.seed().seed())); + void* alloc = policy.get_char_alloc(common); // Since the table is not empty, it will not be sampled. // The decision to sample was already made during the first insertion. // @@ -2276,21 +2376,21 @@ const auto [new_ctrl, new_slots] = AllocBackingArray( common, policy, kNewCapacity, /*has_infoz=*/false, alloc, /*blocked_element_count=*/0); - + // Group `common` updates together after `AllocBackingArray` so `inline_data_` + // is modified in a single load-store sequence. + common.set_capacity(kNewCapacity); + common.InitGrowthLeftNoDeleted(kNewCapacity - 2, kNewCapacity); PrepareInsertCommon(common); ABSL_SWISSTABLE_ASSERT(common.size() == 2); - common.InitGrowthLeftNoDeleted(kNewCapacity - 2, kNewCapacity); - common.generate_new_seed(/*has_infoz=*/false); - const h2_t soo_slot_h2 = H2(policy.hash_slot( - policy.hash_fn(common), common.soo_data(), common.seed().seed())); - const size_t new_hash = get_hash(hasher, common.seed().seed()); const size_t offset = Resize1To3NewOffset(new_hash); InitializeThreeElementsControlBytes(soo_slot_h2, H2(new_hash), offset, new_ctrl); + const size_t slot_size = policy.slot_size; SanitizerPoisonMemoryRegion(new_slots, slot_size * kNewCapacity); void* target_slot = SlotAddress(new_slots, SooSlotIndex(), slot_size); + void* new_slot = SlotAddress(new_slots, offset, slot_size); SanitizerUnpoisonMemoryRegion(target_slot, slot_size); if constexpr (TransferUsesMemcpy) { // Target slot is placed at index 1, but capacity is at @@ -2314,7 +2414,6 @@ // Full SOO table couldn't be sampled. If SOO table is sampled, it would // have been resized to the next capacity. ABSL_SWISSTABLE_ASSERT(!common.infoz().IsSampled()); - void* new_slot = SlotAddress(new_slots, offset, slot_size); SanitizerUnpoisonMemoryRegion(new_slot, slot_size); return new_slot; } @@ -2634,9 +2733,8 @@ const DtorPolicy& policy); // Extern template instantiation for inline functions. -template size_t TryFindNewIndexWithoutProbing(size_t h1, size_t old_index, +template size_t TryFindNewIndexWithoutProbing(size_t hash, size_t old_index, size_t old_capacity, - ctrl_t* new_ctrl, size_t new_capacity); template void AssertIsFull(const ctrl_t* const& ctrl, const void* slot, GenerationType generation,
diff --git a/absl/container/internal/raw_hash_set.h b/absl/container/internal/raw_hash_set.h index ad79652..2d9e654 100644 --- a/absl/container/internal/raw_hash_set.h +++ b/absl/container/internal/raw_hash_set.h
@@ -1273,6 +1273,13 @@ alloc_size_ = slot_offset_ + (capacity - blocked_element_count) * slot_size; } + // Creates a layout for a small non-SOO table growing from capacity + // 0 -> 1 (`size == 0`) or 1 -> 3 (`size == 1`). + static RawHashSetLayout CreateForSmallNonSooGrowth(size_t size, + size_t slot_size, + size_t slot_align, + bool has_infoz); + // Returns precomputed offset from the start of the backing allocation of // control. size_t control_offset() const { return control_offset_; } @@ -1290,6 +1297,8 @@ size_t alloc_size() const { return alloc_size_; } private: + explicit RawHashSetLayout(uninitialized_tag_t) {} + size_t control_offset_; size_t generation_offset_; size_t slot_offset_; @@ -1952,54 +1961,28 @@ constexpr size_t kProbedElementIndexSentinel = ~size_t{}; // Implementation detail of transfer_unprobed_elements_to_next_capacity_fn. -// Tries to find the new index for an element whose hash corresponds to -// `h1` for growth to the next capacity. -// Returns kProbedElementIndexSentinel if full probing is required. +// Tries to find the new index for an element with `hash` for growth to the next +// capacity. +// Returns kProbedElementIndexSentinel if element was probed in the table before +// growth. // // If element is located in the first probing group in the table before growth, // returns one of two positions: `old_index` or `old_index + old_capacity + 1`. -// -// Otherwise, we will try to insert it into the first probe group of the new -// table. We only attempt to do so if the first probe group is already -// initialized. template <typename = void> -inline size_t TryFindNewIndexWithoutProbing(size_t h1, size_t old_index, +inline size_t TryFindNewIndexWithoutProbing(size_t hash, size_t old_index, size_t old_capacity, - ctrl_t* new_ctrl, size_t new_capacity) { - size_t index_diff = old_index - h1; + size_t h1 = H1(hash); // The first probe group starts with h1 & capacity. // All following groups start at (h1 + Group::kWidth * K) & capacity. - // We can find an index within the floating group as index_diff modulo - // Group::kWidth. - // Both old and new capacity are larger than Group::kWidth so we can avoid - // computing `& capacity`. - size_t in_floating_group_index = index_diff & (Group::kWidth - 1); - // By subtracting we will get the difference between the first probe group - // and the probe group corresponding to old_index. - index_diff -= in_floating_group_index; - if (ABSL_PREDICT_TRUE((index_diff & old_capacity) == 0)) { - size_t new_index = (h1 + in_floating_group_index) & new_capacity; + // If old_index is in the first probe group, (old_index - h1) & old_capacity + // is the index within the floating group (< Group::kWidth). + size_t index_diff = (old_index - h1) & old_capacity; + if (ABSL_PREDICT_TRUE(index_diff < Group::kWidth)) { + size_t new_index = (h1 + index_diff) & new_capacity; ABSL_ASSUME(new_index != kProbedElementIndexSentinel); return new_index; } - ABSL_SWISSTABLE_ASSERT(((old_index - h1) & old_capacity) >= Group::kWidth); - // Try to insert element into the first probe group. - // new_ctrl is not yet fully initialized so we can't use regular search via - // find_first_non_full. - - // We can search in the first probe group only if it is located in already - // initialized part of the table. - if (ABSL_PREDICT_FALSE((h1 & old_capacity) >= old_index)) { - return kProbedElementIndexSentinel; - } - size_t offset = h1 & new_capacity; - Group new_g(new_ctrl + offset); - if (auto mask = new_g.MaskNonFull(); ABSL_PREDICT_TRUE(mask)) { - size_t result = offset + mask.LowestBitSet(); - ABSL_ASSUME(result != kProbedElementIndexSentinel); - return result; - } return kProbedElementIndexSentinel; } @@ -2090,11 +2073,11 @@ // Different encoding is used depending on the capacity of the table. // See ProbedItem*Bytes classes for details. void (*transfer_unprobed_elements_to_next_capacity)( - CommonFields& common, const ctrl_t* old_ctrl, void* old_slots, + CommonFields& common, const ctrl_t* old_ctrl, // TODO(b/382423690): Try to use absl::FunctionRef here. - void* probed_storage, - void (*encode_probed_element)(void* probed_storage, h2_t h2, - size_t source_offset, size_t h1)); + void* probed_storage, void* old_slots, size_t old_capacity, + void (*encode_probed_element)(void* probed_storage, size_t hash, + size_t source_offset)); uint8_t soo_capacity() const { return static_cast<uint8_t>(soo_enabled ? SooCapacity() : 0); @@ -4249,15 +4232,15 @@ // by GetKey + Hash for memcpyable types. // TODO(b/382423690): Try to type erase for big slots: sizeof(slot_type) > 16. static void transfer_unprobed_elements_to_next_capacity_fn( - CommonFields& common, const ctrl_t* old_ctrl, void* old_slots, - void* probed_storage, - 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); + CommonFields& common, const ctrl_t* old_ctrl, void* probed_storage, + void* old_slots, size_t old_capacity, + void (*encode_probed_element)(void* probed_storage, size_t hash, + size_t source_offset)) { ABSL_ASSUME(old_capacity + 1 >= Group::kWidth); ABSL_ASSUME((old_capacity + 1) % Group::kWidth == 0); + const size_t new_capacity = NextCapacity(old_capacity); + ABSL_ASSUME(new_capacity > kMaxSmallCapacity); + ABSL_SWISSTABLE_ASSERT(common.capacity() == new_capacity); auto* set = reinterpret_cast<raw_hash_set*>(&common); slot_type* old_slots_ptr = to_slot(old_slots); @@ -4279,19 +4262,17 @@ // TODO(b/382423690): try to avoid entire hash calculation since we need // only one new bit of h1. size_t hash = set->hash_of(old_slot); - size_t h1 = H1(hash); - h2_t h2 = H2(hash); size_t new_index = TryFindNewIndexWithoutProbing( - h1, old_index, old_capacity, new_ctrl, new_capacity); + hash, old_index, old_capacity, new_capacity); // Note that encode_probed_element is allowed to use old_ctrl buffer // till and included the old_index. if (ABSL_PREDICT_FALSE(new_index == kProbedElementIndexSentinel)) { - encode_probed_element(probed_storage, h2, old_index, h1); + encode_probed_element(probed_storage, hash, old_index); continue; } ABSL_SWISSTABLE_ASSERT((new_index & old_capacity) <= old_index); ABSL_SWISSTABLE_ASSERT(IsEmpty(new_ctrl[new_index])); - new_ctrl[new_index] = static_cast<ctrl_t>(h2); + new_ctrl[new_index] = static_cast<ctrl_t>(H2(hash)); auto* new_slot = new_slots + new_index; SanitizerUnpoisonMemoryRegion(new_slot, sizeof(slot_type)); set->transfer(new_slot, old_slot); @@ -4541,10 +4522,9 @@ // Extern template for inline function keeps possibility of inlining. // When compiler decided to not inline, no symbols will be added to the // corresponding translation unit. -extern template size_t TryFindNewIndexWithoutProbing(size_t h1, +extern template size_t TryFindNewIndexWithoutProbing(size_t hash, size_t old_index, size_t old_capacity, - ctrl_t* new_ctrl, size_t new_capacity); extern template void AssertIsFull(const ctrl_t* const& ctrl, const void* slot, GenerationType generation,
diff --git a/absl/container/internal/raw_hash_set_test.cc b/absl/container/internal/raw_hash_set_test.cc index 429d7cf..ed25c05 100644 --- a/absl/container/internal/raw_hash_set_test.cc +++ b/absl/container/internal/raw_hash_set_test.cc
@@ -147,6 +147,26 @@ } } +TEST(RawHashSetLayout, CreateForSmallNonSooGrowth) { + for (size_t size : {0u, 1u}) { + for (size_t slot_align : {1u, 2u, 4u, 8u, 16u}) { + for (size_t slot_size : {slot_align, slot_align * 2, slot_align * 3}) { + for (bool has_infoz : {false, true}) { + RawHashSetLayout expected(NextCapacity(size), slot_size, slot_align, + has_infoz, /*blocked_element_count=*/0); + RawHashSetLayout actual = + RawHashSetLayout::CreateForSmallNonSooGrowth( + size, slot_size, slot_align, has_infoz); + EXPECT_EQ(actual.control_offset(), expected.control_offset()); + EXPECT_EQ(actual.generation_offset(), expected.generation_offset()); + EXPECT_EQ(actual.slot_offset(), expected.slot_offset()); + EXPECT_EQ(actual.alloc_size(), expected.alloc_size()); + } + } + } + } +} + void VerifyMiddleSizeTableLayout(size_t capacity, size_t slot_size, size_t slot_align, bool has_infoz, size_t blocked_element_count,