1 /*
2  * Copyright 2020 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 #pragma once
18 
19 #include <ftl/initializer_list.h>
20 #include <ftl/optional.h>
21 #include <ftl/small_vector.h>
22 
23 #include <algorithm>
24 #include <functional>
25 #include <type_traits>
26 #include <utility>
27 
28 namespace android::ftl {
29 
30 // Associative container with unique, unordered keys. Unlike std::unordered_map, key-value pairs are
31 // stored in contiguous storage for cache efficiency. The map is allocated statically until its size
32 // exceeds N, at which point mappings are relocated to dynamic memory. The try_emplace operation has
33 // a non-standard analogue try_replace that destructively emplaces. The API also defines an in-place
34 // counterpart to insert_or_assign: emplace_or_replace. Lookup is done not via a subscript operator,
35 // but immutable getters that can optionally transform the value.
36 //
37 // SmallMap<K, V, 0> unconditionally allocates on the heap.
38 //
39 // Example usage:
40 //
41 //   ftl::SmallMap<int, std::string, 3> map;
42 //   assert(map.empty());
43 //   assert(!map.dynamic());
44 //
45 //   map = ftl::init::map<int, std::string>(123, "abc")(-1)(42, 3u, '?');
46 //   assert(map.size() == 3u);
47 //   assert(!map.dynamic());
48 //
49 //   assert(map.contains(123));
50 //   assert(map.get(42).transform([](const std::string& s) { return s.size(); }) == 3u);
51 //
52 //   const auto opt = map.get(-1);
53 //   assert(opt);
54 //
55 //   std::string& ref = *opt;
56 //   assert(ref.empty());
57 //   ref = "xyz";
58 //
59 //   map.emplace_or_replace(0, "vanilla", 2u, 3u);
60 //   assert(map.dynamic());
61 //
62 //   assert(map == SmallMap(ftl::init::map(-1, "xyz"sv)(0, "nil"sv)(42, "???"sv)(123, "abc"sv)));
63 //
64 template <typename K, typename V, std::size_t N, typename KeyEqual = std::equal_to<K>>
65 class SmallMap final {
66   using Map = SmallVector<std::pair<const K, V>, N>;
67 
68   template <typename, typename, std::size_t, typename>
69   friend class SmallMap;
70 
71  public:
72   using key_type = K;
73   using mapped_type = V;
74 
75   using value_type = typename Map::value_type;
76   using size_type = typename Map::size_type;
77   using difference_type = typename Map::difference_type;
78 
79   using reference = typename Map::reference;
80   using iterator = typename Map::iterator;
81 
82   using const_reference = typename Map::const_reference;
83   using const_iterator = typename Map::const_iterator;
84 
85   // Creates an empty map.
86   SmallMap() = default;
87 
88   // Constructs at most N key-value pairs in place by forwarding per-pair constructor arguments.
89   // The template arguments K, V, and N are inferred using the deduction guide defined below.
90   // The syntax for listing pairs is as follows:
91   //
92   //   ftl::SmallMap map = ftl::init::map<int, std::string>(123, "abc")(-1)(42, 3u, '?');
93   //   static_assert(std::is_same_v<decltype(map), ftl::SmallMap<int, std::string, 3>>);
94   //
95   // The types of the key and value are deduced if the first pair contains exactly two arguments:
96   //
97   //   ftl::SmallMap map = ftl::init::map(0, 'a')(1, 'b')(2, 'c');
98   //   static_assert(std::is_same_v<decltype(map), ftl::SmallMap<int, char, 3>>);
99   //
100   template <typename U, std::size_t... Sizes, typename... Types>
SmallMap(InitializerList<U,std::index_sequence<Sizes...>,Types...> && list)101   SmallMap(InitializerList<U, std::index_sequence<Sizes...>, Types...>&& list)
102       : map_(std::move(list)) {
103     deduplicate();
104   }
105 
106   // Copies or moves key-value pairs from a convertible map.
107   template <typename Q, typename W, std::size_t M, typename E>
SmallMap(SmallMap<Q,W,M,E> other)108   SmallMap(SmallMap<Q, W, M, E> other) : map_(std::move(other.map_)) {}
109 
static_capacity()110   static constexpr size_type static_capacity() { return N; }
111 
max_size()112   size_type max_size() const { return map_.max_size(); }
size()113   size_type size() const { return map_.size(); }
empty()114   bool empty() const { return map_.empty(); }
115 
116   // Returns whether the map is backed by static or dynamic storage.
dynamic()117   bool dynamic() const {
118     if constexpr (static_capacity() > 0) {
119       return map_.dynamic();
120     } else {
121       return true;
122     }
123   }
124 
begin()125   iterator begin() { return map_.begin(); }
begin()126   const_iterator begin() const { return cbegin(); }
cbegin()127   const_iterator cbegin() const { return map_.cbegin(); }
128 
end()129   iterator end() { return map_.end(); }
end()130   const_iterator end() const { return cend(); }
cend()131   const_iterator cend() const { return map_.cend(); }
132 
133   // Returns whether a mapping exists for the given key.
contains(const key_type & key)134   bool contains(const key_type& key) const { return get(key).has_value(); }
135 
136   // Returns a reference to the value for the given key, or std::nullopt if the key was not found.
137   //
138   //   ftl::SmallMap map = ftl::init::map('a', 'A')('b', 'B')('c', 'C');
139   //
140   //   const auto opt = map.get('c');
141   //   assert(opt == 'C');
142   //
143   //   char d = 'd';
144   //   const auto ref = map.get('d').value_or(std::ref(d));
145   //   ref.get() = 'D';
146   //   assert(d == 'D');
147   //
148   auto get(const key_type& key) const -> Optional<std::reference_wrapper<const mapped_type>> {
149     for (const auto& [k, v] : *this) {
150       if (KeyEqual{}(k, key)) {
151         return std::cref(v);
152       }
153     }
154     return {};
155   }
156 
157   auto get(const key_type& key) -> Optional<std::reference_wrapper<mapped_type>> {
158     for (auto& [k, v] : *this) {
159       if (KeyEqual{}(k, key)) {
160         return std::ref(v);
161       }
162     }
163     return {};
164   }
165 
166   // Returns an iterator to an existing mapping for the given key, or the end() iterator otherwise.
find(const key_type & key)167   const_iterator find(const key_type& key) const { return const_cast<SmallMap&>(*this).find(key); }
find(const key_type & key)168   iterator find(const key_type& key) { return find(key, begin()); }
169 
170   // Inserts a mapping unless it exists. Returns an iterator to the inserted or existing mapping,
171   // and whether the mapping was inserted.
172   //
173   // On emplace, if the map reaches its static or dynamic capacity, then all iterators are
174   // invalidated. Otherwise, only the end() iterator is invalidated.
175   //
176   template <typename... Args>
try_emplace(const key_type & key,Args &&...args)177   std::pair<iterator, bool> try_emplace(const key_type& key, Args&&... args) {
178     if (const auto it = find(key); it != end()) {
179       return {it, false};
180     }
181 
182     decltype(auto) ref_or_it =
183         map_.emplace_back(std::piecewise_construct, std::forward_as_tuple(key),
184                           std::forward_as_tuple(std::forward<Args>(args)...));
185 
186     if constexpr (static_capacity() > 0) {
187       return {&ref_or_it, true};
188     } else {
189       return {ref_or_it, true};
190     }
191   }
192 
193   // Replaces a mapping if it exists, and returns an iterator to it. Returns the end() iterator
194   // otherwise.
195   //
196   // The value is replaced via move constructor, so type V does not need to define copy/move
197   // assignment, e.g. its data members may be const.
198   //
199   // The arguments may directly or indirectly refer to the mapping being replaced.
200   //
201   // Iterators to the replaced mapping point to its replacement, and others remain valid.
202   //
203   template <typename... Args>
try_replace(const key_type & key,Args &&...args)204   iterator try_replace(const key_type& key, Args&&... args) {
205     const auto it = find(key);
206     if (it == end()) return it;
207     map_.replace(it, std::piecewise_construct, std::forward_as_tuple(key),
208                  std::forward_as_tuple(std::forward<Args>(args)...));
209     return it;
210   }
211 
212   // In-place counterpart of std::unordered_map's insert_or_assign. Returns true on emplace, or
213   // false on replace.
214   //
215   // The value is emplaced and replaced via move constructor, so type V does not need to define
216   // copy/move assignment, e.g. its data members may be const.
217   //
218   // On emplace, if the map reaches its static or dynamic capacity, then all iterators are
219   // invalidated. Otherwise, only the end() iterator is invalidated. On replace, iterators
220   // to the replaced mapping point to its replacement, and others remain valid.
221   //
222   template <typename... Args>
emplace_or_replace(const key_type & key,Args &&...args)223   std::pair<iterator, bool> emplace_or_replace(const key_type& key, Args&&... args) {
224     const auto [it, ok] = try_emplace(key, std::forward<Args>(args)...);
225     if (ok) return {it, ok};
226     map_.replace(it, std::piecewise_construct, std::forward_as_tuple(key),
227                  std::forward_as_tuple(std::forward<Args>(args)...));
228     return {it, ok};
229   }
230 
231   // Removes a mapping if it exists, and returns whether it did.
232   //
233   // The last() and end() iterators, as well as those to the erased mapping, are invalidated.
234   //
erase(const key_type & key)235   bool erase(const key_type& key) { return erase(key, begin()); }
236 
237   // Removes all mappings.
238   //
239   // All iterators are invalidated.
240   //
clear()241   void clear() { map_.clear(); }
242 
243  private:
find(const key_type & key,iterator first)244   iterator find(const key_type& key, iterator first) {
245     return std::find_if(first, end(),
246                         [&key](const auto& pair) { return KeyEqual{}(pair.first, key); });
247   }
248 
erase(const key_type & key,iterator first)249   bool erase(const key_type& key, iterator first) {
250     const auto it = find(key, first);
251     if (it == end()) return false;
252     map_.unstable_erase(it);
253     return true;
254   }
255 
deduplicate()256   void deduplicate() {
257     for (auto it = begin(); it != end();) {
258       if (const auto key = it->first; ++it != end()) {
259         while (erase(key, it));
260       }
261     }
262   }
263 
264   Map map_;
265 };
266 
267 // Deduction guide for in-place constructor.
268 template <typename K, typename V, typename E, std::size_t... Sizes, typename... Types>
269 SmallMap(InitializerList<KeyValue<K, V, E>, std::index_sequence<Sizes...>, Types...>&&)
270     -> SmallMap<K, V, sizeof...(Sizes), E>;
271 
272 // Returns whether the key-value pairs of two maps are equal.
273 template <typename K, typename V, std::size_t N, typename Q, typename W, std::size_t M, typename E>
274 bool operator==(const SmallMap<K, V, N, E>& lhs, const SmallMap<Q, W, M, E>& rhs) {
275   if (lhs.size() != rhs.size()) return false;
276 
277   for (const auto& [k, v] : lhs) {
278     const auto& lv = v;
279     if (!rhs.get(k).transform([&lv](const W& rv) { return lv == rv; }).value_or(false)) {
280       return false;
281     }
282   }
283 
284   return true;
285 }
286 
287 // TODO: Remove in C++20.
288 template <typename K, typename V, std::size_t N, typename Q, typename W, std::size_t M, typename E>
289 inline bool operator!=(const SmallMap<K, V, N, E>& lhs, const SmallMap<Q, W, M, E>& rhs) {
290   return !(lhs == rhs);
291 }
292 
293 }  // namespace android::ftl
294