rustc_data_structures/sso/
map.rs

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

rustc_data_structures/sso/
map.rs
