Add tests for btrees in which slot_type is overaligned and slot_type is equal to field_type.

Also do some minor refactoring in btree.h.

PiperOrigin-RevId: 529460378
Change-Id: I278833ada93bbb7652e149fceed08ce3485e4312

2 files changed, 68 insertions, 34 deletions

diff --git a/absl/container/btree_test.cc b/absl/container/btree_test.cc
index 15335c8c..72f446b2 100644
--- a/absl/container/btree_test.cc
+++ b/absl/container/btree_test.cc
@@ -1233,8 +1233,10 @@ class BtreeNodePeer {
   }
 
   template <typename Btree>
-  constexpr static bool UsesGenerations() {
-    return Btree::params_type::kEnableGenerations;
+  constexpr static bool FieldTypeEqualsSlotType() {
+    return std::is_same<
+        typename btree_node<typename Btree::params_type>::field_type,
+        typename btree_node<typename Btree::params_type>::slot_type>::value;
   }
 };
 
@@ -1463,7 +1465,7 @@ class SizedBtreeSet
   using Base = typename SizedBtreeSet::btree_set_container;
 
  public:
-  SizedBtreeSet() {}
+  SizedBtreeSet() = default;
   using Base::Base;
 };
 
@@ -1509,10 +1511,9 @@ TEST(Btree, MovesComparisonsCopiesSwapsTracking) {
   EXPECT_EQ(BtreeNodePeer::GetNumSlotsPerNode<decltype(set61)>(), 61);
   EXPECT_EQ(BtreeNodePeer::GetNumSlotsPerNode<decltype(set100)>(), 100);
   if (sizeof(void *) == 8) {
-    EXPECT_EQ(
-        BtreeNodePeer::GetNumSlotsPerNode<absl::btree_set<int32_t>>(),
-        // When we have generations, there is one fewer slot.
-        BtreeNodePeer::UsesGenerations<absl::btree_set<int32_t>>() ? 60 : 61);
+    EXPECT_EQ(BtreeNodePeer::GetNumSlotsPerNode<absl::btree_set<int32_t>>(),
+              // When we have generations, there is one fewer slot.
+              BtreeGenerationsEnabled() ? 60 : 61);
   }
 
   // Test key insertion/deletion in random order.
@@ -1568,10 +1569,9 @@ TEST(Btree, MovesComparisonsCopiesSwapsTrackingThreeWayCompare) {
   EXPECT_EQ(BtreeNodePeer::GetNumSlotsPerNode<decltype(set61)>(), 61);
   EXPECT_EQ(BtreeNodePeer::GetNumSlotsPerNode<decltype(set100)>(), 100);
   if (sizeof(void *) == 8) {
-    EXPECT_EQ(
-        BtreeNodePeer::GetNumSlotsPerNode<absl::btree_set<int32_t>>(),
-        // When we have generations, there is one fewer slot.
-        BtreeNodePeer::UsesGenerations<absl::btree_set<int32_t>>() ? 60 : 61);
+    EXPECT_EQ(BtreeNodePeer::GetNumSlotsPerNode<absl::btree_set<int32_t>>(),
+              // When we have generations, there is one fewer slot.
+              BtreeGenerationsEnabled() ? 60 : 61);
   }
 
   // Test key insertion/deletion in random order.
@@ -3226,7 +3226,7 @@ TEST(Btree, MutatedKeysCaught) {
 #ifndef _MSC_VER
 // This test crashes on MSVC.
 TEST(Btree, InvalidIteratorUse) {
-  if (!BtreeNodePeer::UsesGenerations<absl::btree_set<int>>())
+  if (!BtreeGenerationsEnabled())
     GTEST_SKIP() << "Generation validation for iterators is disabled.";
 
   // Invalid memory use can trigger heap-use-after-free in ASan or invalidated
@@ -3569,6 +3569,41 @@ TEST(Btree, InvalidPointerUse) {
   EXPECT_DEATH(std::cout << *ptr, "heap-use-after-free");
 }
 
+template<typename Set>
+void TestBasicFunctionality(Set set) {
+  using value_type = typename Set::value_type;
+  for (int i = 0; i < 100; ++i) { set.insert(value_type(i)); }
+  for (int i = 50; i < 100; ++i) { set.erase(value_type(i)); }
+  auto it = set.begin();
+  for (int i = 0; i < 50; ++i, ++it) {
+    ASSERT_EQ(set.find(value_type(i)), it) << i;
+  }
+}
+
+template<size_t align>
+struct alignas(align) OveralignedKey {
+  explicit OveralignedKey(int i) : key(i) {}
+  bool operator<(const OveralignedKey &other) const { return key < other.key; }
+  int key = 0;
+};
+
+TEST(Btree, OveralignedKey) {
+  // Test basic functionality with both even and odd numbers of slots per node.
+  // The goal here is to detect cases where alignment may be incorrect.
+  TestBasicFunctionality(
+      SizedBtreeSet<OveralignedKey<16>, /*TargetValuesPerNode=*/8>());
+  TestBasicFunctionality(
+      SizedBtreeSet<OveralignedKey<16>, /*TargetValuesPerNode=*/9>());
+}
+
+TEST(Btree, FieldTypeEqualsSlotType) {
+  // This breaks if we try to do layout_type::Pointer<slot_type> because
+  // slot_type is the same as field_type.
+  using set_type = absl::btree_set<uint8_t>;
+  static_assert(BtreeNodePeer::FieldTypeEqualsSlotType<set_type>(), "");
+  TestBasicFunctionality(set_type());
+}
+
 }  // namespace
 }  // namespace container_internal
 ABSL_NAMESPACE_END
diff --git a/absl/container/internal/btree.h b/absl/container/internal/btree.h
index 18409a09..569faa07 100644
--- a/absl/container/internal/btree.h
+++ b/absl/container/internal/btree.h
@@ -86,6 +86,12 @@ namespace container_internal {
 #define ABSL_BTREE_ENABLE_GENERATIONS
 #endif
 
+#ifdef ABSL_BTREE_ENABLE_GENERATIONS
+constexpr bool BtreeGenerationsEnabled() { return true; }
+#else
+constexpr bool BtreeGenerationsEnabled() { return false; }
+#endif
+
 template <typename Compare, typename T, typename U>
 using compare_result_t = absl::result_of_t<const Compare(const T &, const U &)>;
 
@@ -378,12 +384,6 @@ struct common_params : common_policy_traits<SlotPolicy> {
       std::is_same<key_compare, StringBtreeDefaultGreater>::value;
   static constexpr bool kIsKeyCompareTransparent =
       IsTransparent<original_key_compare>::value || kIsKeyCompareStringAdapted;
-  static constexpr bool kEnableGenerations =
-#ifdef ABSL_BTREE_ENABLE_GENERATIONS
-      true;
-#else
-      false;
-#endif
 
   // A type which indicates if we have a key-compare-to functor or a plain old
   // key-compare functor.
@@ -589,7 +589,7 @@ class btree_node {
   constexpr static size_type SizeWithNSlots(size_type n) {
     return layout_type(
                /*parent*/ 1,
-               /*generation*/ params_type::kEnableGenerations ? 1 : 0,
+               /*generation*/ BtreeGenerationsEnabled() ? 1 : 0,
                /*position, start, finish, max_count*/ 4,
                /*slots*/ n,
                /*children*/ 0)
@@ -629,23 +629,22 @@ class btree_node {
   // has this value.
   constexpr static field_type kInternalNodeMaxCount = 0;
 
-  // Leaves can have less than kNodeSlots values.
-  constexpr static layout_type LeafLayout(
-      const size_type slot_count = kNodeSlots) {
+  constexpr static layout_type Layout(const size_type slot_count,
+                                      const size_type child_count) {
     return layout_type(
         /*parent*/ 1,
-        /*generation*/ params_type::kEnableGenerations ? 1 : 0,
+        /*generation*/ BtreeGenerationsEnabled() ? 1 : 0,
         /*position, start, finish, max_count*/ 4,
         /*slots*/ slot_count,
-        /*children*/ 0);
+        /*children*/ child_count);
+  }
+  // Leaves can have less than kNodeSlots values.
+  constexpr static layout_type LeafLayout(
+      const size_type slot_count = kNodeSlots) {
+    return Layout(slot_count, 0);
   }
   constexpr static layout_type InternalLayout() {
-    return layout_type(
-        /*parent*/ 1,
-        /*generation*/ params_type::kEnableGenerations ? 1 : 0,
-        /*position, start, finish, max_count*/ 4,
-        /*slots*/ kNodeSlots,
-        /*children*/ kNodeSlots + 1);
+    return Layout(kNodeSlots, kNodeSlots + 1);
   }
   constexpr static size_type LeafSize(const size_type slot_count = kNodeSlots) {
     return LeafLayout(slot_count).AllocSize();
@@ -729,7 +728,7 @@ class btree_node {
 
   // Gets the root node's generation integer, which is the one used by the tree.
   uint32_t *get_root_generation() const {
-    assert(params_type::kEnableGenerations);
+    assert(BtreeGenerationsEnabled());
     const btree_node *curr = this;
     for (; !curr->is_root(); curr = curr->parent()) continue;
     return const_cast<uint32_t *>(&curr->GetField<1>()[0]);
@@ -737,16 +736,16 @@ class btree_node {
 
   // Returns the generation for iterator validation.
   uint32_t generation() const {
-    return params_type::kEnableGenerations ? *get_root_generation() : 0;
+    return BtreeGenerationsEnabled() ? *get_root_generation() : 0;
   }
   // Updates generation. Should only be called on a root node or during node
   // initialization.
   void set_generation(uint32_t generation) {
-    if (params_type::kEnableGenerations) GetField<1>()[0] = generation;
+    if (BtreeGenerationsEnabled()) GetField<1>()[0] = generation;
   }
   // Updates the generation. We do this whenever the node is mutated.
   void next_generation() {
-    if (params_type::kEnableGenerations) ++*get_root_generation();
+    if (BtreeGenerationsEnabled()) ++*get_root_generation();
   }
 
   // Getters for the key/value at position i in the node.