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

tl/inline/
hashmap.rs
