Move CountingAllocator into test_allocator.h and add some other allocators that can be shared between different container tests.

PiperOrigin-RevId: 565693736
Change-Id: I59af987e30da03a805ce59ff0fb7eeae3fc08293

7 files changed, 126 insertions, 63 deletions

diff --git a/absl/container/BUILD.bazel b/absl/container/BUILD.bazel
index 51318379..5be58b1c 100644
--- a/absl/container/BUILD.bazel
+++ b/absl/container/BUILD.bazel
@@ -73,8 +73,8 @@ cc_test(
     copts = ABSL_TEST_COPTS,
     linkopts = ABSL_DEFAULT_LINKOPTS,
     deps = [
-        ":counting_allocator",
         ":fixed_array",
+        ":test_allocator",
         "//absl/base:config",
         "//absl/base:exception_testing",
         "//absl/hash:hash_testing",
@@ -139,9 +139,9 @@ cc_library(
 )
 
 cc_library(
-    name = "counting_allocator",
+    name = "test_allocator",
     testonly = 1,
-    hdrs = ["internal/counting_allocator.h"],
+    hdrs = ["internal/test_allocator.h"],
     copts = ABSL_DEFAULT_COPTS,
     linkopts = ABSL_DEFAULT_LINKOPTS,
     visibility = ["//visibility:private"],
@@ -154,8 +154,8 @@ cc_test(
     copts = ABSL_TEST_COPTS,
     linkopts = ABSL_DEFAULT_LINKOPTS,
     deps = [
-        ":counting_allocator",
         ":inlined_vector",
+        ":test_allocator",
         ":test_instance_tracker",
         "//absl/base:config",
         "//absl/base:core_headers",
@@ -656,6 +656,7 @@ cc_test(
         ":hash_policy_testing",
         ":hashtable_debug",
         ":raw_hash_set",
+        ":test_allocator",
         "//absl/base",
         "//absl/base:config",
         "//absl/base:core_headers",
@@ -993,7 +994,7 @@ cc_test(
     deps = [
         ":btree",
         ":btree_test_common",
-        ":counting_allocator",
+        ":test_allocator",
         ":test_instance_tracker",
         "//absl/algorithm:container",
         "//absl/base:core_headers",
diff --git a/absl/container/CMakeLists.txt b/absl/container/CMakeLists.txt
index f5ef0ddb..bfe0634b 100644
--- a/absl/container/CMakeLists.txt
+++ b/absl/container/CMakeLists.txt
@@ -77,13 +77,13 @@ absl_cc_test(
     absl::btree_test_common
     absl::compare
     absl::core_headers
-    absl::counting_allocator
     absl::flags
     absl::hash_testing
     absl::optional
     absl::random_random
     absl::raw_logging_internal
     absl::strings
+    absl::test_allocator
     absl::test_instance_tracker
     GTest::gmock_main
 )
@@ -145,11 +145,11 @@ absl_cc_test(
     ${ABSL_TEST_COPTS}
   DEPS
     absl::fixed_array
-    absl::counting_allocator
     absl::config
     absl::exception_testing
     absl::hash_testing
     absl::memory
+    absl::test_allocator
     GTest::gmock_main
 )
 
@@ -204,9 +204,9 @@ absl_cc_library(
 # Internal-only target, do not depend on directly.
 absl_cc_library(
   NAME
-    counting_allocator
+    test_allocator
   HDRS
-    "internal/counting_allocator.h"
+    "internal/test_allocator.h"
   COPTS
     ${ABSL_DEFAULT_COPTS}
   DEPS
@@ -224,12 +224,12 @@ absl_cc_test(
     absl::check
     absl::config
     absl::core_headers
-    absl::counting_allocator
     absl::exception_testing
     absl::hash_testing
     absl::inlined_vector
     absl::memory
     absl::strings
+    absl::test_allocator
     absl::test_instance_tracker
     GTest::gmock_main
 )
@@ -747,6 +747,7 @@ absl_cc_test(
     absl::prefetch
     absl::raw_hash_set
     absl::strings
+    absl::test_allocator
     GTest::gmock_main
 )
 
diff --git a/absl/container/btree_test.cc b/absl/container/btree_test.cc
index a204c6d5..3e239c50 100644
--- a/absl/container/btree_test.cc
+++ b/absl/container/btree_test.cc
@@ -37,7 +37,7 @@
 #include "absl/base/macros.h"
 #include "absl/container/btree_map.h"
 #include "absl/container/btree_set.h"
-#include "absl/container/internal/counting_allocator.h"
+#include "absl/container/internal/test_allocator.h"
 #include "absl/container/internal/test_instance_tracker.h"
 #include "absl/flags/flag.h"
 #include "absl/hash/hash_testing.h"
@@ -667,25 +667,10 @@ void BtreeMultiTest() {
   DoTest("identical: ", &container, identical_values);
 }
 
-template <typename T>
-struct PropagatingCountingAlloc : public CountingAllocator<T> {
-  using propagate_on_container_copy_assignment = std::true_type;
-  using propagate_on_container_move_assignment = std::true_type;
-  using propagate_on_container_swap = std::true_type;
-
-  using Base = CountingAllocator<T>;
-  using Base::Base;
-
-  template <typename U>
-  explicit PropagatingCountingAlloc(const PropagatingCountingAlloc<U> &other)
-      : Base(other.bytes_used_) {}
-
-  template <typename U>
-  struct rebind {
-    using other = PropagatingCountingAlloc<U>;
-  };
-};
-
+// TODO(ezb): get rid of BtreeAllocatorTest and replace with test cases using
+// specific propagating allocs (e.g. CopyAssignPropagatingCountingAlloc) and
+// also a test for MinimumAlignmentAlloc. Motivation is better test coverage and
+// faster compilation time.
 template <typename T>
 void BtreeAllocatorTest() {
   using value_type = typename T::value_type;
diff --git a/absl/container/fixed_array_test.cc b/absl/container/fixed_array_test.cc
index 9dbf2a84..2421b5fc 100644
--- a/absl/container/fixed_array_test.cc
+++ b/absl/container/fixed_array_test.cc
@@ -30,7 +30,7 @@
 #include "absl/base/config.h"
 #include "absl/base/internal/exception_testing.h"
 #include "absl/base/options.h"
-#include "absl/container/internal/counting_allocator.h"
+#include "absl/container/internal/test_allocator.h"
 #include "absl/hash/hash_testing.h"
 #include "absl/memory/memory.h"
 
diff --git a/absl/container/inlined_vector_test.cc b/absl/container/inlined_vector_test.cc
index b9a79f5b..241389ae 100644
--- a/absl/container/inlined_vector_test.cc
+++ b/absl/container/inlined_vector_test.cc
@@ -33,7 +33,7 @@
 #include "absl/base/internal/exception_testing.h"
 #include "absl/base/macros.h"
 #include "absl/base/options.h"
-#include "absl/container/internal/counting_allocator.h"
+#include "absl/container/internal/test_allocator.h"
 #include "absl/container/internal/test_instance_tracker.h"
 #include "absl/hash/hash_testing.h"
 #include "absl/log/check.h"
diff --git a/absl/container/internal/raw_hash_set_test.cc b/absl/container/internal/raw_hash_set_test.cc
index 55c6f62e..7588120a 100644
--- a/absl/container/internal/raw_hash_set_test.cc
+++ b/absl/container/internal/raw_hash_set_test.cc
@@ -48,6 +48,7 @@
 #include "absl/container/internal/hash_function_defaults.h"
 #include "absl/container/internal/hash_policy_testing.h"
 #include "absl/container/internal/hashtable_debug.h"
+#include "absl/container/internal/test_allocator.h"
 #include "absl/log/log.h"
 #include "absl/strings/string_view.h"
 
@@ -441,34 +442,6 @@ struct CustomAllocIntTable
   using Base::Base;
 };
 
-// Tries to allocate memory at the minimum alignment even when the default
-// allocator uses a higher alignment.
-template <typename T>
-struct MinimumAlignmentAlloc : std::allocator<T> {
-  MinimumAlignmentAlloc() = default;
-
-  template <typename U>
-  explicit MinimumAlignmentAlloc(const MinimumAlignmentAlloc<U>& /*other*/) {}
-
-  template <class U>
-  struct rebind {
-    using other = MinimumAlignmentAlloc<U>;
-  };
-
-  T* allocate(size_t n) {
-    T* ptr = std::allocator<T>::allocate(n + 1);
-    char* cptr = reinterpret_cast<char*>(ptr);
-    cptr += alignof(T);
-    return reinterpret_cast<T*>(cptr);
-  }
-
-  void deallocate(T* ptr, size_t n) {
-    char* cptr = reinterpret_cast<char*>(ptr);
-    cptr -= alignof(T);
-    std::allocator<T>::deallocate(reinterpret_cast<T*>(cptr), n + 1);
-  }
-};
-
 struct MinimumAlignmentUint8Table
     : raw_hash_set<Uint8Policy, container_internal::hash_default_hash<uint8_t>,
                    std::equal_to<uint8_t>, MinimumAlignmentAlloc<uint8_t>> {
diff --git a/absl/container/internal/counting_allocator.h b/absl/container/internal/test_allocator.h
index 66068a5a..153da427 100644
--- a/absl/container/internal/counting_allocator.h
+++ b/absl/container/internal/test_allocator.h
@@ -12,11 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-#ifndef ABSL_CONTAINER_INTERNAL_COUNTING_ALLOCATOR_H_
-#define ABSL_CONTAINER_INTERNAL_COUNTING_ALLOCATOR_H_
+#ifndef ABSL_CONTAINER_INTERNAL_TEST_ALLOCATOR_H_
+#define ABSL_CONTAINER_INTERNAL_TEST_ALLOCATOR_H_
 
+#include <cstddef>
 #include <cstdint>
 #include <memory>
+#include <type_traits>
 
 #include "absl/base/config.h"
 
@@ -115,8 +117,109 @@ class CountingAllocator {
   int64_t* instance_count_ = nullptr;
 };
 
+template <typename T>
+struct CopyAssignPropagatingCountingAlloc : public CountingAllocator<T> {
+  using propagate_on_container_copy_assignment = std::true_type;
+
+  using Base = CountingAllocator<T>;
+  using Base::Base;
+
+  template <typename U>
+  explicit CopyAssignPropagatingCountingAlloc(
+      const CopyAssignPropagatingCountingAlloc<U>& other)
+      : Base(other.bytes_used_, other.instance_count_) {}
+
+  template <typename U>
+  struct rebind {
+    using other = CopyAssignPropagatingCountingAlloc<U>;
+  };
+};
+
+template <typename T>
+struct MoveAssignPropagatingCountingAlloc : public CountingAllocator<T> {
+  using propagate_on_container_move_assignment = std::true_type;
+
+  using Base = CountingAllocator<T>;
+  using Base::Base;
+
+  template <typename U>
+  explicit MoveAssignPropagatingCountingAlloc(
+      const MoveAssignPropagatingCountingAlloc<U>& other)
+      : Base(other.bytes_used_, other.instance_count_) {}
+
+  template <typename U>
+  struct rebind {
+    using other = MoveAssignPropagatingCountingAlloc<U>;
+  };
+};
+
+template <typename T>
+struct SwapPropagatingCountingAlloc : public CountingAllocator<T> {
+  using propagate_on_container_swap = std::true_type;
+
+  using Base = CountingAllocator<T>;
+  using Base::Base;
+
+  template <typename U>
+  explicit SwapPropagatingCountingAlloc(
+      const SwapPropagatingCountingAlloc<U>& other)
+      : Base(other.bytes_used_, other.instance_count_) {}
+
+  template <typename U>
+  struct rebind {
+    using other = SwapPropagatingCountingAlloc<U>;
+  };
+};
+
+template <typename T>
+struct PropagatingCountingAlloc : public CountingAllocator<T> {
+  using propagate_on_container_copy_assignment = std::true_type;
+  using propagate_on_container_move_assignment = std::true_type;
+  using propagate_on_container_swap = std::true_type;
+
+  using Base = CountingAllocator<T>;
+  using Base::Base;
+
+  template <typename U>
+  explicit PropagatingCountingAlloc(const PropagatingCountingAlloc<U> &other)
+      : Base(other.bytes_used_, other.instance_count_) {}
+
+  template <typename U>
+  struct rebind {
+    using other = PropagatingCountingAlloc<U>;
+  };
+};
+
+// Tries to allocate memory at the minimum alignment even when the default
+// allocator uses a higher alignment.
+template <typename T>
+struct MinimumAlignmentAlloc : std::allocator<T> {
+  MinimumAlignmentAlloc() = default;
+
+  template <typename U>
+  explicit MinimumAlignmentAlloc(const MinimumAlignmentAlloc<U>& /*other*/) {}
+
+  template <class U>
+  struct rebind {
+    using other = MinimumAlignmentAlloc<U>;
+  };
+
+  T* allocate(size_t n) {
+    T* ptr = std::allocator<T>::allocate(n + 1);
+    char* cptr = reinterpret_cast<char*>(ptr);
+    cptr += alignof(T);
+    return reinterpret_cast<T*>(cptr);
+  }
+
+  void deallocate(T* ptr, size_t n) {
+    char* cptr = reinterpret_cast<char*>(ptr);
+    cptr -= alignof(T);
+    std::allocator<T>::deallocate(reinterpret_cast<T*>(cptr), n + 1);
+  }
+};
+
 }  // namespace container_internal
 ABSL_NAMESPACE_END
 }  // namespace absl
 
-#endif  // ABSL_CONTAINER_INTERNAL_COUNTING_ALLOCATOR_H_
+#endif  // ABSL_CONTAINER_INTERNAL_TEST_ALLOCATOR_H_