]>
Commit | Line | Data |
---|---|---|
29967ef6 XL |
1 | use super::either_iter::EitherIter; |
2 | use crate::fx::FxHashMap; | |
3 | use arrayvec::ArrayVec; | |
4 | use std::fmt; | |
5 | use std::hash::Hash; | |
6 | use std::iter::FromIterator; | |
7 | use std::ops::Index; | |
8 | ||
9 | // For pointer-sized arguments arrays | |
10 | // are faster than set/map for up to 64 | |
11 | // arguments. | |
12 | // | |
13 | // On the other hand such a big array | |
14 | // hurts cache performance, makes passing | |
15 | // sso structures around very expensive. | |
16 | // | |
17 | // Biggest performance benefit is gained | |
18 | // for reasonably small arrays that stay | |
19 | // small in vast majority of cases. | |
20 | // | |
cdc7bbd5 | 21 | // '8' is chosen as a sane default, to be |
29967ef6 | 22 | // reevaluated later. |
29967ef6 XL |
23 | const SSO_ARRAY_SIZE: usize = 8; |
24 | ||
25 | /// Small-storage-optimized implementation of a map. | |
26 | /// | |
27 | /// Stores elements in a small array up to a certain length | |
28 | /// and switches to `HashMap` when that length is exceeded. | |
29 | // | |
30 | // FIXME: Implements subset of HashMap API. | |
31 | // | |
32 | // Missing HashMap API: | |
33 | // all hasher-related | |
c295e0f8 | 34 | // try_reserve |
29967ef6 XL |
35 | // shrink_to (unstable) |
36 | // drain_filter (unstable) | |
37 | // into_keys/into_values (unstable) | |
38 | // all raw_entry-related | |
39 | // PartialEq/Eq (requires sorting the array) | |
fc512014 | 40 | // Entry::or_insert_with_key |
29967ef6 XL |
41 | // Vacant/Occupied entries and related |
42 | // | |
43 | // FIXME: In HashMap most methods accepting key reference | |
44 | // accept reference to generic `Q` where `K: Borrow<Q>`. | |
45 | // | |
46 | // However, using this approach in `HashMap::get` apparently | |
47 | // breaks inlining and noticeably reduces performance. | |
48 | // | |
49 | // Performance *should* be the same given that borrow is | |
50 | // a NOP in most cases, but in practice that's not the case. | |
51 | // | |
52 | // Further investigation is required. | |
53 | // | |
54 | // Affected methods: | |
55 | // SsoHashMap::get | |
56 | // SsoHashMap::get_mut | |
57 | // SsoHashMap::get_entry | |
58 | // SsoHashMap::get_key_value | |
59 | // SsoHashMap::contains_key | |
60 | // SsoHashMap::remove | |
61 | // SsoHashMap::remove_entry | |
62 | // Index::index | |
63 | // SsoHashSet::take | |
64 | // SsoHashSet::get | |
65 | // SsoHashSet::remove | |
66 | // SsoHashSet::contains | |
67 | ||
68 | #[derive(Clone)] | |
69 | pub enum SsoHashMap<K, V> { | |
cdc7bbd5 | 70 | Array(ArrayVec<(K, V), SSO_ARRAY_SIZE>), |
29967ef6 XL |
71 | Map(FxHashMap<K, V>), |
72 | } | |
73 | ||
74 | impl<K, V> SsoHashMap<K, V> { | |
75 | /// Creates an empty `SsoHashMap`. | |
76 | #[inline] | |
77 | pub fn new() -> Self { | |
78 | SsoHashMap::Array(ArrayVec::new()) | |
79 | } | |
80 | ||
81 | /// Creates an empty `SsoHashMap` with the specified capacity. | |
82 | pub fn with_capacity(cap: usize) -> Self { | |
83 | if cap <= SSO_ARRAY_SIZE { | |
84 | Self::new() | |
85 | } else { | |
86 | SsoHashMap::Map(FxHashMap::with_capacity_and_hasher(cap, Default::default())) | |
87 | } | |
88 | } | |
89 | ||
90 | /// Clears the map, removing all key-value pairs. Keeps the allocated memory | |
91 | /// for reuse. | |
92 | pub fn clear(&mut self) { | |
93 | match self { | |
94 | SsoHashMap::Array(array) => array.clear(), | |
95 | SsoHashMap::Map(map) => map.clear(), | |
96 | } | |
97 | } | |
98 | ||
99 | /// Returns the number of elements the map can hold without reallocating. | |
100 | pub fn capacity(&self) -> usize { | |
101 | match self { | |
102 | SsoHashMap::Array(_) => SSO_ARRAY_SIZE, | |
103 | SsoHashMap::Map(map) => map.capacity(), | |
104 | } | |
105 | } | |
106 | ||
107 | /// Returns the number of elements in the map. | |
108 | pub fn len(&self) -> usize { | |
109 | match self { | |
110 | SsoHashMap::Array(array) => array.len(), | |
111 | SsoHashMap::Map(map) => map.len(), | |
112 | } | |
113 | } | |
114 | ||
115 | /// Returns `true` if the map contains no elements. | |
116 | pub fn is_empty(&self) -> bool { | |
117 | match self { | |
118 | SsoHashMap::Array(array) => array.is_empty(), | |
119 | SsoHashMap::Map(map) => map.is_empty(), | |
120 | } | |
121 | } | |
122 | ||
123 | /// An iterator visiting all key-value pairs in arbitrary order. | |
124 | /// The iterator element type is `(&'a K, &'a V)`. | |
125 | #[inline] | |
126 | pub fn iter(&self) -> <&Self as IntoIterator>::IntoIter { | |
127 | self.into_iter() | |
128 | } | |
129 | ||
130 | /// An iterator visiting all key-value pairs in arbitrary order, | |
131 | /// with mutable references to the values. | |
132 | /// The iterator element type is `(&'a K, &'a mut V)`. | |
133 | #[inline] | |
134 | pub fn iter_mut(&mut self) -> impl Iterator<Item = (&'_ K, &'_ mut V)> { | |
135 | self.into_iter() | |
136 | } | |
137 | ||
138 | /// An iterator visiting all keys in arbitrary order. | |
139 | /// The iterator element type is `&'a K`. | |
140 | pub fn keys(&self) -> impl Iterator<Item = &'_ K> { | |
141 | match self { | |
142 | SsoHashMap::Array(array) => EitherIter::Left(array.iter().map(|(k, _v)| k)), | |
143 | SsoHashMap::Map(map) => EitherIter::Right(map.keys()), | |
144 | } | |
145 | } | |
146 | ||
147 | /// An iterator visiting all values in arbitrary order. | |
148 | /// The iterator element type is `&'a V`. | |
149 | pub fn values(&self) -> impl Iterator<Item = &'_ V> { | |
150 | match self { | |
151 | SsoHashMap::Array(array) => EitherIter::Left(array.iter().map(|(_k, v)| v)), | |
152 | SsoHashMap::Map(map) => EitherIter::Right(map.values()), | |
153 | } | |
154 | } | |
155 | ||
156 | /// An iterator visiting all values mutably in arbitrary order. | |
157 | /// The iterator element type is `&'a mut V`. | |
158 | pub fn values_mut(&mut self) -> impl Iterator<Item = &'_ mut V> { | |
159 | match self { | |
160 | SsoHashMap::Array(array) => EitherIter::Left(array.iter_mut().map(|(_k, v)| v)), | |
161 | SsoHashMap::Map(map) => EitherIter::Right(map.values_mut()), | |
162 | } | |
163 | } | |
164 | ||
165 | /// Clears the map, returning all key-value pairs as an iterator. Keeps the | |
166 | /// allocated memory for reuse. | |
167 | pub fn drain(&mut self) -> impl Iterator<Item = (K, V)> + '_ { | |
168 | match self { | |
169 | SsoHashMap::Array(array) => EitherIter::Left(array.drain(..)), | |
170 | SsoHashMap::Map(map) => EitherIter::Right(map.drain()), | |
171 | } | |
172 | } | |
173 | } | |
174 | ||
175 | impl<K: Eq + Hash, V> SsoHashMap<K, V> { | |
176 | /// Changes underlying storage from array to hashmap | |
177 | /// if array is full. | |
178 | fn migrate_if_full(&mut self) { | |
179 | if let SsoHashMap::Array(array) = self { | |
180 | if array.is_full() { | |
181 | *self = SsoHashMap::Map(array.drain(..).collect()); | |
182 | } | |
183 | } | |
184 | } | |
185 | ||
186 | /// Reserves capacity for at least `additional` more elements to be inserted | |
187 | /// in the `SsoHashMap`. The collection may reserve more space to avoid | |
188 | /// frequent reallocations. | |
189 | pub fn reserve(&mut self, additional: usize) { | |
190 | match self { | |
191 | SsoHashMap::Array(array) => { | |
192 | if SSO_ARRAY_SIZE < (array.len() + additional) { | |
193 | let mut map: FxHashMap<K, V> = array.drain(..).collect(); | |
194 | map.reserve(additional); | |
195 | *self = SsoHashMap::Map(map); | |
196 | } | |
197 | } | |
198 | SsoHashMap::Map(map) => map.reserve(additional), | |
199 | } | |
200 | } | |
201 | ||
202 | /// Shrinks the capacity of the map as much as possible. It will drop | |
203 | /// down as much as possible while maintaining the internal rules | |
204 | /// and possibly leaving some space in accordance with the resize policy. | |
205 | pub fn shrink_to_fit(&mut self) { | |
206 | if let SsoHashMap::Map(map) = self { | |
207 | if map.len() <= SSO_ARRAY_SIZE { | |
208 | *self = SsoHashMap::Array(map.drain().collect()); | |
209 | } else { | |
210 | map.shrink_to_fit(); | |
211 | } | |
212 | } | |
213 | } | |
214 | ||
215 | /// Retains only the elements specified by the predicate. | |
216 | pub fn retain<F>(&mut self, mut f: F) | |
217 | where | |
218 | F: FnMut(&K, &mut V) -> bool, | |
219 | { | |
220 | match self { | |
221 | SsoHashMap::Array(array) => array.retain(|(k, v)| f(k, v)), | |
222 | SsoHashMap::Map(map) => map.retain(f), | |
223 | } | |
224 | } | |
225 | ||
226 | /// Inserts a key-value pair into the map. | |
227 | /// | |
228 | /// If the map did not have this key present, [`None`] is returned. | |
229 | /// | |
230 | /// If the map did have this key present, the value is updated, and the old | |
231 | /// value is returned. The key is not updated, though; this matters for | |
232 | /// types that can be `==` without being identical. See the [module-level | |
233 | /// documentation] for more. | |
234 | pub fn insert(&mut self, key: K, value: V) -> Option<V> { | |
235 | match self { | |
236 | SsoHashMap::Array(array) => { | |
237 | for (k, v) in array.iter_mut() { | |
238 | if *k == key { | |
239 | let old_value = std::mem::replace(v, value); | |
240 | return Some(old_value); | |
241 | } | |
242 | } | |
243 | if let Err(error) = array.try_push((key, value)) { | |
244 | let mut map: FxHashMap<K, V> = array.drain(..).collect(); | |
245 | let (key, value) = error.element(); | |
246 | map.insert(key, value); | |
247 | *self = SsoHashMap::Map(map); | |
248 | } | |
249 | None | |
250 | } | |
251 | SsoHashMap::Map(map) => map.insert(key, value), | |
252 | } | |
253 | } | |
254 | ||
255 | /// Removes a key from the map, returning the value at the key if the key | |
256 | /// was previously in the map. | |
257 | pub fn remove(&mut self, key: &K) -> Option<V> { | |
258 | match self { | |
259 | SsoHashMap::Array(array) => { | |
c295e0f8 | 260 | array.iter().position(|(k, _v)| k == key).map(|index| array.swap_remove(index).1) |
29967ef6 XL |
261 | } |
262 | SsoHashMap::Map(map) => map.remove(key), | |
263 | } | |
264 | } | |
265 | ||
266 | /// Removes a key from the map, returning the stored key and value if the | |
267 | /// key was previously in the map. | |
268 | pub fn remove_entry(&mut self, key: &K) -> Option<(K, V)> { | |
269 | match self { | |
270 | SsoHashMap::Array(array) => { | |
c295e0f8 | 271 | array.iter().position(|(k, _v)| k == key).map(|index| array.swap_remove(index)) |
29967ef6 XL |
272 | } |
273 | SsoHashMap::Map(map) => map.remove_entry(key), | |
274 | } | |
275 | } | |
276 | ||
277 | /// Returns a reference to the value corresponding to the key. | |
278 | pub fn get(&self, key: &K) -> Option<&V> { | |
279 | match self { | |
280 | SsoHashMap::Array(array) => { | |
281 | for (k, v) in array { | |
282 | if k == key { | |
283 | return Some(v); | |
284 | } | |
285 | } | |
286 | None | |
287 | } | |
288 | SsoHashMap::Map(map) => map.get(key), | |
289 | } | |
290 | } | |
291 | ||
292 | /// Returns a mutable reference to the value corresponding to the key. | |
293 | pub fn get_mut(&mut self, key: &K) -> Option<&mut V> { | |
294 | match self { | |
295 | SsoHashMap::Array(array) => { | |
296 | for (k, v) in array { | |
297 | if k == key { | |
298 | return Some(v); | |
299 | } | |
300 | } | |
301 | None | |
302 | } | |
303 | SsoHashMap::Map(map) => map.get_mut(key), | |
304 | } | |
305 | } | |
306 | ||
307 | /// Returns the key-value pair corresponding to the supplied key. | |
308 | pub fn get_key_value(&self, key: &K) -> Option<(&K, &V)> { | |
309 | match self { | |
310 | SsoHashMap::Array(array) => { | |
311 | for (k, v) in array { | |
312 | if k == key { | |
313 | return Some((k, v)); | |
314 | } | |
315 | } | |
316 | None | |
317 | } | |
318 | SsoHashMap::Map(map) => map.get_key_value(key), | |
319 | } | |
320 | } | |
321 | ||
322 | /// Returns `true` if the map contains a value for the specified key. | |
323 | pub fn contains_key(&self, key: &K) -> bool { | |
324 | match self { | |
325 | SsoHashMap::Array(array) => array.iter().any(|(k, _v)| k == key), | |
326 | SsoHashMap::Map(map) => map.contains_key(key), | |
327 | } | |
328 | } | |
329 | ||
330 | /// Gets the given key's corresponding entry in the map for in-place manipulation. | |
331 | #[inline] | |
332 | pub fn entry(&mut self, key: K) -> Entry<'_, K, V> { | |
333 | Entry { ssomap: self, key } | |
334 | } | |
335 | } | |
336 | ||
337 | impl<K, V> Default for SsoHashMap<K, V> { | |
338 | #[inline] | |
339 | fn default() -> Self { | |
340 | Self::new() | |
341 | } | |
342 | } | |
343 | ||
344 | impl<K: Eq + Hash, V> FromIterator<(K, V)> for SsoHashMap<K, V> { | |
345 | fn from_iter<I: IntoIterator<Item = (K, V)>>(iter: I) -> SsoHashMap<K, V> { | |
346 | let mut map: SsoHashMap<K, V> = Default::default(); | |
347 | map.extend(iter); | |
348 | map | |
349 | } | |
350 | } | |
351 | ||
352 | impl<K: Eq + Hash, V> Extend<(K, V)> for SsoHashMap<K, V> { | |
353 | fn extend<I>(&mut self, iter: I) | |
354 | where | |
355 | I: IntoIterator<Item = (K, V)>, | |
356 | { | |
357 | for (key, value) in iter.into_iter() { | |
358 | self.insert(key, value); | |
359 | } | |
360 | } | |
361 | ||
362 | #[inline] | |
363 | fn extend_one(&mut self, (k, v): (K, V)) { | |
364 | self.insert(k, v); | |
365 | } | |
366 | ||
367 | fn extend_reserve(&mut self, additional: usize) { | |
368 | match self { | |
369 | SsoHashMap::Array(array) => { | |
370 | if SSO_ARRAY_SIZE < (array.len() + additional) { | |
371 | let mut map: FxHashMap<K, V> = array.drain(..).collect(); | |
372 | map.extend_reserve(additional); | |
373 | *self = SsoHashMap::Map(map); | |
374 | } | |
375 | } | |
376 | SsoHashMap::Map(map) => map.extend_reserve(additional), | |
377 | } | |
378 | } | |
379 | } | |
380 | ||
381 | impl<'a, K, V> Extend<(&'a K, &'a V)> for SsoHashMap<K, V> | |
382 | where | |
383 | K: Eq + Hash + Copy, | |
384 | V: Copy, | |
385 | { | |
386 | fn extend<T: IntoIterator<Item = (&'a K, &'a V)>>(&mut self, iter: T) { | |
387 | self.extend(iter.into_iter().map(|(k, v)| (*k, *v))) | |
388 | } | |
389 | ||
390 | #[inline] | |
391 | fn extend_one(&mut self, (&k, &v): (&'a K, &'a V)) { | |
392 | self.insert(k, v); | |
393 | } | |
394 | ||
395 | #[inline] | |
396 | fn extend_reserve(&mut self, additional: usize) { | |
397 | Extend::<(K, V)>::extend_reserve(self, additional) | |
398 | } | |
399 | } | |
400 | ||
401 | impl<K, V> IntoIterator for SsoHashMap<K, V> { | |
402 | type IntoIter = EitherIter< | |
cdc7bbd5 | 403 | <ArrayVec<(K, V), 8> as IntoIterator>::IntoIter, |
29967ef6 XL |
404 | <FxHashMap<K, V> as IntoIterator>::IntoIter, |
405 | >; | |
406 | type Item = <Self::IntoIter as Iterator>::Item; | |
407 | ||
408 | fn into_iter(self) -> Self::IntoIter { | |
409 | match self { | |
410 | SsoHashMap::Array(array) => EitherIter::Left(array.into_iter()), | |
411 | SsoHashMap::Map(map) => EitherIter::Right(map.into_iter()), | |
412 | } | |
413 | } | |
414 | } | |
415 | ||
416 | /// adapts Item of array reference iterator to Item of hashmap reference iterator. | |
417 | #[inline(always)] | |
c295e0f8 | 418 | fn adapt_array_ref_it<K, V>(pair: &(K, V)) -> (&K, &V) { |
29967ef6 XL |
419 | let (a, b) = pair; |
420 | (a, b) | |
421 | } | |
422 | ||
423 | /// adapts Item of array mut reference iterator to Item of hashmap mut reference iterator. | |
424 | #[inline(always)] | |
c295e0f8 | 425 | fn adapt_array_mut_it<K, V>(pair: &mut (K, V)) -> (&K, &mut V) { |
29967ef6 XL |
426 | let (a, b) = pair; |
427 | (a, b) | |
428 | } | |
429 | ||
430 | impl<'a, K, V> IntoIterator for &'a SsoHashMap<K, V> { | |
431 | type IntoIter = EitherIter< | |
432 | std::iter::Map< | |
cdc7bbd5 | 433 | <&'a ArrayVec<(K, V), 8> as IntoIterator>::IntoIter, |
29967ef6 XL |
434 | fn(&'a (K, V)) -> (&'a K, &'a V), |
435 | >, | |
436 | <&'a FxHashMap<K, V> as IntoIterator>::IntoIter, | |
437 | >; | |
438 | type Item = <Self::IntoIter as Iterator>::Item; | |
439 | ||
440 | fn into_iter(self) -> Self::IntoIter { | |
441 | match self { | |
442 | SsoHashMap::Array(array) => EitherIter::Left(array.into_iter().map(adapt_array_ref_it)), | |
443 | SsoHashMap::Map(map) => EitherIter::Right(map.iter()), | |
444 | } | |
445 | } | |
446 | } | |
447 | ||
448 | impl<'a, K, V> IntoIterator for &'a mut SsoHashMap<K, V> { | |
449 | type IntoIter = EitherIter< | |
450 | std::iter::Map< | |
cdc7bbd5 | 451 | <&'a mut ArrayVec<(K, V), 8> as IntoIterator>::IntoIter, |
29967ef6 XL |
452 | fn(&'a mut (K, V)) -> (&'a K, &'a mut V), |
453 | >, | |
454 | <&'a mut FxHashMap<K, V> as IntoIterator>::IntoIter, | |
455 | >; | |
456 | type Item = <Self::IntoIter as Iterator>::Item; | |
457 | ||
458 | fn into_iter(self) -> Self::IntoIter { | |
459 | match self { | |
460 | SsoHashMap::Array(array) => EitherIter::Left(array.into_iter().map(adapt_array_mut_it)), | |
461 | SsoHashMap::Map(map) => EitherIter::Right(map.iter_mut()), | |
462 | } | |
463 | } | |
464 | } | |
465 | ||
466 | impl<K, V> fmt::Debug for SsoHashMap<K, V> | |
467 | where | |
468 | K: fmt::Debug, | |
469 | V: fmt::Debug, | |
470 | { | |
471 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { | |
472 | f.debug_map().entries(self.iter()).finish() | |
473 | } | |
474 | } | |
475 | ||
476 | impl<'a, K, V> Index<&'a K> for SsoHashMap<K, V> | |
477 | where | |
478 | K: Eq + Hash, | |
479 | { | |
480 | type Output = V; | |
481 | ||
482 | #[inline] | |
483 | fn index(&self, key: &K) -> &V { | |
484 | self.get(key).expect("no entry found for key") | |
485 | } | |
486 | } | |
487 | ||
488 | /// A view into a single entry in a map. | |
489 | pub struct Entry<'a, K, V> { | |
490 | ssomap: &'a mut SsoHashMap<K, V>, | |
491 | key: K, | |
492 | } | |
493 | ||
494 | impl<'a, K: Eq + Hash, V> Entry<'a, K, V> { | |
495 | /// Provides in-place mutable access to an occupied entry before any | |
496 | /// potential inserts into the map. | |
497 | pub fn and_modify<F>(self, f: F) -> Self | |
498 | where | |
499 | F: FnOnce(&mut V), | |
500 | { | |
501 | if let Some(value) = self.ssomap.get_mut(&self.key) { | |
502 | f(value); | |
503 | } | |
504 | self | |
505 | } | |
506 | ||
507 | /// Ensures a value is in the entry by inserting the default if empty, and returns | |
508 | /// a mutable reference to the value in the entry. | |
509 | #[inline] | |
510 | pub fn or_insert(self, value: V) -> &'a mut V { | |
511 | self.or_insert_with(|| value) | |
512 | } | |
513 | ||
514 | /// Ensures a value is in the entry by inserting the result of the default function if empty, | |
515 | /// and returns a mutable reference to the value in the entry. | |
516 | pub fn or_insert_with<F: FnOnce() -> V>(self, default: F) -> &'a mut V { | |
517 | self.ssomap.migrate_if_full(); | |
518 | match self.ssomap { | |
519 | SsoHashMap::Array(array) => { | |
520 | let key_ref = &self.key; | |
521 | let found_index = array.iter().position(|(k, _v)| k == key_ref); | |
522 | let index = if let Some(index) = found_index { | |
523 | index | |
524 | } else { | |
525 | let index = array.len(); | |
526 | array.try_push((self.key, default())).unwrap(); | |
527 | index | |
528 | }; | |
529 | &mut array[index].1 | |
530 | } | |
531 | SsoHashMap::Map(map) => map.entry(self.key).or_insert_with(default), | |
532 | } | |
533 | } | |
534 | ||
535 | /// Returns a reference to this entry's key. | |
536 | #[inline] | |
537 | pub fn key(&self) -> &K { | |
538 | &self.key | |
539 | } | |
540 | } | |
541 | ||
542 | impl<'a, K: Eq + Hash, V: Default> Entry<'a, K, V> { | |
543 | /// Ensures a value is in the entry by inserting the default value if empty, | |
544 | /// and returns a mutable reference to the value in the entry. | |
545 | #[inline] | |
546 | pub fn or_default(self) -> &'a mut V { | |
547 | self.or_insert_with(Default::default) | |
548 | } | |
549 | } |