small_map/
lib.rs

1#![allow(clippy::manual_map)]
2
3use core::{
4    hash::{BuildHasher, Hash},
5    hint::unreachable_unchecked,
6    iter::FusedIterator,
7};
8use std::{
9    collections::hash_map::RandomState,
10    fmt::{self, Debug},
11};
12
13use hashbrown::HashMap;
14use inline::Inline;
15
16#[macro_use]
17mod macros;
18mod inline;
19mod raw;
20
21pub use hashbrown::Equivalent;
22
23#[cfg(feature = "serde")]
24mod serde;
25
26#[cfg(feature = "fxhash")]
27pub type FxSmallMap<const N: usize, K, V> =
28    SmallMap<N, K, V, core::hash::BuildHasherDefault<rustc_hash::FxHasher>>;
29#[cfg(feature = "ahash")]
30pub type ASmallMap<const N: usize, K, V> =
31    SmallMap<N, K, V, core::hash::BuildHasherDefault<ahash::AHasher>>;
32
33#[derive(Clone)]
34pub enum SmallMap<const N: usize, K, V, S = RandomState> {
35    Heap(HashMap<K, V, S>),
36    Inline(Inline<N, K, V, S>),
37}
38
39impl<const N: usize, K, V, S> Debug for SmallMap<N, K, V, S>
40where
41    K: Debug,
42    V: Debug,
43{
44    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
45        f.debug_map().entries(self.iter()).finish()
46    }
47}
48
49impl<const N: usize, K, V, S: Default> Default for SmallMap<N, K, V, S> {
50    /// Creates an empty `SmallMap<N, K, V, S>`, with the `Default` value for the hasher.
51    ///
52    /// # Examples
53    ///
54    /// ```
55    /// use std::collections::hash_map::RandomState;
56    ///
57    /// use small_map::SmallMap;
58    ///
59    /// // You can specify all types of SmallMap, including N and hasher.
60    /// // Created map is empty and don't allocate memory
61    /// let map: SmallMap<8, u32, String> = SmallMap::default();
62    /// assert_eq!(map.capacity(), 8);
63    /// let map: SmallMap<8, u32, String, RandomState> = SmallMap::default();
64    /// assert_eq!(map.capacity(), 8);
65    /// ```
66    #[inline]
67    fn default() -> Self {
68        Self::with_hasher(Default::default())
69    }
70}
71
72impl<const N: usize, K, V, S: Default> SmallMap<N, K, V, S> {
73    /// Creates an empty `SmallMap`.
74    #[inline]
75    pub fn new() -> Self {
76        Self::with_hasher(Default::default())
77    }
78
79    /// Creates an empty `SmallMap` with the specified capacity.
80    ///
81    /// The hash map will be able to hold at least `capacity` elements without
82    /// reallocating. If `capacity` is smaller than N, the hash map will not allocate.
83    #[inline]
84    pub fn with_capacity(capacity: usize) -> Self {
85        if capacity > N {
86            Self::Heap(HashMap::with_capacity_and_hasher(
87                capacity,
88                Default::default(),
89            ))
90        } else {
91            Self::Inline(Inline::new(Default::default()))
92        }
93    }
94}
95
96impl<const N: usize, K, V, S> SmallMap<N, K, V, S> {
97    /// Creates an empty `SmallMap` which will use the given hash builder to hash
98    /// keys. It will be allocated with the given allocator.
99    ///
100    /// The hash map is initially created with a capacity of N, so it will not allocate until it
101    /// its size bigger than inline size N.
102    /// # Examples
103    ///
104    /// ```
105    /// use core::hash::BuildHasherDefault;
106    ///
107    /// use small_map::SmallMap;
108    ///
109    /// let s = BuildHasherDefault::<ahash::AHasher>::default();
110    /// let mut map = SmallMap::<8, _, _, _>::with_hasher(s);
111    /// map.insert(1, 2);
112    /// ```
113    #[inline]
114    pub const fn with_hasher(hash_builder: S) -> Self {
115        Self::Inline(Inline::new(hash_builder))
116    }
117
118    /// Creates an empty `SmallMap` with the specified capacity, using `hash_builder`
119    /// to hash the keys.
120    ///
121    /// The hash map will be able to hold at least `capacity` elements without
122    /// reallocating. If `capacity` is smaller than or eq to N, the hash map will not allocate.
123    #[inline]
124    pub fn with_capacity_and_hasher(capacity: usize, hash_builder: S) -> Self {
125        if capacity > N {
126            Self::Heap(HashMap::with_capacity_and_hasher(capacity, hash_builder))
127        } else {
128            Self::Inline(Inline::new(hash_builder))
129        }
130    }
131
132    /// What branch it is now.
133    #[inline]
134    pub fn is_inline(&self) -> bool {
135        matches!(self, Self::Inline(..))
136    }
137
138    /// Returns the number of elements the map can hold without reallocating.
139    ///
140    /// This number is a lower bound; the `SmallMap<N, K, V>` might be able to hold
141    /// more, but is guaranteed to be able to hold at least this many.
142    ///
143    /// # Examples
144    ///
145    /// ```
146    /// use small_map::SmallMap;
147    ///
148    /// let map: SmallMap<8, i32, i32> = SmallMap::with_capacity(100);
149    /// assert_eq!(map.len(), 0);
150    /// assert!(map.capacity() >= 100);
151    ///
152    /// let map: SmallMap<8, i32, i32> = SmallMap::with_capacity(2);
153    /// assert_eq!(map.len(), 0);
154    /// assert!(map.capacity() >= 8);
155    /// ```
156    #[inline]
157    pub fn capacity(&self) -> usize {
158        match self {
159            SmallMap::Heap(inner) => inner.capacity(),
160            SmallMap::Inline(_) => N,
161        }
162    }
163}
164
165impl<const N: usize, K, V, S> SmallMap<N, K, V, S>
166where
167    K: Eq + Hash,
168    S: BuildHasher,
169{
170    #[inline]
171    pub fn get<Q: ?Sized>(&self, k: &Q) -> Option<&V>
172    where
173        Q: Hash + Equivalent<K>,
174    {
175        match self {
176            SmallMap::Heap(inner) => inner.get(k),
177            SmallMap::Inline(inner) => inner.get(k),
178        }
179    }
180
181    /// Returns a mutable reference to the value corresponding to the key.
182    #[inline]
183    pub fn get_mut<Q: ?Sized>(&mut self, k: &Q) -> Option<&mut V>
184    where
185        Q: Hash + Equivalent<K>,
186    {
187        match self {
188            SmallMap::Heap(inner) => inner.get_mut(k),
189            SmallMap::Inline(inner) => inner.get_mut(k),
190        }
191    }
192
193    /// Returns the key-value pair corresponding to the supplied key.
194    #[inline]
195    pub fn get_key_value<Q: ?Sized>(&self, k: &Q) -> Option<(&K, &V)>
196    where
197        Q: Hash + Equivalent<K>,
198    {
199        match self {
200            SmallMap::Heap(inner) => inner.get_key_value(k),
201            SmallMap::Inline(inner) => inner.get_key_value(k),
202        }
203    }
204
205    #[inline]
206    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
207        match self {
208            SmallMap::Heap(inner) => inner.insert(key, value),
209            SmallMap::Inline(inner) => {
210                if raw::util::likely(!inner.is_full()) {
211                    return inner.insert(key, value);
212                }
213                let heap = HashMap::with_capacity_and_hasher(N + 1, unsafe { inner.take_hasher() });
214
215                let heap = match core::mem::replace(self, SmallMap::Heap(heap)) {
216                    SmallMap::Heap(_) => unsafe { unreachable_unchecked() },
217                    SmallMap::Inline(inline) => {
218                        let heap = match self {
219                            SmallMap::Heap(heap) => heap,
220                            SmallMap::Inline(_) => unsafe { unreachable_unchecked() },
221                        };
222                        for (k, v) in inline.into_iter() {
223                            heap.insert_unique_unchecked(k, v);
224                        }
225                        heap
226                    }
227                };
228                heap.insert(key, value)
229            }
230        }
231    }
232
233    #[inline]
234    pub fn remove<Q: ?Sized>(&mut self, k: &Q) -> Option<V>
235    where
236        Q: Hash + Equivalent<K>,
237    {
238        // Avoid `Option::map` because it bloats LLVM IR.
239        match self.remove_entry(k) {
240            Some((_, v)) => Some(v),
241            None => None,
242        }
243    }
244
245    #[inline]
246    pub fn remove_entry<Q: ?Sized>(&mut self, k: &Q) -> Option<(K, V)>
247    where
248        Q: Hash + Equivalent<K>,
249    {
250        match self {
251            SmallMap::Heap(inner) => inner.remove_entry(k),
252            SmallMap::Inline(inner) => inner.remove_entry(k),
253        }
254    }
255}
256
257pub enum Iter<'a, const N: usize, K, V> {
258    Heap(hashbrown::hash_map::Iter<'a, K, V>),
259    Inline(inline::Iter<'a, N, K, V>),
260}
261
262impl<'a, const N: usize, K, V> Iterator for Iter<'a, N, K, V> {
263    type Item = (&'a K, &'a V);
264
265    #[inline]
266    fn next(&mut self) -> Option<(&'a K, &'a V)> {
267        match self {
268            Iter::Heap(inner) => inner.next(),
269            Iter::Inline(inner) => inner.next(),
270        }
271    }
272    #[inline]
273    fn size_hint(&self) -> (usize, Option<usize>) {
274        match self {
275            Iter::Heap(inner) => inner.size_hint(),
276            Iter::Inline(inner) => inner.size_hint(),
277        }
278    }
279}
280
281impl<const N: usize, K, V> ExactSizeIterator for Iter<'_, N, K, V> {
282    #[inline]
283    fn len(&self) -> usize {
284        match self {
285            Iter::Heap(inner) => inner.len(),
286            Iter::Inline(inner) => inner.len(),
287        }
288    }
289}
290
291pub enum IntoIter<const N: usize, K, V> {
292    Heap(hashbrown::hash_map::IntoIter<K, V>),
293    Inline(inline::IntoIter<N, K, V>),
294}
295
296impl<const N: usize, K, V> Iterator for IntoIter<N, K, V> {
297    type Item = (K, V);
298
299    #[inline]
300    fn next(&mut self) -> Option<(K, V)> {
301        match self {
302            IntoIter::Heap(inner) => inner.next(),
303            IntoIter::Inline(inner) => inner.next(),
304        }
305    }
306    #[inline]
307    fn size_hint(&self) -> (usize, Option<usize>) {
308        match self {
309            IntoIter::Heap(inner) => inner.size_hint(),
310            IntoIter::Inline(inner) => inner.size_hint(),
311        }
312    }
313}
314impl<const N: usize, K, V> ExactSizeIterator for IntoIter<N, K, V> {
315    #[inline]
316    fn len(&self) -> usize {
317        match self {
318            IntoIter::Heap(inner) => inner.len(),
319            IntoIter::Inline(inner) => inner.len(),
320        }
321    }
322}
323impl<const N: usize, K, V> FusedIterator for IntoIter<N, K, V> {}
324
325impl<const N: usize, K, V, S> IntoIterator for SmallMap<N, K, V, S> {
326    type Item = (K, V);
327    type IntoIter = IntoIter<N, K, V>;
328
329    /// Creates a consuming iterator, that is, one that moves each key-value
330    /// pair out of the map in arbitrary order. The map cannot be used after
331    /// calling this.
332    #[inline]
333    fn into_iter(self) -> IntoIter<N, K, V> {
334        match self {
335            SmallMap::Heap(inner) => IntoIter::Heap(inner.into_iter()),
336            SmallMap::Inline(inner) => IntoIter::Inline(inner.into_iter()),
337        }
338    }
339}
340
341impl<'a, const N: usize, K, V, S> IntoIterator for &'a SmallMap<N, K, V, S> {
342    type Item = (&'a K, &'a V);
343    type IntoIter = Iter<'a, N, K, V>;
344
345    /// Creates an iterator over the entries of a `HashMap` in arbitrary order.
346    /// The iterator element type is `(&'a K, &'a V)`.
347    #[inline]
348    fn into_iter(self) -> Iter<'a, N, K, V> {
349        match self {
350            SmallMap::Heap(inner) => Iter::Heap(inner.iter()),
351            SmallMap::Inline(inner) => Iter::Inline(inner.iter()),
352        }
353    }
354}
355
356impl<const N: usize, K, V, S> SmallMap<N, K, V, S> {
357    #[inline]
358    pub fn len(&self) -> usize {
359        match self {
360            SmallMap::Heap(inner) => inner.len(),
361            SmallMap::Inline(inner) => inner.len(),
362        }
363    }
364
365    #[inline]
366    pub fn is_empty(&self) -> bool {
367        match self {
368            SmallMap::Heap(inner) => inner.is_empty(),
369            SmallMap::Inline(inner) => inner.is_empty(),
370        }
371    }
372
373    #[inline]
374    pub fn iter(&self) -> Iter<'_, N, K, V> {
375        match self {
376            SmallMap::Heap(inner) => Iter::Heap(inner.iter()),
377            SmallMap::Inline(inner) => Iter::Inline(inner.iter()),
378        }
379    }
380}
381
382#[cfg(test)]
383mod tests {
384    use std::{cell::RefCell, rc::Rc};
385
386    use ahash::RandomState;
387    use hashbrown::HashMap;
388    use rand::Rng;
389
390    use crate::SmallMap;
391
392    #[test]
393    fn basic_op() {
394        let mut map = SmallMap::<16, String, String>::default();
395        map.insert("hello".to_string(), "world".to_string());
396        assert_eq!(map.len(), 1);
397        assert_eq!(map.get("hello").unwrap(), "world");
398        map.insert("hello2".to_string(), "world2".to_string());
399        assert_eq!(map.get("hello2").unwrap(), "world2");
400        assert_eq!(map.len(), 2);
401
402        assert_eq!(
403            map.remove_entry("hello").unwrap(),
404            ("hello".to_string(), "world".to_string())
405        );
406        assert_eq!(map.len(), 1);
407        assert_eq!(map.get("hello2").unwrap(), "world2");
408        assert_eq!(map.remove("hello2").unwrap(), "world2".to_string());
409        assert_eq!(map.len(), 0);
410        assert!(map.get("hello").is_none());
411    }
412
413    #[test]
414    fn multi_group_with_padding() {
415        let mut map = SmallMap::<33, i32, i32>::default();
416        for i in 0..33 {
417            map.insert(i, i * 2);
418        }
419        for i in 0..33 {
420            assert_eq!(*map.get(&i).unwrap(), i * 2);
421        }
422        assert!(map.is_inline());
423
424        for i in 33..64 {
425            map.insert(i, i * 2);
426        }
427        assert!(!map.is_inline());
428        for i in 0..64 {
429            assert_eq!(*map.get(&i).unwrap(), i * 2);
430        }
431    }
432
433    #[test]
434    fn fuzzing() {
435        let mut smallmap = SmallMap::<16, i32, i32>::default();
436        let mut hashmap = HashMap::<i32, i32, RandomState>::default();
437        for _ in 0..1000000 {
438            let op = Operation::random();
439            op.exec(&mut smallmap, &mut hashmap);
440        }
441
442        enum Operation {
443            Insert(i32, i32),
444            Remove(i32),
445            Get(i32),
446            ModifyIfExist(i32, i32),
447        }
448        impl Operation {
449            fn random() -> Self {
450                let mut rng = rand::thread_rng();
451
452                let choice: u8 = rng.gen();
453                match choice % 4 {
454                    0 => Operation::Insert(rng.gen_range(0..32), rng.gen()),
455                    1 => Operation::Remove(rng.gen_range(0..32)),
456                    2 => Operation::Get(rng.gen_range(0..32)),
457                    3 => Operation::ModifyIfExist(rng.gen_range(0..32), rng.gen()),
458                    _ => unreachable!(),
459                }
460            }
461
462            fn exec<const N: usize, S1: core::hash::BuildHasher, S2: core::hash::BuildHasher>(
463                self,
464                sm: &mut SmallMap<N, i32, i32, S1>,
465                hm: &mut HashMap<i32, i32, S2>,
466            ) {
467                match self {
468                    Operation::Insert(k, v) => {
469                        if sm.len() == sm.capacity() {
470                            return;
471                        }
472                        assert_eq!(sm.insert(k, v), hm.insert(k, v));
473                        assert!(sm.is_inline());
474                    }
475                    Operation::Remove(k) => {
476                        assert_eq!(sm.remove(&k), hm.remove(&k));
477                    }
478                    Operation::Get(k) => {
479                        assert_eq!(sm.get(&k), hm.get(&k));
480                    }
481                    Operation::ModifyIfExist(k, nv) => {
482                        let (sv, hv) = (sm.get_mut(&k), hm.get_mut(&k));
483                        assert_eq!(sv, hv);
484                        if let Some(v) = sv {
485                            *v = nv;
486                        }
487                        if let Some(v) = hv {
488                            *v = nv;
489                        }
490                    }
491                }
492            }
493        }
494    }
495
496    #[test]
497    fn drop_chk() {
498        let (probe1, checker1) = drop_checker();
499        let (probe2, checker2) = drop_checker();
500        let (probe3, checker3) = drop_checker();
501        let mut map = SmallMap::<8, _, _>::default();
502        map.insert(1, probe1);
503        map.insert(2, probe2);
504        map.insert(3, probe3);
505        assert_eq!(map.len(), 3);
506        let mut it = map.into_iter();
507        drop(it.next());
508        drop(it);
509        checker1.assert_drop();
510        checker2.assert_drop();
511        checker3.assert_drop();
512
513        fn drop_checker() -> (DropProbe, DropChecker) {
514            let flag = Rc::new(RefCell::new(false));
515            (DropProbe { flag: flag.clone() }, DropChecker { flag })
516        }
517
518        struct DropChecker {
519            flag: Rc<RefCell<bool>>,
520        }
521
522        impl DropChecker {
523            fn assert_drop(self) {
524                assert!(*self.flag.borrow())
525            }
526        }
527
528        struct DropProbe {
529            flag: Rc<RefCell<bool>>,
530        }
531
532        impl Drop for DropProbe {
533            fn drop(&mut self) {
534                *self.flag.borrow_mut() = true;
535            }
536        }
537    }
538}
small_map/lib.rs

small_map/
lib.rs
