Small table growth optimization.

Details:
- In case the table entirely fits into a single group size (`capacity <= Group::kWidth`), the order of elements is not important.
- For growing upto Group::kWidth we rotate control bytes and slots deterministically (using memcpy).
- We also avoid second find_first_non_full right after resize for small growing.
PiperOrigin-RevId: 588825966
Change-Id: I09bd7fd489e3868dcf56c36b436805d08dae7ab5

absl/container/internal/raw_hash_set.cc

@@ -17,11 +17,13 @@
 #include <atomic>
 #include <atomic>
 #include <cassert>
 #include <cassert>
 #include <cstddef>
 #include <cstddef>
+#include <cstdint>
 #include <cstring>
 #include <cstring>
 
 
 #include "absl/base/attributes.h"
 #include "absl/base/attributes.h"
 #include "absl/base/config.h"
 #include "absl/base/config.h"
 #include "absl/base/dynamic_annotations.h"
 #include "absl/base/dynamic_annotations.h"
+#include "absl/container/internal/container_memory.h"
 #include "absl/hash/hash.h"
 #include "absl/hash/hash.h"
 
 
 namespace absl {
 namespace absl {
@@ -126,14 +128,6 @@ FindInfo find_first_non_full_outofline(const CommonFields& common,
   return find_first_non_full(common, hash);
   return find_first_non_full(common, hash);
 }
 }
 
 
-// Returns the address of the ith slot in slots where each slot occupies
-// slot_size.
-static inline void* SlotAddress(void* slot_array, size_t slot,
-                                size_t slot_size) {
-  return reinterpret_cast<void*>(reinterpret_cast<uintptr_t>(slot_array) +
-                                 (slot * slot_size));
-}
-
 // Returns the address of the slot just after slot assuming each slot has the
 // Returns the address of the slot just after slot assuming each slot has the
 // specified size.
 // specified size.
 static inline void* NextSlot(void* slot, size_t slot_size) {
 static inline void* NextSlot(void* slot, size_t slot_size) {
@@ -254,6 +248,7 @@ void ClearBackingArray(CommonFields& c, const PolicyFunctions& policy,
   c.set_size(0);
   c.set_size(0);
   if (reuse) {
   if (reuse) {
     ResetCtrl(c, policy.slot_size);
     ResetCtrl(c, policy.slot_size);
+    ResetGrowthLeft(c);
     c.infoz().RecordStorageChanged(0, c.capacity());
     c.infoz().RecordStorageChanged(0, c.capacity());
   } else {
   } else {
     // We need to record infoz before calling dealloc, which will unregister
     // We need to record infoz before calling dealloc, which will unregister
@@ -268,6 +263,109 @@ void ClearBackingArray(CommonFields& c, const PolicyFunctions& policy,
   }
   }
 }
 }
 
 
+void HashSetResizeHelper::GrowIntoSingleGroupShuffleControlBytes(
+    ctrl_t* new_ctrl, size_t new_capacity) const {
+  assert(is_single_group(new_capacity));
+  constexpr size_t kHalfWidth = Group::kWidth / 2;
+  assert(old_capacity_ < kHalfWidth);
+
+  const size_t half_old_capacity = old_capacity_ / 2;
+
+  // NOTE: operations are done with compile time known size = kHalfWidth.
+  // Compiler optimizes that into single ASM operation.
+
+  // Copy second half of bytes to the beginning.
+  // We potentially copy more bytes in order to have compile time known size.
+  // Mirrored bytes from the old_ctrl_ will also be copied.
+  // In case of old_capacity_ == 3, we will copy 1st element twice.
+  // Examples:
+  // old_ctrl = 0S0EEEEEEE...
+  // new_ctrl = S0EEEEEEEE...
+  //
+  // old_ctrl = 01S01EEEEE...
+  // new_ctrl = 1S01EEEEEE...
+  //
+  // old_ctrl = 0123456S0123456EE...
+  // new_ctrl = 456S0123?????????...
+  std::memcpy(new_ctrl, old_ctrl_ + half_old_capacity + 1, kHalfWidth);
+  // Clean up copied kSentinel from old_ctrl.
+  new_ctrl[half_old_capacity] = ctrl_t::kEmpty;
+
+  // Clean up damaged or uninitialized bytes.
+
+  // Clean bytes after the intended size of the copy.
+  // Example:
+  // new_ctrl = 1E01EEEEEEE????
+  // *new_ctrl= 1E0EEEEEEEE????
+  // position      /
+  std::memset(new_ctrl + old_capacity_ + 1, static_cast<int8_t>(ctrl_t::kEmpty),
+              kHalfWidth);
+  // Clean non-mirrored bytes that are not initialized.
+  // For small old_capacity that may be inside of mirrored bytes zone.
+  // Examples:
+  // new_ctrl = 1E0EEEEEEEE??????????....
+  // *new_ctrl= 1E0EEEEEEEEEEEEE?????....
+  // position           /
+  //
+  // new_ctrl = 456E0123???????????...
+  // *new_ctrl= 456E0123EEEEEEEE???...
+  // position           /
+  std::memset(new_ctrl + kHalfWidth, static_cast<int8_t>(ctrl_t::kEmpty),
+              kHalfWidth);
+  // Clean last mirrored bytes that are not initialized
+  // and will not be overwritten by mirroring.
+  // Examples:
+  // new_ctrl = 1E0EEEEEEEEEEEEE????????
+  // *new_ctrl= 1E0EEEEEEEEEEEEEEEEEEEEE
+  // position           S       /
+  //
+  // new_ctrl = 456E0123EEEEEEEE???????????????
+  // *new_ctrl= 456E0123EEEEEEEE???????EEEEEEEE
+  // position                  S       /
+  std::memset(new_ctrl + new_capacity + kHalfWidth,
+              static_cast<int8_t>(ctrl_t::kEmpty), kHalfWidth);
+
+  // Create mirrored bytes. old_capacity_ < kHalfWidth
+  // Example:
+  // new_ctrl = 456E0123EEEEEEEE???????EEEEEEEE
+  // *new_ctrl= 456E0123EEEEEEEE456E0123EEEEEEE
+  // position                  S/
+  ctrl_t g[kHalfWidth];
+  std::memcpy(g, new_ctrl, kHalfWidth);
+  std::memcpy(new_ctrl + new_capacity + 1, g, kHalfWidth);
+
+  // Finally set sentinel to its place.
+  new_ctrl[new_capacity] = ctrl_t::kSentinel;
+}
+
+void HashSetResizeHelper::GrowIntoSingleGroupShuffleTransferableSlots(
+    void* old_slots, void* new_slots, size_t slot_size) const {
+  assert(old_capacity_ > 0);
+  const size_t half_old_capacity = old_capacity_ / 2;
+
+  SanitizerUnpoisonMemoryRegion(old_slots, slot_size * old_capacity_);
+  std::memcpy(new_slots,
+              SlotAddress(old_slots, half_old_capacity + 1, slot_size),
+              slot_size * half_old_capacity);
+  std::memcpy(SlotAddress(new_slots, half_old_capacity + 1, slot_size),
+              old_slots, slot_size * (half_old_capacity + 1));
+}
+
+void HashSetResizeHelper::GrowSizeIntoSingleGroupTransferable(
+    CommonFields& c, void* old_slots, size_t slot_size) {
+  assert(old_capacity_ < Group::kWidth / 2);
+  assert(is_single_group(c.capacity()));
+  assert(IsGrowingIntoSingleGroupApplicable(old_capacity_, c.capacity()));
+
+  GrowIntoSingleGroupShuffleControlBytes(c.control(), c.capacity());
+  GrowIntoSingleGroupShuffleTransferableSlots(old_slots, c.slot_array(),
+                                              slot_size);
+
+  // We poison since GrowIntoSingleGroupShuffleTransferableSlots
+  // may leave empty slots unpoisoned.
+  PoisonSingleGroupEmptySlots(c, slot_size);
+}
+
 }  // namespace container_internal
 }  // namespace container_internal
 ABSL_NAMESPACE_END
 ABSL_NAMESPACE_END
 }  // namespace absl
 }  // namespace absl

absl/container/internal/raw_hash_set_test.cc

@@ -30,6 +30,7 @@
 #include <ostream>
 #include <ostream>
 #include <random>
 #include <random>
 #include <string>
 #include <string>
+#include <tuple>
 #include <type_traits>
 #include <type_traits>
 #include <unordered_map>
 #include <unordered_map>
 #include <unordered_set>
 #include <unordered_set>
@@ -299,7 +300,7 @@ TEST(Group, CountLeadingEmptyOrDeleted) {
   }
   }
 }
 }
 
 
-template <class T>
+template <class T, bool kTransferable = false>
 struct ValuePolicy {
 struct ValuePolicy {
   using slot_type = T;
   using slot_type = T;
   using key_type = T;
   using key_type = T;
@@ -317,10 +318,11 @@ struct ValuePolicy {
   }
   }
 
 
   template <class Allocator>
   template <class Allocator>
-  static void transfer(Allocator* alloc, slot_type* new_slot,
-                       slot_type* old_slot) {
+  static std::integral_constant<bool, kTransferable> transfer(
+      Allocator* alloc, slot_type* new_slot, slot_type* old_slot) {
     construct(alloc, new_slot, std::move(*old_slot));
     construct(alloc, new_slot, std::move(*old_slot));
     destroy(alloc, old_slot);
     destroy(alloc, old_slot);
+    return {};
   }
   }
 
 
   static T& element(slot_type* slot) { return *slot; }
   static T& element(slot_type* slot) { return *slot; }
@@ -337,6 +339,8 @@ struct ValuePolicy {
 using IntPolicy = ValuePolicy<int64_t>;
 using IntPolicy = ValuePolicy<int64_t>;
 using Uint8Policy = ValuePolicy<uint8_t>;
 using Uint8Policy = ValuePolicy<uint8_t>;
 
 
+using TranferableIntPolicy = ValuePolicy<int64_t, /*kTransferable=*/true>;
+
 class StringPolicy {
 class StringPolicy {
   template <class F, class K, class V,
   template <class F, class K, class V,
             class = typename std::enable_if<
             class = typename std::enable_if<
@@ -409,9 +413,10 @@ struct StringTable
   using Base::Base;
   using Base::Base;
 };
 };
 
 
-template <typename T>
-struct ValueTable : raw_hash_set<ValuePolicy<T>, hash_default_hash<T>,
-                                 std::equal_to<T>, std::allocator<T>> {
+template <typename T, bool kTransferable = false>
+struct ValueTable
+    : raw_hash_set<ValuePolicy<T, kTransferable>, hash_default_hash<T>,
+                   std::equal_to<T>, std::allocator<T>> {
   using Base = typename ValueTable::raw_hash_set;
   using Base = typename ValueTable::raw_hash_set;
   using Base::Base;
   using Base::Base;
 };
 };
@@ -419,6 +424,8 @@ struct ValueTable : raw_hash_set<ValuePolicy<T>, hash_default_hash<T>,
 using IntTable = ValueTable<int64_t>;
 using IntTable = ValueTable<int64_t>;
 using Uint8Table = ValueTable<uint8_t>;
 using Uint8Table = ValueTable<uint8_t>;
 
 
+using TransferableIntTable = ValueTable<int64_t, /*kTransferable=*/true>;
+
 template <typename T>
 template <typename T>
 struct CustomAlloc : std::allocator<T> {
 struct CustomAlloc : std::allocator<T> {
   CustomAlloc() = default;
   CustomAlloc() = default;
@@ -653,6 +660,68 @@ TEST(Table, InsertWithinCapacity) {
   EXPECT_THAT(addr(0), original_addr_0);
   EXPECT_THAT(addr(0), original_addr_0);
 }
 }
 
 
+template <class TableType>
+class SmallTableResizeTest : public testing::Test {};
+
+TYPED_TEST_SUITE_P(SmallTableResizeTest);
+
+TYPED_TEST_P(SmallTableResizeTest, InsertIntoSmallTable) {
+  TypeParam t;
+  for (int i = 0; i < 32; ++i) {
+    t.insert(i);
+    ASSERT_EQ(t.size(), i + 1);
+    for (int j = 0; j < i + 1; ++j) {
+      EXPECT_TRUE(t.find(j) != t.end());
+      EXPECT_EQ(*t.find(j), j);
+    }
+  }
+}
+
+TYPED_TEST_P(SmallTableResizeTest, ResizeGrowSmallTables) {
+  TypeParam t;
+  for (size_t source_size = 0; source_size < 32; ++source_size) {
+    for (size_t target_size = source_size; target_size < 32; ++target_size) {
+      for (bool rehash : {false, true}) {
+        for (size_t i = 0; i < source_size; ++i) {
+          t.insert(static_cast<int>(i));
+        }
+        if (rehash) {
+          t.rehash(target_size);
+        } else {
+          t.reserve(target_size);
+        }
+        for (size_t i = 0; i < source_size; ++i) {
+          EXPECT_TRUE(t.find(static_cast<int>(i)) != t.end());
+          EXPECT_EQ(*t.find(static_cast<int>(i)), static_cast<int>(i));
+        }
+      }
+    }
+  }
+}
+
+TYPED_TEST_P(SmallTableResizeTest, ResizeReduceSmallTables) {
+  TypeParam t;
+  for (size_t source_size = 0; source_size < 32; ++source_size) {
+    for (size_t target_size = 0; target_size <= source_size; ++target_size) {
+      size_t inserted_count = std::min<size_t>(source_size, 5);
+      for (size_t i = 0; i < inserted_count; ++i) {
+        t.insert(static_cast<int>(i));
+      }
+      t.rehash(target_size);
+      for (size_t i = 0; i < inserted_count; ++i) {
+        EXPECT_TRUE(t.find(static_cast<int>(i)) != t.end());
+        EXPECT_EQ(*t.find(static_cast<int>(i)), static_cast<int>(i));
+      }
+    }
+  }
+}
+
+REGISTER_TYPED_TEST_SUITE_P(SmallTableResizeTest, InsertIntoSmallTable,
+                            ResizeGrowSmallTables, ResizeReduceSmallTables);
+using SmallTableTypes = ::testing::Types<IntTable, TransferableIntTable>;
+INSTANTIATE_TYPED_TEST_SUITE_P(InstanceSmallTableResizeTest,
+                               SmallTableResizeTest, SmallTableTypes);
+
 TEST(Table, LazyEmplace) {
 TEST(Table, LazyEmplace) {
   StringTable t;
   StringTable t;
   bool called = false;
   bool called = false;
@@ -1071,7 +1140,7 @@ TEST(Table, Erase) {
 TEST(Table, EraseMaintainsValidIterator) {
 TEST(Table, EraseMaintainsValidIterator) {
   IntTable t;
   IntTable t;
   const int kNumElements = 100;
   const int kNumElements = 100;
-  for (int i = 0; i < kNumElements; i ++) {
+  for (int i = 0; i < kNumElements; i++) {
     EXPECT_TRUE(t.emplace(i).second);
     EXPECT_TRUE(t.emplace(i).second);
   }
   }
   EXPECT_EQ(t.size(), kNumElements);
   EXPECT_EQ(t.size(), kNumElements);
@@ -2258,21 +2327,34 @@ TEST(RawHashSamplerTest, DoNotSampleCustomAllocators) {
 }
 }
 
 
 #ifdef ABSL_HAVE_ADDRESS_SANITIZER
 #ifdef ABSL_HAVE_ADDRESS_SANITIZER
-TEST(Sanitizer, PoisoningUnused) {
-  IntTable t;
-  t.reserve(5);
-  // Insert something to force an allocation.
-  int64_t& v1 = *t.insert(0).first;
+template <class TableType>
+class SanitizerTest : public testing::Test {};
 
 
-  // Make sure there is something to test.
-  ASSERT_GT(t.capacity(), 1);
+TYPED_TEST_SUITE_P(SanitizerTest);
 
 
-  int64_t* slots = RawHashSetTestOnlyAccess::GetSlots(t);
-  for (size_t i = 0; i < t.capacity(); ++i) {
-    EXPECT_EQ(slots + i != &v1, __asan_address_is_poisoned(slots + i));
+TYPED_TEST_P(SanitizerTest, PoisoningUnused) {
+  TypeParam t;
+  for (size_t reserve_size = 2; reserve_size < 1024;
+       reserve_size = reserve_size * 3 / 2) {
+    t.reserve(reserve_size);
+    // Insert something to force an allocation.
+    int64_t& v = *t.insert(0).first;
+
+    // Make sure there is something to test.
+    ASSERT_GT(t.capacity(), 1);
+
+    int64_t* slots = RawHashSetTestOnlyAccess::GetSlots(t);
+    for (size_t i = 0; i < t.capacity(); ++i) {
+      EXPECT_EQ(slots + i != &v, __asan_address_is_poisoned(slots + i)) << i;
+    }
   }
   }
 }
 }
 
 
+REGISTER_TYPED_TEST_SUITE_P(SanitizerTest, PoisoningUnused);
+using SanitizerTableTypes = ::testing::Types<IntTable, TransferableIntTable>;
+INSTANTIATE_TYPED_TEST_SUITE_P(InstanceSanitizerTest, SanitizerTest,
+                               SanitizerTableTypes);
+
 TEST(Sanitizer, PoisoningOnErase) {
 TEST(Sanitizer, PoisoningOnErase) {
   IntTable t;
   IntTable t;
   int64_t& v = *t.insert(0).first;
   int64_t& v = *t.insert(0).first;