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;
23use smallvec::SmallVec;
24
25use crate::{MAX_LIST_SIZE, MIN_HASH_SIZE};
26
27#[derive(Clone)]
28pub enum AutoMap<K, V, H = BuildHasherDefault<FxHasher>, const I: usize = 0> {
29    List(SmallVec<[(K, V); I]>),
30    Map(Box<HashMap<K, V, H>>),
31}
32
33impl<K, V, H, const I: usize> Default for AutoMap<K, V, H, I> {
34    fn default() -> Self {
35        Self::List(Default::default())
36    }
37}
38
39impl<K: Debug, V: Debug, H, const I: usize> Debug for AutoMap<K, V, H, I> {
40    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        f.debug_map().entries(self.iter()).finish()
42    }
43}
44
45impl<K, V> AutoMap<K, V, BuildHasherDefault<FxHasher>, 0> {
46    /// see [HashMap::new](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.new)
47    pub const fn new() -> Self {
48        AutoMap::List(SmallVec::new_const())
49    }
50
51    /// see [HashMap::with_capacity](https://doc.rust-lang.org/std/collections/struct.HashMap.html#method.with_capacity)
52    pub fn with_capacity(capacity: usize) -> Self {
53        if capacity < MAX_LIST_SIZE {
54            AutoMap::List(SmallVec::with_capacity(capacity))
55        } else {
56            AutoMap::Map(Box::new(HashMap::with_capacity_and_hasher(
57                capacity,
58                Default::default(),
59            )))
60        }
61    }
62}
63
64impl<K, V, H, const I: usize> AutoMap<K, V, H, I> {
65    /// see [HashMap::with_hasher](https://doc.rust-lang.org/std/collections/hash_map/struct.HashMap.html#method.with_hasher)
66    pub const fn with_hasher() -> Self {
67        AutoMap::List(SmallVec::new_const())
68    }
69
70    /// see [HashMap::with_capacity_and_hasher](https://doc.rust-lang.org/std/collections/hash_map/struct.HashMap.html#method.with_capacity_and_hasher)
71    pub fn with_capacity_and_hasher(capacity: usize, hasher: H) -> Self {
72        if capacity <= MAX_LIST_SIZE {
73            AutoMap::List(SmallVec::with_capacity(capacity))
74        } else {
75            AutoMap::Map(Box::new(HashMap::with_capacity_and_hasher(
76                capacity, hasher,
77            )))
78        }
79    }
80
81    /// see [HashMap::clear](https://doc.rust-lang.org/std/collections/hash_map/struct.HashMap.html#method.clear)
82    pub fn clear(&mut self) {
83        match self {
84            AutoMap::List(list) => list.clear(),
85            AutoMap::Map(map) => map.clear(),
86        }
87    }
88}
89
90impl<K: Eq + Hash, V, H: BuildHasher + Default, const I: usize> AutoMap<K, V, H, I> {
91    fn convert_to_map(&mut self) -> &mut HashMap<K, V, H> {
92        if let AutoMap::List(list) = self {
93            let mut map = HashMap::with_capacity_and_hasher(MAX_LIST_SIZE * 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 SmallVec<[(K, V); I]> {
105        if let AutoMap::Map(map) = self {
106            let mut list = SmallVec::with_capacity(MAX_LIST_SIZE);
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_LIST_SIZE {
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_LIST_SIZE {
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_LIST_SIZE {
260                    let mut list = SmallVec::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                if list.capacity() > list.len() * 3 {
274                    // You need to call "grow" to shrink a SmallVec to a specific capacity
275                    // There is no Vec::shrink_to() equivalent for SmallVec
276                    list.grow(list.len().next_power_of_two());
277                }
278            }
279            AutoMap::Map(map) => {
280                if map.len() <= MIN_HASH_SIZE {
281                    let mut list = SmallVec::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(smallvec::IntoIter<[(K, V); I]>),
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(smallvec::IntoIter<[(K, V); I]>),
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 SmallVec<[(K, V); I]>,
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 SmallVec<[(K, V); I]>,
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_LIST_SIZE {
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 SmallVec<[(K, V); I]>,
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 SmallVec<[(K, V); I]>,
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_LIST_SIZE {
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                if let Some(size) = m.size_hint() {
797                    if size < MAX_LIST_SIZE {
798                        let mut list = SmallVec::with_capacity(size);
799                        while let Some((k, v)) = m.next_entry()? {
800                            list.push((k, v));
801                        }
802                        return Ok(AutoMap::List(list));
803                    } else {
804                        let mut map =
805                            Box::new(HashMap::with_capacity_and_hasher(size, H::default()));
806                        while let Some((k, v)) = m.next_entry()? {
807                            map.insert(k, v);
808                        }
809                        return Ok(AutoMap::Map(map));
810                    }
811                }
812                let mut map = AutoMap::with_hasher();
813                while let Some((k, v)) = m.next_entry()? {
814                    map.insert(k, v);
815                }
816                Ok(map)
817            }
818        }
819
820        deserializer.deserialize_map(AutoMapVisitor {
821            phantom: PhantomData::<AutoMap<K, V, H, I>>,
822        })
823    }
824}
825
826impl<K, V, H, const I: usize> Encode for AutoMap<K, V, H, I>
827where
828    K: Encode,
829    V: Encode,
830{
831    fn encode<E: Encoder>(&self, encoder: &mut E) -> Result<(), EncodeError> {
832        // Encode the length first
833        self.len().encode(encoder)?;
834        // Then encode each key-value tuple
835        for entry in self.iter() {
836            entry.encode(encoder)?;
837        }
838        Ok(())
839    }
840}
841
842impl<Context, K, V, H, const I: usize> Decode<Context> for AutoMap<K, V, H, I>
843where
844    K: Decode<Context> + Eq + Hash,
845    V: Decode<Context>,
846    H: BuildHasher + Default,
847{
848    fn decode<D: Decoder<Context = Context>>(decoder: &mut D) -> Result<Self, DecodeError> {
849        let len = usize::decode(decoder)?;
850        if len <= MAX_LIST_SIZE {
851            let mut list = SmallVec::with_capacity(len);
852            for _ in 0..len {
853                let entry = <(K, V)>::decode(decoder)?;
854                list.push(entry);
855            }
856            Ok(AutoMap::List(list))
857        } else {
858            let mut map = HashMap::with_capacity_and_hasher(len, H::default());
859            for _ in 0..len {
860                let (key, value) = <(K, V)>::decode(decoder)?;
861                map.insert(key, value);
862            }
863            Ok(AutoMap::Map(Box::new(map)))
864        }
865    }
866}
867
868impl_borrow_decode!(
869    AutoMap<K, V, H, I>,
870    K: Decode<__Context> + Eq + Hash,
871    V: Decode<__Context>,
872    H: BuildHasher + Default,
873    const I: usize,
874);
875
876impl<K: Eq + Hash, V: Eq, H: BuildHasher, const I: usize> PartialEq for AutoMap<K, V, H, I> {
877    fn eq(&self, other: &Self) -> bool {
878        match (self, other) {
879            (AutoMap::Map(a), AutoMap::Map(b)) => a == b,
880            (AutoMap::List(a), b) => {
881                if a.len() != b.len() {
882                    return false;
883                }
884                a.iter().all(|(k, v)| b.get(k) == Some(v))
885            }
886            (a, AutoMap::List(b)) => {
887                if a.len() != b.len() {
888                    return false;
889                }
890                b.iter().all(|(k, v)| a.get(k) == Some(v))
891            }
892        }
893    }
894}
895
896impl<K: Eq + Hash, V: Eq, H: BuildHasher, const I: usize> Eq for AutoMap<K, V, H, I>
897where
898    K: Eq,
899    V: Eq,
900{
901}
902
903impl<K: Eq + Hash, V: Eq + Hash, MH: BuildHasher + Default, const I: usize> Hash
904    for AutoMap<K, V, MH, I>
905{
906    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
907        // Hash the length first to distinguish maps of different sizes
908        self.len().hash(state);
909
910        // Use a commutative approach to ensure equal maps have equal hashes
911        // regardless of iteration order
912        let mut combined_hash = 0u64;
913
914        let hash_builder = MH::default();
915        for (k, v) in self.iter() {
916            use std::hash::Hasher;
917
918            // Hash each key-value pair independently
919            let mut entry_hasher = hash_builder.build_hasher();
920            k.hash(&mut entry_hasher);
921            v.hash(&mut entry_hasher);
922            let entry_hash = entry_hasher.finish();
923
924            // Combine using addition to make it order-independent (wrapping_add is commutative)
925            combined_hash = combined_hash.wrapping_add(entry_hash);
926        }
927
928        // Hash the combined result
929        combined_hash.hash(state);
930    }
931}
932
933impl<K, V, H, const I: usize> FromIterator<(K, V)> for AutoMap<K, V, H, I>
934where
935    K: Eq + Hash,
936    H: BuildHasher + Default,
937{
938    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
939        let iter = iter.into_iter();
940        let (lower, _) = iter.size_hint();
941        if lower > MAX_LIST_SIZE {
942            let map = iter.collect::<HashMap<K, V, H>>();
943            // The hint is not enforced
944            if map.len() < MIN_HASH_SIZE {
945                return AutoMap::List(map.into_iter().collect());
946            }
947            return AutoMap::Map(Box::new(map));
948        }
949        let mut map = AutoMap::with_hasher();
950        for (k, v) in iter {
951            map.insert(k, v);
952        }
953        map
954    }
955}
956
957pub enum ExtractIfIter<'l, K, V, const I: usize, F>
958where
959    F: for<'a, 'b> FnMut(&'a K, &'b mut V) -> bool,
960{
961    List {
962        list: &'l mut SmallVec<[(K, V); I]>,
963        index: usize,
964        f: F,
965    },
966    Map(hashbrown::hash_map::ExtractIf<'l, K, V, F>),
967}
968
969impl<K, V, const I: usize, F> Iterator for ExtractIfIter<'_, K, V, I, F>
970where
971    F: for<'a, 'b> FnMut(&'a K, &'b mut V) -> bool,
972{
973    type Item = (K, V);
974
975    fn next(&mut self) -> Option<Self::Item> {
976        match self {
977            ExtractIfIter::List { list, index, f } => {
978                while *index < list.len() {
979                    let (key, value) = &mut list[*index];
980                    if f(key, value) {
981                        let item = list.swap_remove(*index);
982                        return Some(item);
983                    } else {
984                        *index += 1;
985                    }
986                }
987                None
988            }
989            ExtractIfIter::Map(extract_if) => extract_if.next(),
990        }
991    }
992}
993
994impl<K, V, H, const I: usize> ShrinkToFit for AutoMap<K, V, H, I>
995where
996    K: Eq + Hash,
997    V: Eq,
998    H: BuildHasher + Default,
999{
1000    fn shrink_to_fit(&mut self) {
1001        if self.len() < MIN_HASH_SIZE {
1002            self.convert_to_list();
1003        }
1004
1005        match self {
1006            AutoMap::List(list) => list.shrink_to_fit(),
1007            AutoMap::Map(map) => {
1008                hashbrown::HashMap::shrink_to_fit(map);
1009            }
1010        }
1011    }
1012}
1013#[cfg(test)]
1014mod tests {
1015    use super::*;
1016
1017    #[test]
1018    fn test_auto_map() {
1019        let mut map = AutoMap::new();
1020        for i in 0..MAX_LIST_SIZE * 2 {
1021            map.insert(i, i);
1022        }
1023        for i in 0..MAX_LIST_SIZE * 2 {
1024            assert_eq!(map.get(&i), Some(&i));
1025        }
1026        assert_eq!(map.get(&(MAX_LIST_SIZE * 2)), None);
1027        for i in 0..MAX_LIST_SIZE * 2 {
1028            assert_eq!(map.remove(&(MAX_LIST_SIZE * 2)), None);
1029            assert_eq!(map.remove(&i), Some(i));
1030        }
1031        assert_eq!(map.remove(&(MAX_LIST_SIZE * 2)), None);
1032    }
1033
1034    #[test]
1035    fn test_extract_if_map() {
1036        let mut map = AutoMap::new();
1037        for i in 0..MAX_LIST_SIZE * 2 {
1038            map.insert(i, i);
1039        }
1040        let iter = map.extract_if(|_, v| *v % 2 == 0);
1041        assert_eq!(iter.count(), MAX_LIST_SIZE);
1042        assert_eq!(map.len(), MAX_LIST_SIZE);
1043    }
1044
1045    #[test]
1046    fn test_extract_if_list() {
1047        let mut map = AutoMap::new();
1048        for i in 0..MIN_HASH_SIZE {
1049            map.insert(i, i);
1050        }
1051        let iter = map.extract_if(|_, v| *v % 2 == 0);
1052        assert_eq!(iter.count(), MIN_HASH_SIZE / 2);
1053        assert_eq!(map.len(), MIN_HASH_SIZE / 2);
1054    }
1055
1056    #[test]
1057    fn test_extract_if_list2() {
1058        let mut map = AutoMap::new();
1059        for i in 0..MIN_HASH_SIZE {
1060            map.insert(i, i);
1061        }
1062        let iter = map.extract_if(|_, v| *v < 5);
1063        assert_eq!(iter.count(), 5);
1064        assert_eq!(map.len(), MIN_HASH_SIZE - 5);
1065    }
1066}
auto_hash_map/map.rs

auto_hash_map/
map.rs
