counting_allocator.h 3.6 KB

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