1 /*
2  * Copyright (c) Facebook, Inc. and its affiliates.
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *     http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 #pragma once
18 
19 #include <map>
20 #include <memory>
21 #include <scoped_allocator>
22 #include <set>
23 #include <string>
24 #include <vector>
25 
26 #include <folly/sorted_vector_types.h>
27 
28 template <class T>
29 struct AlwaysThrowAllocator : private std::allocator<T> {
30   using value_type = T;
31 
32   AlwaysThrowAllocator() = default;
33   AlwaysThrowAllocator(const AlwaysThrowAllocator&) noexcept = default;
34   AlwaysThrowAllocator& operator=(const AlwaysThrowAllocator&) noexcept =
35       default;
36   template <class U>
AlwaysThrowAllocatorAlwaysThrowAllocator37   explicit AlwaysThrowAllocator(const AlwaysThrowAllocator<U>&) noexcept {}
38   ~AlwaysThrowAllocator() = default;
39 
allocateAlwaysThrowAllocator40   T* allocate(size_t) { throw std::bad_alloc(); }
41 
deallocateAlwaysThrowAllocator42   void deallocate(T*, size_t) {}
43 
44   template <class U>
45   friend bool operator==(
46       AlwaysThrowAllocator<T> const&, AlwaysThrowAllocator<U> const&) noexcept {
47     return true;
48   }
49 
50   template <class U>
51   friend bool operator!=(
52       AlwaysThrowAllocator<T> const&, AlwaysThrowAllocator<U> const&) noexcept {
53     return false;
54   }
55 };
56 
57 using ScopedAlwaysThrowAlloc =
58     std::scoped_allocator_adaptor<AlwaysThrowAllocator<char>>;
59 
60 template <class T>
61 using AlwaysThrowVector = std::vector<T, ScopedAlwaysThrowAlloc>;
62 
63 template <class T>
64 using AlwaysThrowSet = std::set<T, std::less<T>, ScopedAlwaysThrowAlloc>;
65 
66 template <class K, class V>
67 using AlwaysThrowMap = std::map<K, V, std::less<K>, ScopedAlwaysThrowAlloc>;
68 
69 using AlwaysThrowString =
70     std::basic_string<char, std::char_traits<char>, ScopedAlwaysThrowAlloc>;
71 
72 template <class T>
73 struct StatefulAlloc : private std::allocator<T> {
74   using value_type = T;
75 
76   StatefulAlloc() = default;
77   StatefulAlloc(const StatefulAlloc&) = default;
78   StatefulAlloc& operator=(const StatefulAlloc&) noexcept = default;
StatefulAllocStatefulAlloc79   explicit StatefulAlloc(int state) : state_(state) {}
80   template <class U>
StatefulAllocStatefulAlloc81   explicit StatefulAlloc(const StatefulAlloc<U>& other) noexcept
82       : state_(other.state_) {}
83 
84   using propagate_on_container_copy_assignment = std::true_type;
85   using propagate_on_container_move_assignment = std::true_type;
86   using propagate_on_container_swap = std::true_type;
87 
88   int state_ = 0;
89 
allocateStatefulAlloc90   T* allocate(size_t size) { return std::allocator<T>::allocate(size); }
91 
deallocateStatefulAlloc92   void deallocate(T* p, size_t size) { std::allocator<T>::deallocate(p, size); }
93 
94   template <class U>
95   friend bool operator==(
96       StatefulAlloc<T> const& a, StatefulAlloc<U> const& b) noexcept {
97     return a.state_ == b.state_;
98   }
99 
100   template <class U>
101   friend bool operator!=(
102       StatefulAlloc<T> const& a, StatefulAlloc<U> const& b) noexcept {
103     return a.state_ != b.state_;
104   }
105 };
106 
107 using ScopedStatefulAlloc = std::scoped_allocator_adaptor<StatefulAlloc<char>>;
108 
109 template <class T>
110 using StatefulAllocVector = std::vector<T, ScopedStatefulAlloc>;
111 
112 template <class T>
113 using StatefulAllocSet = std::set<T, std::less<T>, ScopedStatefulAlloc>;
114 
115 template <class K, class V>
116 using StatefulAllocMap = std::map<K, V, std::less<K>, ScopedStatefulAlloc>;
117 
118 template <class T>
119 using StatefulAllocSortedVectorSet =
120     folly::sorted_vector_set<T, std::less<T>, ScopedStatefulAlloc>;
121 
122 template <class K, class V>
123 using StatefulAllocSortedVectorMap =
124     folly::sorted_vector_map<K, V, std::less<K>, ScopedStatefulAlloc>;
125 
126 template <class T>
127 struct CountingAlloc : private std::allocator<T> {
128   using value_type = T;
129 
CountingAllocCountingAlloc130   CountingAlloc() : counter_(std::make_shared<int>(0)) {}
131 
132   CountingAlloc(const CountingAlloc&) = default;
133   CountingAlloc& operator=(const CountingAlloc&) noexcept = default;
134   template <class U>
CountingAllocCountingAlloc135   explicit CountingAlloc(const CountingAlloc<U>& other) noexcept
136       : counter_(other.counter_) {}
137 
138   std::shared_ptr<int> counter_;
139 
getCountCountingAlloc140   int getCount() const { return *counter_; }
141 
allocateCountingAlloc142   T* allocate(size_t size) {
143     (*counter_)++;
144     return std::allocator<T>::allocate(size);
145   }
146 
deallocateCountingAlloc147   void deallocate(T* p, size_t size) { std::allocator<T>::deallocate(p, size); }
148 
149   template <class U>
150   friend bool operator==(
151       CountingAlloc<T> const&, CountingAlloc<U> const&) noexcept {
152     return true;
153   }
154 
155   template <class U>
156   friend bool operator!=(
157       CountingAlloc<T> const&, CountingAlloc<U> const&) noexcept {
158     return false;
159   }
160 };
161 
162 using ScopedCountingAlloc = std::scoped_allocator_adaptor<CountingAlloc<char>>;
163 
164 template <class T>
165 using CountingVector = std::vector<T, ScopedCountingAlloc>;
166 
167 template <class T>
168 using CountingSet = std::set<T, std::less<T>, ScopedCountingAlloc>;
169 
170 template <class K, class V>
171 using CountingMap = std::map<K, V, std::less<K>, ScopedCountingAlloc>;
172 
173 using CountingString =
174     std::basic_string<char, std::char_traits<char>, ScopedCountingAlloc>;
175