auto_hash_map/
map.rs

1use std::{
2    borrow::Borrow,
3    fmt::{Debug, Formatter},
4    hash::{BuildHasher, BuildHasherDefault, Hash},
5    marker::PhantomData,
6};
7
8use bincode::{
9    Decode, Encode,
10    de::Decoder,
11    enc::Encoder,
12    error::{DecodeError, EncodeError},
13    impl_borrow_decode,
14};
15use hashbrown::hash_map::HashMap;
16use rustc_hash::FxHasher;
17use serde::{
18    Deserialize, Deserializer, Serialize, Serializer,
19    de::{MapAccess, Visitor},
20    ser::SerializeMap,
21};
22use shrink_to_fit::ShrinkToFit;
23
24use crate::{MAX_USEFUL_LINEAR_SCAN, MIN_HASH_SIZE, TinyVec, tiny_vec};
25
26#[derive(Clone)]
27pub enum AutoMap<K, V, H = BuildHasherDefault<FxHasher>, const I: usize = 0> {
28    List(TinyVec<(K, V), I, MAX_USEFUL_LINEAR_SCAN>),
29    Map(Box<HashMap<K, V, H>>),
30}
31
32impl<K, V, H, const I: usize> Default for AutoMap<K, V, H, I> {
33    fn default() -> Self {
34        Self::List(Default::default())
35    }
36}
37
38impl<K: Debug, V: Debug, H, const I: usize> Debug for AutoMap<K, V, H, I> {
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        f.debug_map().entries(self.iter()).finish()
41    }
42}
43
44impl<K, V> AutoMap<K, V, BuildHasherDefault<FxHasher>, 0> {
45    /// see [HashMap::new](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.new)
46    pub const fn new() -> Self {
47        AutoMap::List(TinyVec::new())
48    }
49
50    /// see [HashMap::with_capacity](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.with_capacity)
51    pub fn with_capacity(capacity: usize) -> Self {
52        if capacity < MAX_USEFUL_LINEAR_SCAN {
53            AutoMap::List(TinyVec::with_capacity(capacity))
54        } else {
55            AutoMap::Map(Box::new(HashMap::with_capacity_and_hasher(
56                capacity,
57                Default::default(),
58            )))
59        }
60    }
61}
62
63impl<K, V, H, const I: usize> AutoMap<K, V, H, I> {
64    /// see [HashMap::with_hasher](https://doc.rust-lang.org/std/collections/hash_map/struct.HashMap.html#method.with_hasher)
65    pub const fn with_hasher() -> Self {
66        AutoMap::List(TinyVec::new())
67    }
68
69    /// see [HashMap::with_capacity_and_hasher](https://doc.rust-lang.org/std/collections/hash_map/struct.HashMap.html#method.with_capacity_and_hasher)
70    pub fn with_capacity_and_hasher(capacity: usize, hasher: H) -> Self {
71        if capacity <= MAX_USEFUL_LINEAR_SCAN {
72            AutoMap::List(TinyVec::with_capacity(capacity))
73        } else {
74            AutoMap::Map(Box::new(HashMap::with_capacity_and_hasher(
75                capacity, hasher,
76            )))
77        }
78    }
79
80    /// see [HashMap::clear](https://doc.rust-lang.org/std/collections/hash_map/struct.HashMap.html#method.clear)
81    pub fn clear(&mut self) {
82        match self {
83            AutoMap::List(list) => list.clear(),
84            AutoMap::Map(map) => map.clear(),
85        }
86    }
87}
88
89impl<K: Eq + Hash, V, H: BuildHasher + Default, const I: usize> AutoMap<K, V, H, I> {
90    fn convert_to_map(&mut self) -> &mut HashMap<K, V, H> {
91        if let AutoMap::List(list) = self {
92            let mut map =
93                HashMap::with_capacity_and_hasher(MAX_USEFUL_LINEAR_SCAN * 2, Default::default());
94            map.extend(list.drain());
95            *self = AutoMap::Map(Box::new(map));
96        }
97        if let AutoMap::Map(map) = self {
98            map
99        } else {
100            unreachable!()
101        }
102    }
103
104    fn convert_to_list(&mut self) -> &mut TinyVec<(K, V), I, MAX_USEFUL_LINEAR_SCAN> {
105        if let AutoMap::Map(map) = self {
106            let mut list = TinyVec::with_capacity(MAX_USEFUL_LINEAR_SCAN);
107            list.extend(map.drain());
108            *self = AutoMap::List(list);
109        }
110        if let AutoMap::List(list) = self {
111            list
112        } else {
113            unreachable!()
114        }
115    }
116
117    /// see [HashMap::insert](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.insert)
118    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
119        match self {
120            AutoMap::List(list) => {
121                for (k, v) in list.iter_mut() {
122                    if *k == key {
123                        return Some(std::mem::replace(v, value));
124                    }
125                }
126                if list.len() >= MAX_USEFUL_LINEAR_SCAN {
127                    let map = self.convert_to_map();
128                    map.insert(key, value);
129                } else {
130                    list.push((key, value));
131                }
132                None
133            }
134            AutoMap::Map(map) => map.insert(key, value),
135        }
136    }
137
138    /// see [HashMap::remove](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.remove)
139    pub fn remove(&mut self, key: &K) -> Option<V> {
140        match self {
141            AutoMap::List(list) => {
142                for i in 0..list.len() {
143                    if list[i].0 == *key {
144                        return Some(list.swap_remove(i).1);
145                    }
146                }
147                None
148            }
149            AutoMap::Map(map) => map.remove(key),
150        }
151    }
152
153    /// see [HashMap::extend](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.extend)
154    pub fn extend(&mut self, iter: impl IntoIterator<Item = (K, V)>) {
155        let iter = iter.into_iter();
156        match self {
157            AutoMap::List(list) => {
158                let (lower, _) = iter.size_hint();
159                if list.len() + lower > MAX_USEFUL_LINEAR_SCAN {
160                    let map = self.convert_to_map();
161                    map.extend(iter);
162                    // The hint is not enforced
163                    if map.len() < MIN_HASH_SIZE {
164                        self.convert_to_list();
165                    }
166                    return;
167                }
168                for (k, v) in iter {
169                    self.insert(k, v);
170                }
171            }
172            AutoMap::Map(map) => {
173                map.extend(iter);
174            }
175        }
176    }
177
178    /// see [HashMap::entry](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.entry)
179    pub fn entry(&mut self, key: K) -> Entry<'_, K, V, H, I> {
180        let this = self as *mut Self;
181        match self {
182            AutoMap::List(list) => match list.iter().position(|(k, _)| *k == key) {
183                Some(index) => Entry::Occupied(OccupiedEntry::List { list, index }),
184                None => Entry::Vacant(VacantEntry::List { this, list, key }),
185            },
186            AutoMap::Map(map) => match map.entry(key) {
187                hashbrown::hash_map::Entry::Occupied(entry) => {
188                    Entry::Occupied(OccupiedEntry::Map { this, entry })
189                }
190                hashbrown::hash_map::Entry::Vacant(entry) => Entry::Vacant(VacantEntry::Map(entry)),
191            },
192        }
193    }
194
195    pub fn raw_entry_mut<Q>(&mut self, key: &Q) -> RawEntry<'_, K, V, H, I>
196    where
197        K: Borrow<Q>,
198        Q: Hash + Eq + ?Sized,
199    {
200        let this = self as *mut Self;
201        match self {
202            AutoMap::List(list) => match list.iter().position(|(k, _)| k.borrow() == key) {
203                Some(index) => RawEntry::Occupied(OccupiedRawEntry::List { list, index }),
204                None => RawEntry::Vacant(VacantRawEntry::List { this, list }),
205            },
206            AutoMap::Map(map) => match map.raw_entry_mut().from_key(key) {
207                hashbrown::hash_map::RawEntryMut::Occupied(entry) => {
208                    RawEntry::Occupied(OccupiedRawEntry::Map { this, entry })
209                }
210                hashbrown::hash_map::RawEntryMut::Vacant(entry) => {
211                    RawEntry::Vacant(VacantRawEntry::Map(entry))
212                }
213            },
214        }
215    }
216
217    /// see [HashMap::retain](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.retain)
218    pub fn retain<F>(&mut self, mut f: F)
219    where
220        F: FnMut(&K, &mut V) -> bool,
221    {
222        match self {
223            AutoMap::List(list) => {
224                // Don't use `Vec::retain`, as that uses a slower algorithm to maintain order,
225                // which we don't care about
226                let mut len = list.len();
227                let mut i = 0;
228                while i < len {
229                    let (key, value) = &mut list[i];
230                    if !f(key, value) {
231                        list.swap_remove(i);
232                        len -= 1;
233                    } else {
234                        i += 1;
235                    }
236                }
237            }
238            AutoMap::Map(map) => {
239                map.retain(f);
240            }
241        }
242    }
243
244    pub fn extract_if<'l, F>(&'l mut self, f: F) -> ExtractIfIter<'l, K, V, I, F>
245    where
246        F: for<'a, 'b> FnMut(&'a K, &'b mut V) -> bool,
247    {
248        match self {
249            AutoMap::List(list) => ExtractIfIter::List { list, index: 0, f },
250            AutoMap::Map(map) => ExtractIfIter::Map(map.extract_if(f)),
251        }
252    }
253
254    /// see [HashMap::shrink_to_fit](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.shrink_to_fit)
255    pub fn shrink_to_fit(&mut self) {
256        match self {
257            AutoMap::List(list) => list.shrink_to_fit(),
258            AutoMap::Map(map) => {
259                if map.len() <= MAX_USEFUL_LINEAR_SCAN {
260                    let mut list = TinyVec::with_capacity(map.len());
261                    list.extend(map.drain());
262                    *self = AutoMap::List(list);
263                } else {
264                    map.shrink_to_fit();
265                }
266            }
267        }
268    }
269
270    pub fn shrink_amortized(&mut self) {
271        match self {
272            AutoMap::List(list) => {
273                // `TinyVec::shrink_to_fit` returns spilled storage to inline
274                // (or a tighter heap buffer) when the occupancy warrants it.
275                if list.capacity() > list.len() * 3 {
276                    list.shrink_to_fit();
277                }
278            }
279            AutoMap::Map(map) => {
280                if map.len() <= MIN_HASH_SIZE {
281                    let mut list = TinyVec::with_capacity(map.len());
282                    list.extend(map.drain());
283                    *self = AutoMap::List(list);
284                } else if map.capacity() > map.len() * 3 {
285                    // HashMaps will always have a capacity that is a power of two
286                    map.shrink_to_fit();
287                }
288            }
289        }
290    }
291}
292
293impl<K: Eq + Hash, V, H: BuildHasher, const I: usize> AutoMap<K, V, H, I> {
294    /// see [HashMap::get](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.get)
295    pub fn get<Q>(&self, key: &Q) -> Option<&V>
296    where
297        K: Borrow<Q>,
298        Q: Hash + Eq + ?Sized,
299    {
300        match self {
301            AutoMap::List(list) => list
302                .iter()
303                .find_map(|(k, v)| if *k.borrow() == *key { Some(v) } else { None }),
304            AutoMap::Map(map) => map.get(key),
305        }
306    }
307
308    /// see [HashMap::get_mut](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.get_mut)
309    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
310        match self {
311            AutoMap::List(list) => list
312                .iter_mut()
313                .find_map(|(k, v)| if *k == *key { Some(v) } else { None }),
314            AutoMap::Map(map) => map.get_mut(key),
315        }
316    }
317
318    /// see [HashMap::contains_key](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.contains_key)
319    pub fn contains_key(&self, key: &K) -> bool {
320        match self {
321            AutoMap::List(list) => list.iter().any(|(k, _)| *k == *key),
322            AutoMap::Map(map) => map.contains_key(key),
323        }
324    }
325}
326
327impl<K, V, H, const I: usize> AutoMap<K, V, H, I> {
328    /// see [HashMap::iter](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.iter)
329    pub fn iter(&self) -> Iter<'_, K, V> {
330        match self {
331            AutoMap::List(list) => Iter::List(list.iter()),
332            AutoMap::Map(map) => Iter::Map(map.iter()),
333        }
334    }
335    /// see [HashMap::iter_mut](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.iter_mut)
336    pub fn iter_mut(&mut self) -> IterMut<'_, K, V> {
337        match self {
338            AutoMap::List(list) => IterMut::List(list.iter_mut()),
339            AutoMap::Map(map) => IterMut::Map(map.iter_mut()),
340        }
341    }
342
343    /// see [HashMap::is_empty](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.is_empty)
344    pub fn is_empty(&self) -> bool {
345        match self {
346            AutoMap::List(list) => list.is_empty(),
347            AutoMap::Map(map) => map.is_empty(),
348        }
349    }
350
351    /// see [HashMap::len](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.len)
352    pub fn len(&self) -> usize {
353        match self {
354            AutoMap::List(list) => list.len(),
355            AutoMap::Map(map) => map.len(),
356        }
357    }
358
359    /// see [HashMap::values_mut](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.values_mut)
360    pub fn values_mut(&mut self) -> ValuesMut<'_, K, V> {
361        match self {
362            AutoMap::List(list) => ValuesMut::List(list.iter_mut()),
363            AutoMap::Map(map) => ValuesMut::Map(map.values_mut()),
364        }
365    }
366
367    /// see [HashMap::values](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.values)
368    pub fn values(&self) -> Values<'_, K, V> {
369        match self {
370            AutoMap::List(list) => Values::List(list.iter()),
371            AutoMap::Map(map) => Values::Map(map.values()),
372        }
373    }
374
375    /// see [HashMap::into_values](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.into_values)
376    pub fn into_values(self) -> IntoValues<K, V, I> {
377        match self {
378            AutoMap::List(list) => IntoValues::List(list.into_iter()),
379            AutoMap::Map(map) => IntoValues::Map(map.into_values()),
380        }
381    }
382}
383
384impl<K, V, H, const I: usize> IntoIterator for AutoMap<K, V, H, I> {
385    type Item = (K, V);
386    type IntoIter = IntoIter<K, V, I>;
387
388    fn into_iter(self) -> Self::IntoIter {
389        match self {
390            AutoMap::List(list) => IntoIter::List(list.into_iter()),
391            AutoMap::Map(map) => IntoIter::Map(map.into_iter()),
392        }
393    }
394}
395
396impl<'a, K, V, H, const I: usize> IntoIterator for &'a AutoMap<K, V, H, I> {
397    type Item = (&'a K, &'a V);
398    type IntoIter = Iter<'a, K, V>;
399
400    fn into_iter(self) -> Self::IntoIter {
401        self.iter()
402    }
403}
404
405pub enum Iter<'a, K, V> {
406    List(std::slice::Iter<'a, (K, V)>),
407    Map(hashbrown::hash_map::Iter<'a, K, V>),
408}
409
410impl<'a, K, V> Iterator for Iter<'a, K, V> {
411    type Item = (&'a K, &'a V);
412
413    // This clippy lint doesn't account for lifetimes
414    #[allow(clippy::map_identity)]
415    fn next(&mut self) -> Option<Self::Item> {
416        match self {
417            Iter::List(iter) => iter.next().map(|(k, v)| (k, v)),
418            Iter::Map(iter) => iter.next(),
419        }
420    }
421
422    fn size_hint(&self) -> (usize, Option<usize>) {
423        match self {
424            Iter::List(iter) => iter.size_hint(),
425            Iter::Map(iter) => iter.size_hint(),
426        }
427    }
428}
429
430impl<K, V> Clone for Iter<'_, K, V> {
431    fn clone(&self) -> Self {
432        match self {
433            Iter::List(iter) => Iter::List(iter.clone()),
434            Iter::Map(iter) => Iter::Map(iter.clone()),
435        }
436    }
437}
438
439pub enum IterMut<'a, K, V> {
440    List(std::slice::IterMut<'a, (K, V)>),
441    Map(hashbrown::hash_map::IterMut<'a, K, V>),
442}
443
444impl<'a, K, V> Iterator for IterMut<'a, K, V> {
445    type Item = (&'a K, &'a mut V);
446
447    fn next(&mut self) -> Option<Self::Item> {
448        match self {
449            IterMut::List(iter) => iter.next().map(|(k, v)| (&*k, v)),
450            IterMut::Map(iter) => iter.next(),
451        }
452    }
453
454    fn size_hint(&self) -> (usize, Option<usize>) {
455        match self {
456            IterMut::List(iter) => iter.size_hint(),
457            IterMut::Map(iter) => iter.size_hint(),
458        }
459    }
460}
461
462pub enum IntoIter<K, V, const I: usize> {
463    List(tiny_vec::IntoIter<(K, V), I, MAX_USEFUL_LINEAR_SCAN>),
464    Map(hashbrown::hash_map::IntoIter<K, V>),
465}
466
467impl<K, V, const I: usize> Iterator for IntoIter<K, V, I> {
468    type Item = (K, V);
469
470    fn next(&mut self) -> Option<Self::Item> {
471        match self {
472            IntoIter::List(iter) => iter.next(),
473            IntoIter::Map(iter) => iter.next(),
474        }
475    }
476
477    fn size_hint(&self) -> (usize, Option<usize>) {
478        match self {
479            IntoIter::List(iter) => iter.size_hint(),
480            IntoIter::Map(iter) => iter.size_hint(),
481        }
482    }
483}
484
485pub enum Values<'a, K, V> {
486    List(std::slice::Iter<'a, (K, V)>),
487    Map(hashbrown::hash_map::Values<'a, K, V>),
488}
489
490impl<'a, K, V> Iterator for Values<'a, K, V> {
491    type Item = &'a V;
492
493    fn next(&mut self) -> Option<Self::Item> {
494        match self {
495            Values::List(iter) => iter.next().map(|(_, v)| v),
496            Values::Map(iter) => iter.next(),
497        }
498    }
499
500    fn size_hint(&self) -> (usize, Option<usize>) {
501        match self {
502            Values::List(iter) => iter.size_hint(),
503            Values::Map(iter) => iter.size_hint(),
504        }
505    }
506}
507
508pub enum ValuesMut<'a, K, V> {
509    List(std::slice::IterMut<'a, (K, V)>),
510    Map(hashbrown::hash_map::ValuesMut<'a, K, V>),
511}
512
513impl<'a, K, V> Iterator for ValuesMut<'a, K, V> {
514    type Item = &'a mut V;
515
516    fn next(&mut self) -> Option<Self::Item> {
517        match self {
518            ValuesMut::List(iter) => iter.next().map(|(_, v)| v),
519            ValuesMut::Map(iter) => iter.next(),
520        }
521    }
522
523    fn size_hint(&self) -> (usize, Option<usize>) {
524        match self {
525            ValuesMut::List(iter) => iter.size_hint(),
526            ValuesMut::Map(iter) => iter.size_hint(),
527        }
528    }
529}
530
531pub enum IntoValues<K, V, const I: usize> {
532    List(tiny_vec::IntoIter<(K, V), I, MAX_USEFUL_LINEAR_SCAN>),
533    Map(hashbrown::hash_map::IntoValues<K, V>),
534}
535
536impl<K, V, const I: usize> Iterator for IntoValues<K, V, I> {
537    type Item = V;
538
539    fn next(&mut self) -> Option<Self::Item> {
540        match self {
541            IntoValues::List(iter) => iter.next().map(|(_, v)| v),
542            IntoValues::Map(iter) => iter.next(),
543        }
544    }
545
546    fn size_hint(&self) -> (usize, Option<usize>) {
547        match self {
548            IntoValues::List(iter) => iter.size_hint(),
549            IntoValues::Map(iter) => iter.size_hint(),
550        }
551    }
552}
553
554pub enum Entry<'a, K, V, H, const I: usize> {
555    Occupied(OccupiedEntry<'a, K, V, H, I>),
556    Vacant(VacantEntry<'a, K, V, H, I>),
557}
558
559impl<'a, K: Eq + Hash, V, H: BuildHasher + Default + 'a, const I: usize> Entry<'a, K, V, H, I> {
560    /// see [HashMap::Entry::or_insert](https://doc.rust-lang.org/std/collections/hash_map/enum.Entry.html#method.or_insert)
561    pub fn or_insert_with(self, default: impl FnOnce() -> V) -> &'a mut V {
562        match self {
563            Entry::Occupied(entry) => entry.into_mut(),
564            Entry::Vacant(entry) => entry.insert(default()),
565        }
566    }
567
568    /// see [HashMap::Entry::or_insert](https://doc.rust-lang.org/std/collections/hash_map/enum.Entry.html#method.or_insert)
569    pub fn or_insert(self, default: V) -> &'a mut V {
570        match self {
571            Entry::Occupied(entry) => entry.into_mut(),
572            Entry::Vacant(entry) => entry.insert(default),
573        }
574    }
575}
576
577impl<'a, K: Eq + Hash, V: Default, H: BuildHasher + Default + 'a, const I: usize>
578    Entry<'a, K, V, H, I>
579{
580    /// see [HashMap::Entry::or_default](https://doc.rust-lang.org/std/collections/hash_map/enum.Entry.html#method.or_default)
581    pub fn or_default(self) -> &'a mut V {
582        match self {
583            Entry::Occupied(entry) => entry.into_mut(),
584            Entry::Vacant(entry) => entry.insert(Default::default()),
585        }
586    }
587}
588
589pub enum OccupiedEntry<'a, K, V, H, const I: usize> {
590    List {
591        list: &'a mut TinyVec<(K, V), I, MAX_USEFUL_LINEAR_SCAN>,
592        index: usize,
593    },
594    Map {
595        this: *mut AutoMap<K, V, H, I>,
596        entry: hashbrown::hash_map::OccupiedEntry<'a, K, V, H>,
597    },
598}
599
600impl<'a, K: Eq + Hash, V, H: BuildHasher, const I: usize> OccupiedEntry<'a, K, V, H, I> {
601    /// see [HashMap::OccupiedEntry::get_mut](https://doc.rust-lang.org/std/collections/hash_map/enum.OccupiedEntry.html#method.get_mut)
602    pub fn get_mut(&mut self) -> &mut V {
603        match self {
604            OccupiedEntry::List { list, index } => &mut list[*index].1,
605            OccupiedEntry::Map { entry, .. } => entry.get_mut(),
606        }
607    }
608
609    /// see [HashMap::OccupiedEntry::into_mut](https://doc.rust-lang.org/std/collections/hash_map/enum.OccupiedEntry.html#method.into_mut)
610    pub fn into_mut(self) -> &'a mut V {
611        match self {
612            OccupiedEntry::List { list, index } => &mut list[index].1,
613            OccupiedEntry::Map { entry, .. } => entry.into_mut(),
614        }
615    }
616}
617
618impl<K: Eq + Hash, V, H: BuildHasher + Default, const I: usize> OccupiedEntry<'_, K, V, H, I> {
619    /// see [HashMap::OccupiedEntry::remove](https://doc.rust-lang.org/std/collections/hash_map/enum.OccupiedEntry.html#method.remove)
620    pub fn remove(self) -> V {
621        match self {
622            OccupiedEntry::List { list, index } => list.swap_remove(index).1,
623            OccupiedEntry::Map { entry, this: _ } => entry.remove(),
624        }
625    }
626
627    pub fn replace_entry_with(self, func: impl FnOnce(&K, V) -> Option<V>) {
628        match self {
629            OccupiedEntry::List { list, index } => {
630                let (key, value) = list.swap_remove(index);
631                if let Some(value) = func(&key, value) {
632                    list.push((key, value));
633                }
634            }
635            OccupiedEntry::Map { entry, .. } => {
636                entry.replace_entry_with(func);
637            }
638        }
639    }
640}
641
642pub enum VacantEntry<'a, K, V, H, const I: usize> {
643    List {
644        this: *mut AutoMap<K, V, H, I>,
645        list: &'a mut TinyVec<(K, V), I, MAX_USEFUL_LINEAR_SCAN>,
646        key: K,
647    },
648    Map(hashbrown::hash_map::VacantEntry<'a, K, V, H>),
649}
650
651impl<'a, K: Eq + Hash, V, H: BuildHasher + Default + 'a, const I: usize>
652    VacantEntry<'a, K, V, H, I>
653{
654    /// see [HashMap::VacantEntry::insert](https://doc.rust-lang.org/std/collections/hash_map/enum.VacantEntry.html#method.insert)
655    pub fn insert(self, value: V) -> &'a mut V {
656        match self {
657            VacantEntry::List { this, list, key } => {
658                if list.len() >= MAX_USEFUL_LINEAR_SCAN {
659                    let this = unsafe { &mut *this };
660                    this.convert_to_map().entry(key).or_insert(value)
661                } else {
662                    list.push((key, value));
663                    &mut list.last_mut().unwrap().1
664                }
665            }
666            VacantEntry::Map(entry) => entry.insert(value),
667        }
668    }
669}
670
671pub enum RawEntry<'a, K, V, H, const I: usize> {
672    Occupied(OccupiedRawEntry<'a, K, V, H, I>),
673    Vacant(VacantRawEntry<'a, K, V, H, I>),
674}
675
676pub enum OccupiedRawEntry<'a, K, V, H, const I: usize> {
677    List {
678        list: &'a mut TinyVec<(K, V), I, MAX_USEFUL_LINEAR_SCAN>,
679        index: usize,
680    },
681    Map {
682        this: *mut AutoMap<K, V, H, I>,
683        entry: hashbrown::hash_map::RawOccupiedEntryMut<'a, K, V, H>,
684    },
685}
686
687impl<'a, K: Eq + Hash, V, H: BuildHasher, const I: usize> OccupiedRawEntry<'a, K, V, H, I> {
688    /// see [HashMap::RawOccupiedEntryMut::get_mut](https://doc.rust-lang.org/std/collections/hash_map/struct.RawOccupiedEntryMut.html#method.get_mut)
689    pub fn get_mut(&mut self) -> &mut V {
690        match self {
691            OccupiedRawEntry::List { list, index } => &mut list[*index].1,
692            OccupiedRawEntry::Map { entry, .. } => entry.get_mut(),
693        }
694    }
695
696    /// see [HashMap::RawOccupiedEntryMut::into_mut](https://doc.rust-lang.org/std/collections/hash_map/struct.RawOccupiedEntryMut.html#method.into_mut)
697    pub fn into_mut(self) -> &'a mut V {
698        match self {
699            OccupiedRawEntry::List { list, index } => &mut list[index].1,
700            OccupiedRawEntry::Map { entry, .. } => entry.into_mut(),
701        }
702    }
703}
704
705impl<K: Eq + Hash, V, H: BuildHasher + Default, const I: usize> OccupiedRawEntry<'_, K, V, H, I> {
706    /// see [HashMap::OccupiedEntry::remove](https://doc.rust-lang.org/std/collections/hash_map/enum.OccupiedEntry.html#method.remove)
707    pub fn remove(self) -> V {
708        match self {
709            OccupiedRawEntry::List { list, index } => list.swap_remove(index).1,
710            OccupiedRawEntry::Map { entry, this: _ } => entry.remove(),
711        }
712    }
713}
714
715pub enum VacantRawEntry<'a, K, V, H, const I: usize> {
716    List {
717        this: *mut AutoMap<K, V, H, I>,
718        list: &'a mut TinyVec<(K, V), I, MAX_USEFUL_LINEAR_SCAN>,
719    },
720    Map(hashbrown::hash_map::RawVacantEntryMut<'a, K, V, H>),
721}
722
723impl<'a, K: Eq + Hash, V, H: BuildHasher + Default + 'a, const I: usize>
724    VacantRawEntry<'a, K, V, H, I>
725{
726    /// see [HashMap::RawVacantEntryMut::insert](https://doc.rust-lang.org/std/collections/hash_map/struct.RawVacantEntryMut.html#method.insert)
727    pub fn insert(self, key: K, value: V) -> &'a mut V {
728        match self {
729            VacantRawEntry::List { this, list } => {
730                if list.len() >= MAX_USEFUL_LINEAR_SCAN {
731                    let this = unsafe { &mut *this };
732                    this.convert_to_map().entry(key).or_insert(value)
733                } else {
734                    list.push((key, value));
735                    &mut list.last_mut().unwrap().1
736                }
737            }
738            VacantRawEntry::Map(entry) => entry.insert(key, value).1,
739        }
740    }
741}
742
743impl<K, V, H, const I: usize> Serialize for AutoMap<K, V, H, I>
744where
745    K: Eq + Hash + Serialize,
746    V: Serialize,
747    H: BuildHasher,
748{
749    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
750    where
751        S: Serializer,
752    {
753        match self {
754            AutoMap::List(list) => {
755                let mut map = serializer.serialize_map(Some(list.len()))?;
756                for (k, v) in list {
757                    map.serialize_entry(k, v)?;
758                }
759                map.end()
760            }
761            AutoMap::Map(map) => (**map).serialize(serializer),
762        }
763    }
764}
765
766impl<'de, K, V, H, const I: usize> Deserialize<'de> for AutoMap<K, V, H, I>
767where
768    K: Eq + Hash + Deserialize<'de>,
769    V: Deserialize<'de>,
770    H: BuildHasher + Default,
771{
772    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
773    where
774        D: Deserializer<'de>,
775    {
776        struct AutoMapVisitor<K, V, H, const I: usize> {
777            phantom: PhantomData<AutoMap<K, V, H, I>>,
778        }
779
780        impl<'de, K, V, H, const I: usize> Visitor<'de> for AutoMapVisitor<K, V, H, I>
781        where
782            K: Eq + Hash + Deserialize<'de>,
783            V: Deserialize<'de>,
784            H: BuildHasher + Default,
785        {
786            type Value = AutoMap<K, V, H, I>;
787
788            fn expecting(&self, formatter: &mut Formatter) -> std::fmt::Result {
789                formatter.write_str("a map")
790            }
791
792            fn visit_map<M>(self, mut m: M) -> Result<Self::Value, M::Error>
793            where
794                M: MapAccess<'de>,
795            {
796                // `size_hint` is only advisory, so use it to reserve map capacity when large, when
797                // small just grow like normal
798                let size = m.size_hint().unwrap_or(0);
799                let mut map = if size < MAX_USEFUL_LINEAR_SCAN {
800                    AutoMap::with_hasher()
801                } else {
802                    AutoMap::Map(Box::new(HashMap::with_capacity_and_hasher(
803                        size,
804                        H::default(),
805                    )))
806                };
807                while let Some((k, v)) = m.next_entry()? {
808                    map.insert(k, v);
809                }
810                Ok(map)
811            }
812        }
813
814        deserializer.deserialize_map(AutoMapVisitor {
815            phantom: PhantomData::<AutoMap<K, V, H, I>>,
816        })
817    }
818}
819
820impl<K, V, H, const I: usize> Encode for AutoMap<K, V, H, I>
821where
822    K: Encode,
823    V: Encode,
824{
825    fn encode<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError> {
826        // Encode the length first
827        self.len().encode(encoder)?;
828        // Then encode each key-value tuple
829        for entry in self.iter() {
830            entry.encode(encoder)?;
831        }
832        Ok(())
833    }
834}
835
836impl<Context, K, V, H, const I: usize> Decode<Context> for AutoMap<K, V, H, I>
837where
838    K: Decode<Context> + Eq + Hash,
839    V: Decode<Context>,
840    H: BuildHasher + Default,
841{
842    fn decode<D: Decoder<Context = Context>>(decoder: &mut D) -> Result<Self, DecodeError> {
843        let len = usize::decode(decoder)?;
844        if len <= MAX_USEFUL_LINEAR_SCAN {
845            let mut list = TinyVec::with_capacity(len);
846            for _ in 0..len {
847                let entry = <(K, V)>::decode(decoder)?;
848                list.push(entry);
849            }
850            Ok(AutoMap::List(list))
851        } else {
852            let mut map = HashMap::with_capacity_and_hasher(len, H::default());
853            for _ in 0..len {
854                let (key, value) = <(K, V)>::decode(decoder)?;
855                map.insert(key, value);
856            }
857            Ok(AutoMap::Map(Box::new(map)))
858        }
859    }
860}
861
862impl_borrow_decode!(
863    AutoMap<K, V, H, I>,
864    K: Decode<__Context> + Eq + Hash,
865    V: Decode<__Context>,
866    H: BuildHasher + Default,
867    const I: usize,
868);
869
870impl<K: Eq + Hash, V: Eq, H: BuildHasher, const I: usize> PartialEq for AutoMap<K, V, H, I> {
871    fn eq(&self, other: &Self) -> bool {
872        match (self, other) {
873            (AutoMap::Map(a), AutoMap::Map(b)) => a == b,
874            (AutoMap::List(a), b) => {
875                if a.len() != b.len() {
876                    return false;
877                }
878                a.iter().all(|(k, v)| b.get(k) == Some(v))
879            }
880            (a, AutoMap::List(b)) => {
881                if a.len() != b.len() {
882                    return false;
883                }
884                b.iter().all(|(k, v)| a.get(k) == Some(v))
885            }
886        }
887    }
888}
889
890impl<K: Eq + Hash, V: Eq, H: BuildHasher, const I: usize> Eq for AutoMap<K, V, H, I>
891where
892    K: Eq,
893    V: Eq,
894{
895}
896
897impl<K: Eq + Hash, V: Eq + Hash, MH: BuildHasher + Default, const I: usize> Hash
898    for AutoMap<K, V, MH, I>
899{
900    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
901        // Hash the length first to distinguish maps of different sizes
902        self.len().hash(state);
903
904        // Use a commutative approach to ensure equal maps have equal hashes
905        // regardless of iteration order
906        let mut combined_hash = 0u64;
907
908        let hash_builder = MH::default();
909        for (k, v) in self.iter() {
910            use std::hash::Hasher;
911
912            // Hash each key-value pair independently
913            let mut entry_hasher = hash_builder.build_hasher();
914            k.hash(&mut entry_hasher);
915            v.hash(&mut entry_hasher);
916            let entry_hash = entry_hasher.finish();
917
918            // Combine using addition to make it order-independent (wrapping_add is commutative)
919            combined_hash = combined_hash.wrapping_add(entry_hash);
920        }
921
922        // Hash the combined result
923        combined_hash.hash(state);
924    }
925}
926
927impl<K, V, H, const I: usize> FromIterator<(K, V)> for AutoMap<K, V, H, I>
928where
929    K: Eq + Hash,
930    H: BuildHasher + Default,
931{
932    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
933        let iter = iter.into_iter();
934        let (lower, _) = iter.size_hint();
935        if lower > MAX_USEFUL_LINEAR_SCAN {
936            let map = iter.collect::<HashMap<K, V, H>>();
937            // The hint is not enforced
938            if map.len() < MIN_HASH_SIZE {
939                return AutoMap::List(map.into_iter().collect());
940            }
941            return AutoMap::Map(Box::new(map));
942        }
943        let mut map = AutoMap::with_hasher();
944        for (k, v) in iter {
945            map.insert(k, v);
946        }
947        map
948    }
949}
950
951pub enum ExtractIfIter<'l, K, V, const I: usize, F>
952where
953    F: for<'a, 'b> FnMut(&'a K, &'b mut V) -> bool,
954{
955    List {
956        list: &'l mut TinyVec<(K, V), I, MAX_USEFUL_LINEAR_SCAN>,
957        index: usize,
958        f: F,
959    },
960    Map(hashbrown::hash_map::ExtractIf<'l, K, V, F>),
961}
962
963impl<K, V, const I: usize, F> Iterator for ExtractIfIter<'_, K, V, I, F>
964where
965    F: for<'a, 'b> FnMut(&'a K, &'b mut V) -> bool,
966{
967    type Item = (K, V);
968
969    fn next(&mut self) -> Option<Self::Item> {
970        match self {
971            ExtractIfIter::List { list, index, f } => {
972                while *index < list.len() {
973                    let (key, value) = &mut list[*index];
974                    if f(key, value) {
975                        let item = list.swap_remove(*index);
976                        return Some(item);
977                    } else {
978                        *index += 1;
979                    }
980                }
981                None
982            }
983            ExtractIfIter::Map(extract_if) => extract_if.next(),
984        }
985    }
986}
987
988impl<K, V, H, const I: usize> ShrinkToFit for AutoMap<K, V, H, I>
989where
990    K: Eq + Hash,
991    V: Eq,
992    H: BuildHasher + Default,
993{
994    fn shrink_to_fit(&mut self) {
995        if self.len() < MIN_HASH_SIZE {
996            self.convert_to_list();
997        }
998
999        match self {
1000            AutoMap::List(list) => list.shrink_to_fit(),
1001            AutoMap::Map(map) => {
1002                hashbrown::HashMap::shrink_to_fit(map);
1003            }
1004        }
1005    }
1006}
1007#[cfg(test)]
1008mod tests {
1009    use super::*;
1010
1011    #[test]
1012    fn test_auto_map() {
1013        let mut map = AutoMap::new();
1014        for i in 0..MAX_USEFUL_LINEAR_SCAN * 2 {
1015            map.insert(i, i);
1016        }
1017        for i in 0..MAX_USEFUL_LINEAR_SCAN * 2 {
1018            assert_eq!(map.get(&i), Some(&i));
1019        }
1020        assert_eq!(map.get(&(MAX_USEFUL_LINEAR_SCAN * 2)), None);
1021        for i in 0..MAX_USEFUL_LINEAR_SCAN * 2 {
1022            assert_eq!(map.remove(&(MAX_USEFUL_LINEAR_SCAN * 2)), None);
1023            assert_eq!(map.remove(&i), Some(i));
1024        }
1025        assert_eq!(map.remove(&(MAX_USEFUL_LINEAR_SCAN * 2)), None);
1026    }
1027
1028    #[test]
1029    fn test_extract_if_map() {
1030        let mut map = AutoMap::new();
1031        for i in 0..MAX_USEFUL_LINEAR_SCAN * 2 {
1032            map.insert(i, i);
1033        }
1034        let iter = map.extract_if(|_, v| *v % 2 == 0);
1035        assert_eq!(iter.count(), MAX_USEFUL_LINEAR_SCAN);
1036        assert_eq!(map.len(), MAX_USEFUL_LINEAR_SCAN);
1037    }
1038
1039    #[test]
1040    fn test_extract_if_list() {
1041        let mut map = AutoMap::new();
1042        for i in 0..MIN_HASH_SIZE {
1043            map.insert(i, i);
1044        }
1045        let iter = map.extract_if(|_, v| *v % 2 == 0);
1046        assert_eq!(iter.count(), MIN_HASH_SIZE / 2);
1047        assert_eq!(map.len(), MIN_HASH_SIZE / 2);
1048    }
1049
1050    #[test]
1051    fn test_extract_if_list2() {
1052        let mut map = AutoMap::new();
1053        for i in 0..MIN_HASH_SIZE {
1054            map.insert(i, i);
1055        }
1056        let iter = map.extract_if(|_, v| *v < 5);
1057        assert_eq!(iter.count(), 5);
1058        assert_eq!(map.len(), MIN_HASH_SIZE - 5);
1059    }
1060}
auto_hash_map/map.rs

auto_hash_map/
map.rs
