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

auto_hash_map/
map.rs
