Skip to main content

tl/inline/
hashmap.rs

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