1 /*
2  * Copyright (C) 2012 The Android Open Source Project
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 #ifndef ART_LIBARTBASE_BASE_SAFE_MAP_H_
18 #define ART_LIBARTBASE_BASE_SAFE_MAP_H_
19 
20 #include <map>
21 #include <memory>
22 #include <type_traits>
23 
24 #include <android-base/logging.h>
25 
26 namespace art {
27 
28 // Equivalent to std::map, but without operator[] and its bug-prone semantics (in particular,
29 // the implicit insertion of a default-constructed value on failed lookups).
30 template <typename K, typename V, typename Comparator = std::less<K>,
31           typename Allocator = std::allocator<std::pair<const K, V>>>
32 class SafeMap {
33  private:
34   using Self = SafeMap<K, V, Comparator, Allocator>;
35   using Impl = std::map<K, V, Comparator, Allocator>;
36 
37  public:
38   using key_compare        = typename Impl::key_compare;
39   using value_compare      = typename Impl::value_compare;
40   using allocator_type     = typename Impl::allocator_type;
41   using iterator           = typename Impl::iterator;
42   using const_iterator     = typename Impl::const_iterator;
43   using size_type          = typename Impl::size_type;
44   using key_type           = typename Impl::key_type;
45   using value_type         = typename Impl::value_type;
46   using node_type          = typename Impl::node_type;
47   using insert_return_type = typename Impl::insert_return_type;
48 
49   SafeMap() = default;
50   SafeMap(const SafeMap&) = default;
51   SafeMap(SafeMap&&) noexcept = default;
SafeMap(const allocator_type & allocator)52   explicit SafeMap(const allocator_type& allocator) : map_(allocator) {}
53   explicit SafeMap(const key_compare& cmp, const allocator_type& allocator = allocator_type())
map_(cmp,allocator)54       : map_(cmp, allocator) {
55   }
56 
57   Self& operator=(const Self& rhs) {
58     map_ = rhs.map_;
59     return *this;
60   }
61 
get_allocator()62   allocator_type get_allocator() const { return map_.get_allocator(); }
key_comp()63   key_compare key_comp() const { return map_.key_comp(); }
value_comp()64   value_compare value_comp() const { return map_.value_comp(); }
65 
begin()66   iterator begin() { return map_.begin(); }
begin()67   const_iterator begin() const { return map_.begin(); }
end()68   iterator end() { return map_.end(); }
end()69   const_iterator end() const { return map_.end(); }
70 
empty()71   bool empty() const { return map_.empty(); }
size()72   size_type size() const { return map_.size(); }
73 
swap(Self & other)74   void swap(Self& other) { map_.swap(other.map_); }
clear()75   void clear() { map_.clear(); }
76 
erase(const_iterator pos)77   iterator erase(const_iterator pos) { return map_.erase(pos); }
erase(iterator pos)78   iterator erase(iterator pos) { return map_.erase(pos); }
erase(iterator first,iterator last)79   iterator erase(iterator first, iterator last) { return map_.erase(first, last); }
erase(const key_type & k)80   size_type erase(const key_type& k) { return map_.erase(k); }
81 
extract(const_iterator pos)82   node_type extract(const_iterator pos) { return map_.extract(pos); }
extract(const key_type & k)83   node_type extract(const key_type& k) { return map_.extract(k); }
84 
insert(value_type && value)85   std::pair<iterator, bool> insert(value_type&& value) { return map_.insert(std::move(value)); }
insert(node_type && node)86   insert_return_type insert(node_type&& node) { return map_.insert(std::move(node)); }
insert(const_iterator hint,node_type && node)87   insert_return_type insert(const_iterator hint, node_type&& node) {
88     return map_.insert(hint, std::move(node));
89   }
90 
find(const Kv & k)91   template<typename Kv> iterator find(const Kv& k) { return map_.find(k); }
find(const Kv & k)92   template<typename Kv> const_iterator find(const Kv& k) const { return map_.find(k); }
93 
lower_bound(const Kv & k)94   template<typename Kv> iterator lower_bound(const Kv& k) { return map_.lower_bound(k); }
lower_bound(const Kv & k)95   template<typename Kv> const_iterator lower_bound(const Kv& k) const {
96     return map_.lower_bound(k);
97   }
98 
upper_bound(const Kv & k)99   template<typename Kv> iterator upper_bound(const Kv& k) { return map_.upper_bound(k); }
upper_bound(const Kv & k)100   template<typename Kv> const_iterator upper_bound(const Kv& k) const {
101     return map_.upper_bound(k);
102   }
103 
count(const Kv & k)104   template<typename Kv> size_type count(const Kv& k) const { return map_.count(k); }
105 
106   // Note that unlike std::map's operator[], this doesn't return a reference to the value.
Get(const K & k)107   V Get(const K& k) const {
108     const_iterator it = map_.find(k);
109     DCHECK(it != map_.end());
110     return it->second;
111   }
112 
113   // Used to insert a new mapping.
Put(const K & k,const V & v)114   iterator Put(const K& k, const V& v) {
115     std::pair<iterator, bool> result = map_.emplace(k, v);
116     DCHECK(result.second);  // Check we didn't accidentally overwrite an existing value.
117     return result.first;
118   }
Put(const K & k,V && v)119   iterator Put(const K& k, V&& v) {
120     std::pair<iterator, bool> result = map_.emplace(k, std::move(v));
121     DCHECK(result.second);  // Check we didn't accidentally overwrite an existing value.
122     return result.first;
123   }
124 
125   // Used to insert a new mapping at a known position for better performance.
PutBefore(const_iterator pos,const K & k,const V & v)126   iterator PutBefore(const_iterator pos, const K& k, const V& v) {
127     // Check that we're using the correct position and the key is not in the map.
128     DCHECK(pos == map_.end() || map_.key_comp()(k, pos->first));
129     DCHECK(pos == map_.begin() || map_.key_comp()((--const_iterator(pos))->first, k));
130     return map_.emplace_hint(pos, k, v);
131   }
PutBefore(const_iterator pos,const K & k,V && v)132   iterator PutBefore(const_iterator pos, const K& k, V&& v) {
133     // Check that we're using the correct position and the key is not in the map.
134     DCHECK(pos == map_.end() || map_.key_comp()(k, pos->first));
135     DCHECK(pos == map_.begin() || map_.key_comp()((--const_iterator(pos))->first, k));
136     return map_.emplace_hint(pos, k, std::move(v));
137   }
138 
139   // Used to insert a new mapping or overwrite an existing mapping. Note that if the value type
140   // of this container is a pointer, any overwritten pointer will be lost and if this container
141   // was the owner, you have a leak. Returns iterator pointing to the new or overwritten entry.
Overwrite(const K & k,const V & v)142   iterator Overwrite(const K& k, const V& v) {
143     std::pair<iterator, bool> result = map_.insert(std::make_pair(k, v));
144     if (!result.second) {
145       // Already there - update the value for the existing key
146       result.first->second = v;
147     }
148     return result.first;
149   }
150 
151   template <typename CreateFn>
GetOrCreate(const K & k,CreateFn && create)152   V& GetOrCreate(const K& k, CreateFn&& create) {
153     static_assert(std::is_same_v<V, std::invoke_result_t<CreateFn>>,
154                   "Argument `create` should return a value of type V.");
155     auto lb = lower_bound(k);
156     if (lb != end() && !key_comp()(k, lb->first)) {
157       return lb->second;
158     }
159     auto it = PutBefore(lb, k, create());
160     return it->second;
161   }
162 
FindOrAdd(const K & k,const V & v)163   iterator FindOrAdd(const K& k, const V& v) {
164     iterator it = find(k);
165     return it == end() ? Put(k, v) : it;
166   }
167 
FindOrAdd(const K & k)168   iterator FindOrAdd(const K& k) {
169     iterator it = find(k);
170     return it == end() ? Put(k, V()) : it;
171   }
172 
Equals(const Self & rhs)173   bool Equals(const Self& rhs) const {
174     return map_ == rhs.map_;
175   }
176 
177   template <class... Args>
emplace(Args &&...args)178   std::pair<iterator, bool> emplace(Args&&... args) {
179     return map_.emplace(std::forward<Args>(args)...);
180   }
181 
182  private:
183   Impl map_;
184 };
185 
186 template <typename K, typename V, typename Comparator, typename Allocator>
187 bool operator==(const SafeMap<K, V, Comparator, Allocator>& lhs,
188                 const SafeMap<K, V, Comparator, Allocator>& rhs) {
189   return lhs.Equals(rhs);
190 }
191 
192 template <typename K, typename V, typename Comparator, typename Allocator>
193 bool operator!=(const SafeMap<K, V, Comparator, Allocator>& lhs,
194                 const SafeMap<K, V, Comparator, Allocator>& rhs) {
195   return !(lhs == rhs);
196 }
197 
198 }  // namespace art
199 
200 #endif  // ART_LIBARTBASE_BASE_SAFE_MAP_H_
201