2 files changed, 56 insertions, 19 deletions

diff --git a/absl/container/internal/btree.h b/absl/container/internal/btree.h
index 4504e9ce6659..b23138f09553 100644
--- a/absl/container/internal/btree.h
+++ b/absl/container/internal/btree.h
@@ -929,9 +929,15 @@ struct btree_iterator {
   void decrement_slow();
 
  public:
+  bool operator==(const iterator &other) const {
+    return node == other.node && position == other.position;
+  }
   bool operator==(const const_iterator &other) const {
     return node == other.node && position == other.position;
   }
+  bool operator!=(const iterator &other) const {
+    return node != other.node || position != other.position;
+  }
   bool operator!=(const const_iterator &other) const {
     return node != other.node || position != other.position;
   }
diff --git a/absl/container/internal/counting_allocator.h b/absl/container/internal/counting_allocator.h
index 9efdc6621363..927cf0825579 100644
--- a/absl/container/internal/counting_allocator.h
+++ b/absl/container/internal/counting_allocator.h
@@ -15,7 +15,6 @@
 #ifndef ABSL_CONTAINER_INTERNAL_COUNTING_ALLOCATOR_H_
 #define ABSL_CONTAINER_INTERNAL_COUNTING_ALLOCATOR_H_
 
-#include <cassert>
 #include <cstdint>
 #include <memory>
 
@@ -31,33 +30,63 @@ namespace container_internal {
 // containers - that chain of allocators uses the same state and is
 // thus easier to query for aggregate allocation information.
 template <typename T>
-class CountingAllocator : public std::allocator<T> {
+class CountingAllocator {
  public:
-  using Alloc = std::allocator<T>;
-  using pointer = typename Alloc::pointer;
-  using size_type = typename Alloc::size_type;
+  using Allocator = std::allocator<T>;
+  using AllocatorTraits = std::allocator_traits<Allocator>;
+  using value_type = typename AllocatorTraits::value_type;
+  using pointer = typename AllocatorTraits::pointer;
+  using const_pointer = typename AllocatorTraits::const_pointer;
+  using size_type = typename AllocatorTraits::size_type;
+  using difference_type = typename AllocatorTraits::difference_type;
 
-  CountingAllocator() : bytes_used_(nullptr) {}
-  explicit CountingAllocator(int64_t* b) : bytes_used_(b) {}
+  CountingAllocator() = default;
+  explicit CountingAllocator(int64_t* bytes_used) : bytes_used_(bytes_used) {}
+  CountingAllocator(int64_t* bytes_used, int64_t* instance_count)
+      : bytes_used_(bytes_used), instance_count_(instance_count) {}
 
   template <typename U>
   CountingAllocator(const CountingAllocator<U>& x)
-      : Alloc(x), bytes_used_(x.bytes_used_) {}
+      : bytes_used_(x.bytes_used_), instance_count_(x.instance_count_) {}
 
-  pointer allocate(size_type n,
-                   std::allocator<void>::const_pointer hint = nullptr) {
-    assert(bytes_used_ != nullptr);
-    *bytes_used_ += n * sizeof(T);
-    return Alloc::allocate(n, hint);
+  pointer allocate(
+      size_type n,
+      typename AllocatorTraits::const_void_pointer hint = nullptr) {
+    Allocator allocator;
+    pointer ptr = AllocatorTraits::allocate(allocator, n, hint);
+    if (bytes_used_ != nullptr) {
+      *bytes_used_ += n * sizeof(T);
+    }
+    return ptr;
   }
 
   void deallocate(pointer p, size_type n) {
-    Alloc::deallocate(p, n);
-    assert(bytes_used_ != nullptr);
-    *bytes_used_ -= n * sizeof(T);
+    Allocator allocator;
+    AllocatorTraits::deallocate(allocator, p, n);
+    if (bytes_used_ != nullptr) {
+      *bytes_used_ -= n * sizeof(T);
+    }
   }
 
-  template<typename U>
+  template <typename U, typename... Args>
+  void construct(U* p, Args&&... args) {
+    Allocator allocator;
+    AllocatorTraits::construct(allocator, p, std::forward<Args>(args)...);
+    if (instance_count_ != nullptr) {
+      *instance_count_ += 1;
+    }
+  }
+
+  template <typename U>
+  void destroy(U* p) {
+    Allocator allocator;
+    AllocatorTraits::destroy(allocator, p);
+    if (instance_count_ != nullptr) {
+      *instance_count_ -= 1;
+    }
+  }
+
+  template <typename U>
   class rebind {
    public:
     using other = CountingAllocator<U>;
@@ -65,7 +94,8 @@ class CountingAllocator : public std::allocator<T> {
 
   friend bool operator==(const CountingAllocator& a,
                          const CountingAllocator& b) {
-    return a.bytes_used_ == b.bytes_used_;
+    return a.bytes_used_ == b.bytes_used_ &&
+           a.instance_count_ == b.instance_count_;
   }
 
   friend bool operator!=(const CountingAllocator& a,
@@ -73,7 +103,8 @@ class CountingAllocator : public std::allocator<T> {
     return !(a == b);
   }
 
-  int64_t* bytes_used_;
+  int64_t* bytes_used_ = nullptr;
+  int64_t* instance_count_ = nullptr;
 };
 
 }  // namespace container_internal