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 <cassert> 19 #include <cstdint> 20 #include <memory> 21 22 #include "absl/base/config.h" 23 24 namespace absl { 25 ABSL_NAMESPACE_BEGIN 26 namespace container_internal { 27 28 // This is a stateful allocator, but the state lives outside of the 29 // allocator (in whatever test is using the allocator). This is odd 30 // but helps in tests where the allocator is propagated into nested 31 // containers - that chain of allocators uses the same state and is 32 // thus easier to query for aggregate allocation information. 33 template <typename T> 34 class CountingAllocator : public std::allocator<T> { 35 public: 36 using Alloc = std::allocator<T>; 37 using pointer = typename Alloc::pointer; 38 using size_type = typename Alloc::size_type; 39 CountingAllocator()40 CountingAllocator() : bytes_used_(nullptr) {} CountingAllocator(int64_t * b)41 explicit CountingAllocator(int64_t* b) : bytes_used_(b) {} 42 43 template <typename U> CountingAllocator(const CountingAllocator<U> & x)44 CountingAllocator(const CountingAllocator<U>& x) 45 : Alloc(x), bytes_used_(x.bytes_used_) {} 46 47 pointer allocate(size_type n, 48 std::allocator<void>::const_pointer hint = nullptr) { 49 assert(bytes_used_ != nullptr); 50 *bytes_used_ += n * sizeof(T); 51 return Alloc::allocate(n, hint); 52 } 53 deallocate(pointer p,size_type n)54 void deallocate(pointer p, size_type n) { 55 Alloc::deallocate(p, n); 56 assert(bytes_used_ != nullptr); 57 *bytes_used_ -= n * sizeof(T); 58 } 59 60 template<typename U> 61 class rebind { 62 public: 63 using other = CountingAllocator<U>; 64 }; 65 66 friend bool operator==(const CountingAllocator& a, 67 const CountingAllocator& b) { 68 return a.bytes_used_ == b.bytes_used_; 69 } 70 71 friend bool operator!=(const CountingAllocator& a, 72 const CountingAllocator& b) { 73 return !(a == b); 74 } 75 76 int64_t* bytes_used_; 77 }; 78 79 } // namespace container_internal 80 ABSL_NAMESPACE_END 81 } // namespace absl 82 83 #endif // ABSL_CONTAINER_INTERNAL_COUNTING_ALLOCATOR_H_ 84