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

tl/inline/
hashmap.rs
