1 // Copyright 2018 The Abseil Authors.
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //      https://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 #ifndef ABSL_CONTAINER_INTERNAL_COUNTING_ALLOCATOR_H_
16 #define ABSL_CONTAINER_INTERNAL_COUNTING_ALLOCATOR_H_
17 
18 #include <cstdint>
19 #include <memory>
20 
21 #include "absl/base/config.h"
22 
23 namespace absl {
24 ABSL_NAMESPACE_BEGIN
25 namespace container_internal {
26 
27 // This is a stateful allocator, but the state lives outside of the
28 // allocator (in whatever test is using the allocator). This is odd
29 // but helps in tests where the allocator is propagated into nested
30 // containers - that chain of allocators uses the same state and is
31 // thus easier to query for aggregate allocation information.
32 template <typename T>
33 class CountingAllocator {
34  public:
35   using Allocator = std::allocator<T>;
36   using AllocatorTraits = std::allocator_traits<Allocator>;
37   using value_type = typename AllocatorTraits::value_type;
38   using pointer = typename AllocatorTraits::pointer;
39   using const_pointer = typename AllocatorTraits::const_pointer;
40   using size_type = typename AllocatorTraits::size_type;
41   using difference_type = typename AllocatorTraits::difference_type;
42 
43   CountingAllocator() = default;
CountingAllocator(int64_t * bytes_used)44   explicit CountingAllocator(int64_t* bytes_used) : bytes_used_(bytes_used) {}
CountingAllocator(int64_t * bytes_used,int64_t * instance_count)45   CountingAllocator(int64_t* bytes_used, int64_t* instance_count)
46       : bytes_used_(bytes_used), instance_count_(instance_count) {}
47 
48   template <typename U>
CountingAllocator(const CountingAllocator<U> & x)49   CountingAllocator(const CountingAllocator<U>& x)
50       : bytes_used_(x.bytes_used_), instance_count_(x.instance_count_) {}
51 
52   pointer allocate(
53       size_type n,
54       typename AllocatorTraits::const_void_pointer hint = nullptr) {
55     Allocator allocator;
56     pointer ptr = AllocatorTraits::allocate(allocator, n, hint);
57     if (bytes_used_ != nullptr) {
58       *bytes_used_ += n * sizeof(T);
59     }
60     return ptr;
61   }
62 
deallocate(pointer p,size_type n)63   void deallocate(pointer p, size_type n) {
64     Allocator allocator;
65     AllocatorTraits::deallocate(allocator, p, n);
66     if (bytes_used_ != nullptr) {
67       *bytes_used_ -= n * sizeof(T);
68     }
69   }
70 
71   template <typename U, typename... Args>
construct(U * p,Args &&...args)72   void construct(U* p, Args&&... args) {
73     Allocator allocator;
74     AllocatorTraits::construct(allocator, p, std::forward<Args>(args)...);
75     if (instance_count_ != nullptr) {
76       *instance_count_ += 1;
77     }
78   }
79 
80   template <typename U>
destroy(U * p)81   void destroy(U* p) {
82     Allocator allocator;
83     // Ignore GCC warning bug.
84 #if ABSL_INTERNAL_HAVE_MIN_GNUC_VERSION(12, 0)
85 #pragma GCC diagnostic push
86 #pragma GCC diagnostic ignored "-Wuse-after-free"
87 #endif
88     AllocatorTraits::destroy(allocator, p);
89 #if ABSL_INTERNAL_HAVE_MIN_GNUC_VERSION(12, 0)
90 #pragma GCC diagnostic pop
91 #endif
92     if (instance_count_ != nullptr) {
93       *instance_count_ -= 1;
94     }
95   }
96 
97   template <typename U>
98   class rebind {
99    public:
100     using other = CountingAllocator<U>;
101   };
102 
103   friend bool operator==(const CountingAllocator& a,
104                          const CountingAllocator& b) {
105     return a.bytes_used_ == b.bytes_used_ &&
106            a.instance_count_ == b.instance_count_;
107   }
108 
109   friend bool operator!=(const CountingAllocator& a,
110                          const CountingAllocator& b) {
111     return !(a == b);
112   }
113 
114   int64_t* bytes_used_ = nullptr;
115   int64_t* instance_count_ = nullptr;
116 };
117 
118 }  // namespace container_internal
119 ABSL_NAMESPACE_END
120 }  // namespace absl
121 
122 #endif  // ABSL_CONTAINER_INTERNAL_COUNTING_ALLOCATOR_H_
123