tl/inline/
hashmap.rs

1use std::fmt::{Debug, Formatter};
2use std::hash::Hash;
3use std::ptr;
4use std::{collections::HashMap, mem::MaybeUninit};
5
6/// Similar to InlineVec, this structure will use an array
7/// if it is small enough to live on the stack, otherwise
8/// it allocates a HashMap on the heap
9///
10/// Hashing can be slower than just iterating through an array
11/// if the array is small, which is where it makes most sense
12#[derive(Debug, Clone)]
13pub struct InlineHashMap<K, V, const N: usize>(InlineHashMapInner<K, V, N>);
14
15impl<K, V, const N: usize> InlineHashMap<K, V, N>
16where
17    K: Hash + Eq,
18{
19    /// Creates a new InlineHashMap
20    pub(crate) fn new() -> Self {
21        Self(InlineHashMapInner::new())
22    }
23
24    /// Returns the number of elements in the map
25    #[inline]
26    pub fn len(&self) -> usize {
27        self.0.len()
28    }
29
30    /// Returns true if the map contains no elements
31    #[inline]
32    pub fn is_empty(&self) -> bool {
33        self.len() == 0
34    }
35
36    /// Returns an iterator over the elements of this map
37    ///
38    /// This function boxes the returned iterator because it can be either of two:
39    /// - The iterator returned by `HashMap::iter()`
40    /// - The iterator over a stack-allocated array
41    #[inline]
42    pub fn iter(&self) -> Box<dyn Iterator<Item = (&K, &V)> + '_> {
43        self.0.iter()
44    }
45
46    /// Copies `self` into a new `HashMap<K, V>`
47    #[inline]
48    pub fn to_map(&self) -> HashMap<K, V>
49    where
50        K: Clone + Hash + Eq,
51        V: Clone,
52    {
53        self.0.to_map()
54    }
55
56    /// Checks whether this vector is allocated on the heap
57    #[inline]
58    pub fn is_heap_allocated(&self) -> bool {
59        self.0.is_heap_allocated()
60    }
61
62    /// Inserts a new element into the map.
63    ///
64    /// If an equal key is already present, the stored key is retained and its value is replaced.
65    /// The map's length does not change.
66    #[inline]
67    pub fn insert(&mut self, key: K, value: V) {
68        self.0.insert(key, value)
69    }
70
71    /// Removes an element from the map, and returns ownership over the value
72    #[inline]
73    pub fn remove(&mut self, key: &K) -> Option<V> {
74        self.0.remove(key)
75    }
76
77    /// Returns a reference to the value corresponding to the key.
78    #[inline]
79    pub fn get(&self, key: &K) -> Option<&V> {
80        self.0.get(key)
81    }
82
83    /// Returns a mutable reference to the value corresponding to the key.
84    #[inline]
85    pub fn get_mut(&mut self, key: &K) -> Option<&mut V> {
86        self.0.get_mut(key)
87    }
88
89    /// Checks whether the map contains a value for the specified key.
90    #[inline]
91    pub fn contains_key(&self, key: &K) -> bool {
92        self.0.contains_key(key)
93    }
94}
95
96enum InlineHashMapInner<K, V, const N: usize> {
97    Inline {
98        len: usize,
99        data: [MaybeUninit<(K, V)>; N],
100    },
101    Heap(HashMap<K, V>),
102}
103
104impl<K, V, const N: usize> Debug for InlineHashMapInner<K, V, N>
105where
106    K: Debug,
107    V: Debug,
108{
109    fn fmt(&self, f: &mut Formatter<'_>) -> Result<(), std::fmt::Error> {
110        write!(f, "InlineHashMap<{} items>", self.len())
111    }
112}
113
114impl<K, V, const N: usize> Clone for InlineHashMapInner<K, V, N>
115where
116    K: Clone,
117    V: Clone,
118{
119    fn clone(&self) -> Self {
120        match self {
121            Self::Heap(m) => Self::Heap(m.clone()),
122            Self::Inline { len, data } => {
123                let mut new_data = super::uninit_array();
124
125                let iter = data.iter().take(*len).enumerate();
126
127                for (idx, element) in iter {
128                    let element = unsafe { &*element.as_ptr() };
129                    let (key, value) = element.clone();
130                    new_data[idx] = MaybeUninit::new((key, value));
131                }
132
133                Self::Inline {
134                    len: *len,
135                    data: new_data,
136                }
137            }
138        }
139    }
140}
141
142impl<K, V, const N: usize> Drop for InlineHashMapInner<K, V, N> {
143    fn drop(&mut self) {
144        if let Self::Inline { len, data } = self {
145            for element in data.iter_mut().take(*len) {
146                unsafe { ptr::drop_in_place(element.as_mut_ptr()) };
147            }
148        }
149    }
150}
151
152impl<K, V, const N: usize> InlineHashMapInner<K, V, N> {
153    #[inline]
154    pub(crate) fn new() -> Self {
155        Self::Inline {
156            len: 0,
157            data: super::uninit_array(),
158        }
159    }
160
161    #[inline]
162    pub fn iter(&self) -> Box<dyn Iterator<Item = (&K, &V)> + '_> {
163        match self {
164            Self::Inline { len, data } => {
165                Box::new(unsafe { InlineHashMapIterator::new(data, *len) })
166            }
167            Self::Heap(h) => Box::new(h.iter()),
168        }
169    }
170
171    #[inline]
172    fn to_map(&self) -> HashMap<K, V>
173    where
174        K: Clone + Hash + Eq,
175        V: Clone,
176    {
177        match &self {
178            InlineHashMapInner::Heap(m) => m.clone(),
179            InlineHashMapInner::Inline { len, data } => {
180                let mut new_data = HashMap::with_capacity(*len);
181
182                let iter = data.iter().take(*len);
183
184                for element in iter {
185                    let element = unsafe { &*element.as_ptr() };
186                    let (key, value) = element.clone();
187                    new_data.insert(key, value);
188                }
189
190                new_data
191            }
192        }
193    }
194
195    #[inline]
196    pub fn len(&self) -> usize {
197        match self {
198            Self::Inline { len, .. } => *len,
199            Self::Heap(map) => map.len(),
200        }
201    }
202
203    #[inline]
204    pub fn is_heap_allocated(&self) -> bool {
205        matches!(self, Self::Heap(_))
206    }
207}
208
209impl<K: Eq + Hash, V, const N: usize> InlineHashMapInner<K, V, N> {
210    pub fn get<'m>(&'m self, k: &K) -> Option<&'m V> {
211        match self {
212            Self::Inline { data, len } => unsafe {
213                InlineHashMapIterator::new(data, *len)
214                    .find(|(key, _)| key.eq(&k))
215                    .map(|(_, value)| value)
216            },
217            Self::Heap(map) => map.get(k),
218        }
219    }
220
221    pub fn get_mut<'m>(&'m mut self, k: &K) -> Option<&'m mut V> {
222        match self {
223            Self::Inline { data, len } => unsafe {
224                InlineHashMapIteratorMut::new(data, *len)
225                    .find(|(key, _)| key.eq(k))
226                    .map(|(_, value)| value)
227            },
228            Self::Heap(map) => map.get_mut(k),
229        }
230    }
231
232    pub fn remove(&mut self, key: &K) -> Option<V> {
233        match self {
234            Self::Inline { data, len } => {
235                let idx = data
236                    .iter()
237                    .take(*len)
238                    .map(|x| unsafe { &*x.as_ptr() })
239                    .position(|x| &x.0 == key)?;
240
241                let element = unsafe {
242                    std::mem::replace(data.get_unchecked_mut(idx), MaybeUninit::uninit())
243                };
244
245                // HashMap order is not guaranteed, so instead of swapping every item like we do with InlineVec,
246                // we can simply swap the last item with the one we want to remove.
247                data.swap(idx, *len - 1);
248                *len -= 1;
249
250                Some(unsafe { element.assume_init().1 })
251            }
252            Self::Heap(h) => h.remove(key),
253        }
254    }
255
256    pub fn insert(&mut self, k: K, v: V) {
257        let (array, len) = match self {
258            Self::Inline { data, len } => (data, len),
259            Self::Heap(map) => {
260                map.insert(k, v);
261                return;
262            }
263        };
264
265        for element in array.iter_mut().take(*len) {
266            // SAFETY: The first `len` elements of the inline array are initialized.
267            let (key, value) = unsafe { element.assume_init_mut() };
268            if (*key).eq(&k) {
269                let old_value = std::mem::replace(value, v);
270                drop(old_value);
271                return;
272            }
273        }
274
275        if *len >= N {
276            let mut map = HashMap::with_capacity(*len + 1);
277
278            // Move old elements to the heap from the end of the initialized prefix. Decrementing
279            // the length first keeps the inline representation valid if hashing or equality panics.
280            while *len != 0 {
281                *len -= 1;
282
283                // SAFETY: The element at the decremented length was initialized, and lowering the
284                // length transfers responsibility for dropping it to the local map.
285                let (key, value) = unsafe { array[*len].assume_init_read() };
286                map.insert(key, value);
287            }
288
289            map.insert(k, v);
290            *self = Self::Heap(map);
291        } else {
292            array[*len].write((k, v));
293            *len += 1;
294        }
295    }
296
297    pub fn contains_key(&self, k: &K) -> bool {
298        match self {
299            Self::Inline { data, len } => unsafe {
300                InlineHashMapIterator::new(data, *len).any(|(key, _)| key.eq(k))
301            },
302            Self::Heap(map) => map.contains_key(k),
303        }
304    }
305}
306
307/// An iterator over the inline array elements of an `InlineHashMap`.
308pub struct InlineHashMapIteratorMut<'a, K, V> {
309    array: &'a mut [MaybeUninit<(K, V)>],
310    idx: usize,
311    len: usize,
312}
313
314impl<'a, K, V> InlineHashMapIteratorMut<'a, K, V> {
315    pub(crate) unsafe fn new(array: &'a mut [MaybeUninit<(K, V)>], len: usize) -> Self {
316        Self { array, idx: 0, len }
317    }
318}
319
320impl<'a, K, V> Iterator for InlineHashMapIteratorMut<'a, K, V> {
321    type Item = &'a mut (K, V);
322
323    fn next(&mut self) -> Option<Self::Item> {
324        if self.idx >= self.len {
325            return None;
326        }
327
328        let element = unsafe { &mut *self.array[self.idx].as_mut_ptr() };
329        self.idx += 1;
330
331        Some(element)
332    }
333}
334
335/// An iterator over the inline array elements of an `InlineHashMap`.
336pub struct InlineHashMapIterator<'a, K, V> {
337    array: &'a [MaybeUninit<(K, V)>],
338    idx: usize,
339    len: usize,
340}
341
342impl<'a, K, V> InlineHashMapIterator<'a, K, V> {
343    pub(crate) unsafe fn new(array: &'a [MaybeUninit<(K, V)>], len: usize) -> Self {
344        Self { array, idx: 0, len }
345    }
346}
347
348impl<'a, K, V> Iterator for InlineHashMapIterator<'a, K, V> {
349    type Item = (&'a K, &'a V);
350
351    fn next(&mut self) -> Option<Self::Item> {
352        if self.idx >= self.len {
353            return None;
354        }
355
356        let (k, v) = unsafe { &*self.array[self.idx].as_ptr() };
357        self.idx += 1;
358
359        Some((k, v))
360    }
361}
362
363#[cfg(test)]
364mod tests {
365    use super::*;
366    use std::cell::Cell;
367    use std::hash::{Hash, Hasher};
368    use std::panic::{catch_unwind, AssertUnwindSafe};
369    use std::rc::Rc;
370
371    struct PanicHash(usize, Rc<Cell<Option<usize>>>);
372
373    impl PartialEq for PanicHash {
374        fn eq(&self, other: &Self) -> bool {
375            self.0 == other.0
376        }
377    }
378
379    impl Eq for PanicHash {}
380
381    impl Hash for PanicHash {
382        fn hash<H: Hasher>(&self, state: &mut H) {
383            assert_ne!(self.1.get(), Some(self.0));
384            self.0.hash(state);
385        }
386    }
387
388    #[test]
389    fn inlinehashmap_iter() {
390        let mut x = InlineHashMap::<String, usize, 5>::new();
391        x.insert("foo".into(), 3);
392        x.insert("bar".into(), 6);
393        x.insert("baz".into(), 7);
394        x.insert("qux".into(), 9);
395
396        let mut iter = x.iter();
397
398        // order is guaranteed as long as:
399        // - `InlineHashMap` is a stack-allocated array
400        // - `x.remove()` is never called
401
402        assert_eq!(iter.next(), Some((&"foo".into(), &3usize)));
403        assert_eq!(iter.next(), Some((&"bar".into(), &6usize)));
404        assert_eq!(iter.next(), Some((&"baz".into(), &7usize)));
405        assert_eq!(iter.next(), Some((&"qux".into(), &9usize)));
406    }
407
408    #[test]
409    fn inlinehashmap_growth_hash_panic_is_unwind_safe() {
410        let panic_on = Rc::new(Cell::new(None));
411        let key = |value| PanicHash(value, Rc::clone(&panic_on));
412        let mut map = InlineHashMap::<PanicHash, usize, 3>::new();
413        for value in 1..=3 {
414            map.insert(key(value), value);
415        }
416
417        panic_on.set(Some(2));
418        assert!(catch_unwind(AssertUnwindSafe(|| {
419            map.insert(key(4), 4);
420        }))
421        .is_err());
422        assert!(!map.is_heap_allocated());
423        assert!(map.iter().map(|(key, _)| key.0).eq([1]));
424
425        panic_on.set(None);
426        map.insert(key(5), 5);
427        assert_eq!(map.len(), 2);
428    }
429
430    #[test]
431    fn inlinehashmap_remove() {
432        let mut x = InlineHashMap::<usize, usize, 4>::new();
433        x.insert(789, 1336);
434        assert_eq!(x.len(), 1);
435        assert_eq!(x.get(&789), Some(&1336));
436        assert_eq!(x.remove(&789), Some(1336));
437        assert_eq!(x.len(), 0);
438
439        assert_eq!(x.remove(&789), None);
440
441        for i in 0..4 {
442            x.insert(i, i * 2);
443        }
444
445        assert!(!x.is_heap_allocated());
446        assert_eq!(x.len(), 4);
447
448        assert_eq!(x.remove(&2), Some(4));
449        assert_eq!(x.len(), 3);
450
451        assert_eq!(x.remove(&3), Some(6));
452        assert_eq!(x.len(), 2);
453
454        assert_eq!(x.remove(&1), Some(2));
455        assert_eq!(x.len(), 1);
456
457        assert_eq!(x.remove(&0), Some(0));
458        assert_eq!(x.len(), 0);
459        assert!(!x.is_heap_allocated());
460
461        // trigger heap allocation
462        for i in 0..8 {
463            x.insert(i, i * 2);
464        }
465        assert!(x.is_heap_allocated());
466        assert_eq!(x.len(), 8);
467
468        assert_eq!(x.remove(&7), Some(14));
469        assert_eq!(x.remove(&0), Some(0));
470    }
471
472    #[test]
473    fn inlinehashmap_remove_heap() {
474        let mut x = InlineHashMap::<usize, String, 4>::new();
475        x.insert(42, "test".into());
476        assert_eq!(x.len(), 1);
477        assert_eq!(x.remove(&42), Some("test".into()));
478        assert_eq!(x.len(), 0);
479    }
480
481    #[test]
482    fn inlinehashmap_clone() {
483        let mut x = InlineHashMapInner::<usize, usize, 4>::new();
484
485        for i in 0..10 {
486            x.insert(i, i * 2);
487        }
488
489        let x = x.clone();
490        assert_eq!(x.len(), 10);
491        assert!(x.is_heap_allocated());
492        assert_eq!(x.get(&9), Some(&18));
493    }
494
495    #[test]
496    fn inlinehashmap_to_map_stack() {
497        let mut x = InlineHashMapInner::<usize, usize, 4>::new();
498
499        for i in 0..4 {
500            x.insert(i, i * 2);
501        }
502
503        assert!(!x.is_heap_allocated());
504        assert_eq!(x.len(), 4);
505
506        let xx = x.to_map();
507        assert_eq!(xx.get(&0), Some(&0));
508        assert_eq!(xx.get(&1), Some(&2));
509        assert_eq!(xx.get(&2), Some(&4));
510        assert_eq!(xx.get(&3), Some(&6));
511        assert_eq!(xx.len(), 4);
512
513        x.insert(42, 1337);
514        assert!(x.is_heap_allocated());
515        assert_eq!(x.len(), 5);
516        assert_eq!(x.get(&42), Some(&1337));
517
518        let xx = x.to_map();
519        assert_eq!(xx.get(&0), Some(&0));
520        assert_eq!(xx.get(&42), Some(&1337));
521        assert_eq!(xx.len(), 5);
522    }
523
524    #[test]
525    fn inlinehashmap_to_map_heap() {
526        let mut x = InlineHashMapInner::<usize, String, 4>::new();
527
528        for i in 0..4 {
529            x.insert(i, i.to_string());
530        }
531
532        assert!(!x.is_heap_allocated());
533        assert_eq!(x.len(), 4);
534
535        let xx = x.to_map();
536        assert_eq!(&*xx[&0], "0");
537        assert_eq!(&*xx[&1], "1");
538        assert_eq!(&*xx[&2], "2");
539        assert_eq!(&*xx[&3], "3");
540        assert_eq!(xx.len(), 4);
541
542        x.insert(42, "1337".into());
543        assert!(x.is_heap_allocated());
544        assert_eq!(x.len(), 5);
545        assert_eq!(x.get(&42).map(|x| &**x), Some("1337"));
546
547        let xx = x.to_map();
548        assert_eq!(&*xx[&0], "0");
549        assert_eq!(&*xx[&42], "1337");
550        assert_eq!(xx.len(), 5);
551    }
552
553    #[test]
554    fn inlinehashmap_drop_stack() {
555        let mut x = InlineHashMapInner::<usize, String, 4>::new();
556
557        for i in 0..3 {
558            x.insert(i, i.to_string());
559        }
560
561        assert_eq!(x.len(), 3);
562        assert!(!x.is_heap_allocated());
563    }
564
565    #[test]
566    fn inlinehashmap_drop_heap() {
567        let mut x = InlineHashMapInner::<usize, String, 4>::new();
568
569        for i in 0..16 {
570            x.insert(i, i.to_string());
571        }
572
573        assert_eq!(x.len(), 16);
574        assert!(x.is_heap_allocated());
575    }
576
577    #[test]
578    fn inlinehashmap() {
579        let mut x = InlineHashMapInner::<&'static str, usize, 4>::new();
580        assert_eq!(x.len(), 0);
581        assert_eq!(x.get(&"hi"), None);
582        assert!(!x.is_heap_allocated());
583
584        x.insert("foo", 1337);
585        x.insert("foo", 1);
586        assert_eq!(x.len(), 1);
587        assert_eq!(x.get(&"foo"), Some(&1));
588        assert!(!x.is_heap_allocated());
589
590        x.insert("foo2", 2);
591        x.insert("foo3", 3);
592        x.insert("foo4", 4);
593
594        x.insert("foo", 2);
595        assert_eq!(x.len(), 4);
596        assert_eq!(x.get(&"foo"), Some(&2));
597        assert!(!x.is_heap_allocated());
598
599        x.insert("foo5", 5);
600        x.insert("foo", 3);
601        assert_eq!(x.len(), 5);
602        assert_eq!(x.get(&"foo"), Some(&3));
603        assert!(x.is_heap_allocated());
604
605        x.insert("foo6", 6);
606        x.insert("foo7", 7);
607        x.insert("foo8", 8);
608        x.insert("foo9", 9);
609        x.insert("foo10", 10);
610        x.insert("foo11", 11);
611    }
612}
tl/inline/hashmap.rs

tl/inline/
hashmap.rs
