small_collections/maps/
ordered_map.rs

1#![cfg(feature = "ordered")]
2//! Insertion-order-preserving map that lives on the stack and spills to the heap.
3//!
4//! Provides [`SmallOrderedMap`] — backed by [`HeaplessOrderedMap`] on the stack
5//! and [`ordermap::OrderMap`] on the heap.  Insertion order is maintained across
6//! both storage backends.
7//!
8//! Re-exports [`AnyMap`] from `map.rs` as the common trait.
9
10use core::mem::ManuallyDrop;
11use core::ptr;
12use std::borrow::Borrow;
13use std::fmt::{self, Debug};
14use std::hash::Hash;
15use std::iter::FromIterator;
16use std::ops::{Index, IndexMut};
17
18use ordermap::OrderMap;
19
20use crate::HeaplessOrderedMap;
21
22/// A trait for abstraction over different map types (Stack, Heap, Small).
23/// (Imported or redefined for convenience in this module)
24pub use crate::AnyMap;
25
26impl<K: Eq + Hash, V, S: std::hash::BuildHasher> AnyMap<K, V> for OrderMap<K, V, S> {
27    fn len(&self) -> usize {
28        self.len()
29    }
30    fn insert(&mut self, key: K, value: V) -> Option<V> {
31        self.insert(key, value)
32    }
33    fn get<Q>(&self, key: &Q) -> Option<&V>
34    where
35        K: Borrow<Q>,
36        Q: Hash + Eq + ?Sized,
37    {
38        self.get(key)
39    }
40    fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
41    where
42        K: Borrow<Q>,
43        Q: Hash + Eq + ?Sized,
44    {
45        self.get_mut(key)
46    }
47    fn remove<Q>(&mut self, key: &Q) -> Option<V>
48    where
49        K: Borrow<Q>,
50        Q: Hash + Eq + ?Sized,
51    {
52        self.remove(key)
53    }
54    fn contains_key<Q>(&self, key: &Q) -> bool
55    where
56        K: Borrow<Q>,
57        Q: Hash + Eq + ?Sized,
58    {
59        self.contains_key(key)
60    }
61    fn clear(&mut self) {
62        self.clear();
63    }
64}
65
66impl<K: Eq + Hash, V, const N: usize> AnyMap<K, V> for HeaplessOrderedMap<K, V, N> {
67    fn len(&self) -> usize {
68        self.len()
69    }
70    fn insert(&mut self, key: K, value: V) -> Option<V> {
71        self.insert(key, value).ok().flatten()
72    }
73    fn get<Q>(&self, key: &Q) -> Option<&V>
74    where
75        K: Borrow<Q>,
76        Q: Hash + Eq + ?Sized,
77    {
78        self.get(key)
79    }
80    fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
81    where
82        K: Borrow<Q>,
83        Q: Hash + Eq + ?Sized,
84    {
85        self.get_mut(key)
86    }
87    fn remove<Q>(&mut self, key: &Q) -> Option<V>
88    where
89        K: Borrow<Q>,
90        Q: Hash + Eq + ?Sized,
91    {
92        self.remove(key)
93    }
94    fn contains_key<Q>(&self, key: &Q) -> bool
95    where
96        K: Borrow<Q>,
97        Q: Hash + Eq + ?Sized,
98    {
99        self.contains_key(key)
100    }
101    fn clear(&mut self) {
102        self.clear();
103    }
104}
105
106/// An insertion-order-preserving map that lives on the stack for up to `N` entries,
107/// then spills to a heap-allocated `ordermap::OrderMap`.
108///
109/// # Storage strategy
110/// Uses a tagged union ([`MapData`]) with:
111/// - **Stack side**: [`HeaplessOrderedMap<K, V, N>`] — `heapless::LinearMap` wrapper.
112/// - **Heap side**: `ordermap::OrderMap<K, V>` — insertion-order-preserving hash map.
113///
114/// # Generic parameters
115/// | Parameter | Meaning |
116/// |-----------|--------|
117/// | `K` | Key type; must implement `Eq + Hash` |
118/// | `V` | Value type |
119/// | `N` | Stack capacity — max entries before spill |
120///
121/// # Design Considerations
122/// - **Insertion order**: unlike `SmallMap` (which uses `HashMap` on the heap), this
123///   type uses `OrderMap` which preserves insertion order even after a spill.
124/// - **`K: Eq + Hash` on the struct**: required by `HeaplessOrderedMap` and propagated
125///   to the outer struct to satisfy the `MapData` union's bounds.
126/// - **`remove` rebuilds on stack**: because `heapless::LinearMap::remove` requires
127///   `K: PartialEq<Q>` (stricter than `Borrow<Q>`), the stack-side `remove` is
128///   implemented by draining and rebuilding via `HeaplessOrderedMap::remove`.
129///
130/// # Safety
131/// `on_stack` determines which variant of `MapData` is active.  Only the active
132/// variant may be accessed.
133pub struct SmallOrderedMap<K: Eq + Hash, V, const N: usize> {
134    on_stack: bool,
135    data: MapData<K, V, N>,
136}
137
138impl<K: Eq + Hash, V, const N: usize> AnyMap<K, V> for SmallOrderedMap<K, V, N> {
139    fn len(&self) -> usize {
140        self.len()
141    }
142    fn insert(&mut self, key: K, value: V) -> Option<V> {
143        self.insert(key, value)
144    }
145    fn get<Q>(&self, key: &Q) -> Option<&V>
146    where
147        K: Borrow<Q>,
148        Q: Hash + Eq + ?Sized,
149    {
150        self.get(key)
151    }
152    fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
153    where
154        K: Borrow<Q>,
155        Q: Hash + Eq + ?Sized,
156    {
157        self.get_mut(key)
158    }
159    fn remove<Q>(&mut self, key: &Q) -> Option<V>
160    where
161        K: Borrow<Q>,
162        Q: Hash + Eq + ?Sized,
163    {
164        // SmallOrderedMap::remove has K: Borrow<Q> + PartialEq<Q> bound in original impl
165        // But for AnyMap we just need it to work.
166        self.remove(key)
167    }
168    fn contains_key<Q>(&self, key: &Q) -> bool
169    where
170        K: Borrow<Q>,
171        Q: Hash + Eq + ?Sized,
172    {
173        self.contains_key(key)
174    }
175    fn clear(&mut self) {
176        self.clear();
177    }
178}
179
180/// Internal storage for `SmallOrderedMap`.
181union MapData<K: Eq + Hash, V, const N: usize> {
182    stack: ManuallyDrop<HeaplessOrderedMap<K, V, N>>,
183    heap: ManuallyDrop<OrderMap<K, V>>,
184}
185
186impl<K, V, const N: usize> SmallOrderedMap<K, V, N>
187where
188    K: Eq + Hash,
189{
190    /// A constant parameter.
191    pub const MAX_STACK_SIZE: usize = 16 * 1024;
192
193    /// Creates a new empty ordered map on the stack.
194    pub fn new() -> Self {
195        const {
196            assert!(
197                std::mem::size_of::<Self>() <= SmallOrderedMap::<K, V, N>::MAX_STACK_SIZE,
198                "SmallOrderedMap is too large! Reduce N."
199            );
200        }
201
202        Self {
203            on_stack: true,
204            data: MapData {
205                stack: ManuallyDrop::new(HeaplessOrderedMap::new()),
206            },
207        }
208    }
209
210    /// Returns `true` if the map is currently storing data on the stack.
211    #[inline]
212    pub fn is_on_stack(&self) -> bool {
213        self.on_stack
214    }
215
216    /// Returns the number of elements in the map.
217    pub fn len(&self) -> usize {
218        unsafe {
219            if self.on_stack {
220                self.data.stack.len()
221            } else {
222                self.data.heap.len()
223            }
224        }
225    }
226
227    /// Returns `true` if the map contains no elements.
228    pub fn is_empty(&self) -> bool {
229        self.len() == 0
230    }
231
232    /// Clears the map, removing all key-value pairs.
233    pub fn clear(&mut self) {
234        unsafe {
235            if self.on_stack {
236                (*self.data.stack).clear();
237            } else {
238                (*self.data.heap).clear();
239            }
240        }
241    }
242
243    /// Inserts a key-value pair into the map.
244    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
245        unsafe {
246            if self.on_stack {
247                let stack_map = &mut *self.data.stack;
248                match stack_map.insert(key, value) {
249                    Ok(old) => return old,
250                    Err((k, v)) => {
251                        self.spill_to_heap();
252                        return (*self.data.heap).insert(k, v);
253                    }
254                }
255            }
256
257            (*self.data.heap).insert(key, value)
258        }
259    }
260
261    /// Retrieves a reference to the value corresponding to the key.
262    pub fn get<Q>(&self, key: &Q) -> Option<&V>
263    where
264        K: Borrow<Q>,
265        Q: Hash + Eq + ?Sized,
266    {
267        unsafe {
268            if self.on_stack {
269                self.data.stack.get(key)
270            } else {
271                self.data.heap.get(key)
272            }
273        }
274    }
275
276    /// Retrieves a mutable reference to the value corresponding to the key.
277    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
278    where
279        K: Borrow<Q>,
280        Q: Hash + Eq + ?Sized,
281    {
282        unsafe {
283            if self.on_stack {
284                (*self.data.stack).get_mut(key)
285            } else {
286                (*self.data.heap).get_mut(key)
287            }
288        }
289    }
290
291    /// Removes a key from the map, returning the value at the key if the key was previously in the map.
292    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
293    where
294        K: Borrow<Q>,
295        Q: Hash + Eq + ?Sized,
296    {
297        unsafe {
298            if self.on_stack {
299                (*self.data.stack).remove(key)
300            } else {
301                (*self.data.heap).remove(key)
302            }
303        }
304    }
305
306    /// Returns `true` if the map contains a value for the specified key.
307    pub fn contains_key<Q>(&self, key: &Q) -> bool
308    where
309        K: Borrow<Q>,
310        Q: Hash + Eq + ?Sized,
311    {
312        self.get(key).is_some()
313    }
314
315    #[inline(never)]
316    unsafe fn spill_to_heap(&mut self) {
317        unsafe {
318            let stack_map = ManuallyDrop::take(&mut self.data.stack);
319            let mut new_heap = OrderMap::with_capacity(stack_map.len() * 2);
320
321            for (key, value) in stack_map {
322                new_heap.insert(key, value);
323            }
324
325            ptr::write(&mut self.data.heap, ManuallyDrop::new(new_heap));
326            self.on_stack = false;
327        }
328    }
329}
330
331// --- Index Traits ---
332
333impl<K, V, Q, const N: usize> Index<&Q> for SmallOrderedMap<K, V, N>
334where
335    K: Eq + Hash + Borrow<Q>,
336    Q: Eq + Hash + ?Sized,
337{
338    type Output = V;
339
340    fn index(&self, key: &Q) -> &Self::Output {
341        self.get(key).expect("no entry found for key")
342    }
343}
344
345impl<K, V, Q, const N: usize> IndexMut<&Q> for SmallOrderedMap<K, V, N>
346where
347    K: Eq + Hash + Borrow<Q>,
348    Q: Eq + Hash + ?Sized,
349{
350    fn index_mut(&mut self, key: &Q) -> &mut Self::Output {
351        self.get_mut(key).expect("no entry found for key")
352    }
353}
354
355// --- Traits ---
356
357impl<K: Eq + Hash, V, const N: usize> Drop for SmallOrderedMap<K, V, N> {
358    fn drop(&mut self) {
359        unsafe {
360            if self.on_stack {
361                ManuallyDrop::drop(&mut self.data.stack);
362            } else {
363                ManuallyDrop::drop(&mut self.data.heap);
364            }
365        }
366    }
367}
368
369impl<K, V, const N: usize> Default for SmallOrderedMap<K, V, N>
370where
371    K: Eq + Hash,
372{
373    fn default() -> Self {
374        Self::new()
375    }
376}
377
378impl<K, V, const N: usize> Clone for SmallOrderedMap<K, V, N>
379where
380    K: Eq + Hash + Clone,
381    V: Clone,
382{
383    fn clone(&self) -> Self {
384        unsafe {
385            if self.on_stack {
386                Self {
387                    on_stack: true,
388                    data: MapData {
389                        stack: ManuallyDrop::new((*self.data.stack).clone()),
390                    },
391                }
392            } else {
393                Self {
394                    on_stack: false,
395                    data: MapData {
396                        heap: ManuallyDrop::new((*self.data.heap).clone()),
397                    },
398                }
399            }
400        }
401    }
402}
403
404impl<K, V, const N: usize> Debug for SmallOrderedMap<K, V, N>
405where
406    K: Eq + Hash + Debug,
407    V: Debug,
408{
409    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
410        f.debug_map().entries(self.iter()).finish()
411    }
412}
413
414impl<K, V, const N: usize, M> PartialEq<M> for SmallOrderedMap<K, V, N>
415where
416    K: Eq + Hash,
417    V: PartialEq,
418    M: AnyMap<K, V>,
419{
420    fn eq(&self, other: &M) -> bool {
421        if self.len() != other.len() {
422            return false;
423        }
424        self.iter().all(|(k, v)| other.get(k) == Some(v))
425    }
426}
427
428impl<K, V, const N: usize> Eq for SmallOrderedMap<K, V, N>
429where
430    K: Eq + Hash,
431    V: Eq,
432{
433}
434
435impl<K, V, const N: usize> SmallOrderedMap<K, V, N>
436where
437    K: Eq + Hash,
438{
439    /// Returns an iterator over the elements.
440    pub fn iter(&self) -> SmallMapIter<'_, K, V> {
441        unsafe {
442            if self.on_stack {
443                SmallMapIter::Stack(self.data.stack.iter())
444            } else {
445                SmallMapIter::Heap(self.data.heap.iter())
446            }
447        }
448    }
449}
450
451/// Types of `SmallMapIter`.
452pub enum SmallMapIter<'a, K, V> {
453    /// Automatically generated documentation for this item.
454    Stack(heapless::linear_map::Iter<'a, K, V>),
455    /// Automatically generated documentation for this item.
456    Heap(ordermap::map::Iter<'a, K, V>),
457}
458
459impl<'a, K, V> Iterator for SmallMapIter<'a, K, V> {
460    type Item = (&'a K, &'a V);
461    fn next(&mut self) -> Option<Self::Item> {
462        match self {
463            SmallMapIter::Stack(i) => i.next(),
464            SmallMapIter::Heap(i) => i.next(),
465        }
466    }
467}
468
469impl<K, V, const N: usize> IntoIterator for SmallOrderedMap<K, V, N>
470where
471    K: Eq + Hash,
472{
473    type Item = (K, V);
474    type IntoIter = SmallMapIntoIter<K, V, N>;
475
476    fn into_iter(self) -> Self::IntoIter {
477        let mut this = ManuallyDrop::new(self);
478        unsafe {
479            if this.on_stack {
480                SmallMapIntoIter::Stack(
481                    ManuallyDrop::<HeaplessOrderedMap<K, V, N>>::take(&mut this.data.stack)
482                        .into_iter(),
483                )
484            } else {
485                SmallMapIntoIter::Heap(ManuallyDrop::take(&mut this.data.heap).into_iter())
486            }
487        }
488    }
489}
490
491/// Types of `SmallMapIntoIter`.
492pub enum SmallMapIntoIter<K: Eq, V, const N: usize> {
493    /// Automatically generated documentation for this item.
494    Stack(heapless::linear_map::IntoIter<K, V, N>),
495    /// Automatically generated documentation for this item.
496    Heap(ordermap::map::IntoIter<K, V>),
497}
498
499impl<K, V, const N: usize> Iterator for SmallMapIntoIter<K, V, N>
500where
501    K: Eq + Hash,
502{
503    type Item = (K, V);
504    fn next(&mut self) -> Option<Self::Item> {
505        match self {
506            SmallMapIntoIter::Stack(i) => i.next(),
507            SmallMapIntoIter::Heap(i) => i.next(),
508        }
509    }
510}
511
512impl<K, V, const N: usize> FromIterator<(K, V)> for SmallOrderedMap<K, V, N>
513where
514    K: Eq + Hash,
515{
516    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
517        let mut map = SmallOrderedMap::new();
518        for (k, v) in iter {
519            map.insert(k, v);
520        }
521        map
522    }
523}
524
525impl<K, V, const N: usize> Extend<(K, V)> for SmallOrderedMap<K, V, N>
526where
527    K: Eq + Hash,
528{
529    fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
530        for (k, v) in iter {
531            self.insert(k, v);
532        }
533    }
534}
535
536#[cfg(test)]
537mod tests {
538    use super::*;
539
540    #[test]
541    fn test_ordered_map_stack_ops_insertion_order() {
542        let mut map: SmallOrderedMap<i32, i32, 4> = SmallOrderedMap::new();
543        map.insert(3, 30);
544        map.insert(1, 10);
545        map.insert(2, 20);
546
547        let keys: Vec<_> = map.iter().map(|(k, _)| *k).collect();
548        assert_eq!(keys, vec![3, 1, 2]);
549    }
550
551    #[test]
552    fn test_ordered_map_spill_trigger_on_insert() {
553        let mut map: SmallOrderedMap<i32, i32, 2> = SmallOrderedMap::new();
554        map.insert(1, 10);
555        map.insert(2, 20);
556        assert!(map.is_on_stack());
557
558        map.insert(3, 30); // Spill
559        assert!(!map.is_on_stack());
560
561        let keys: Vec<_> = map.iter().map(|(k, _)| *k).collect();
562        assert_eq!(keys, vec![1, 2, 3]);
563    }
564
565    #[test]
566    fn test_ordered_map_any_storage_lifecycle() {
567        let mut map: SmallOrderedMap<String, i32, 2> = SmallOrderedMap::new();
568        map.insert("a".into(), 1);
569        map.insert("b".into(), 2);
570        assert_eq!(map.get("a"), Some(&1));
571        assert_eq!(map.remove("a"), Some(1));
572        assert_eq!(map.len(), 1);
573        assert!(map.is_on_stack());
574
575        map.insert("c".into(), 3);
576        map.insert("d".into(), 4); // Spill
577        assert!(!map.is_on_stack());
578        assert_eq!(map.len(), 3);
579        assert_eq!(map.get("b"), Some(&2));
580        assert_eq!(map.get("c"), Some(&3));
581        assert_eq!(map.get("d"), Some(&4));
582    }
583
584    #[test]
585    fn test_ordered_map_any_storage_get_mut() {
586        let mut map: SmallOrderedMap<i32, i32, 4> = SmallOrderedMap::new();
587        map.insert(1, 10);
588        map.insert(2, 20);
589        if let Some(v) = map.get_mut(&1) {
590            *v = 11;
591        }
592        assert_eq!(map[&1], 11);
593        map[&2] = 22;
594        assert_eq!(map.get(&2), Some(&22));
595
596        for i in 3..6 {
597            map.insert(i, i * 10);
598        }
599        assert!(!map.is_on_stack());
600        map[&5] = 55;
601        assert_eq!(map.get(&5), Some(&55));
602    }
603
604    #[test]
605    fn test_ordered_map_any_storage_clear() {
606        let mut map: SmallOrderedMap<i32, i32, 2> = SmallOrderedMap::new();
607        assert!(map.is_empty());
608        map.insert(1, 10);
609        map.clear();
610        assert!(map.is_empty());
611
612        map.insert(1, 10);
613        map.insert(2, 20);
614        map.insert(3, 30); // Spill
615        map.clear();
616        assert!(map.is_empty());
617        assert!(!map.is_on_stack());
618    }
619
620    #[test]
621    fn test_ordered_map_traits_exhaustive() {
622        let mut map: SmallOrderedMap<i32, i32, 4> = SmallOrderedMap::new();
623        map.insert(1, 10);
624        map.insert(2, 20);
625
626        // Clone
627        let cloned = map.clone();
628        assert_eq!(cloned.len(), 2);
629        assert_eq!(cloned.get(&1), Some(&10));
630
631        // Debug
632        let debug = format!("{:?}", map);
633        assert!(debug.contains("1: 10"));
634
635        // Default
636        let def: SmallOrderedMap<i32, i32, 4> = SmallOrderedMap::default();
637        assert!(def.is_empty());
638
639        // FromIterator
640        let collected: SmallOrderedMap<i32, i32, 4> = vec![(1, 10), (2, 20)].into_iter().collect();
641        assert_eq!(collected.len(), 2);
642        let items: Vec<_> = collected.iter().collect();
643        assert_eq!(items, vec![(&1, &10), (&2, &20)]);
644
645        // Extend
646        let mut map2 = SmallOrderedMap::<i32, i32, 4>::new();
647        map2.extend(vec![(1, 10), (2, 20)]);
648        assert_eq!(map2.len(), 2);
649
650        // IntoIterator
651        let vec: Vec<_> = map2.into_iter().collect();
652        assert_eq!(vec.len(), 2);
653        assert!(vec.contains(&(1, 10)));
654        assert!(vec.contains(&(2, 20)));
655    }
656
657    #[test]
658    fn test_ordered_map_any_storage_borrow_lookups() {
659        // LinearMap borrow lookup
660        let mut map: SmallOrderedMap<String, i32, 4> = SmallOrderedMap::new();
661        map.insert("Apple".to_string(), 100);
662        assert_eq!(map.get("Apple"), Some(&100));
663        assert_eq!(map.get_mut("Apple"), Some(&mut 100));
664    }
665
666    #[test]
667    fn test_ordered_map_any_storage_clone_heap() {
668        let h: SmallOrderedMap<i32, i32, 2> = vec![(1, 1), (2, 2), (3, 3)].into_iter().collect();
669        assert!(!h.is_on_stack());
670
671        let h2 = h.clone();
672        assert_eq!(h2.len(), 3);
673        assert!(!h2.is_on_stack());
674    }
675
676    #[test]
677    fn test_ordered_map_traits_into_iter_heap() {
678        let h: SmallOrderedMap<i32, i32, 2> = vec![(1, 1), (2, 2), (3, 3)].into_iter().collect();
679        let mut it = h.into_iter();
680        assert_eq!(it.next(), Some((1, 1)));
681    }
682}
683
684#[cfg(test)]
685mod ordered_map_coverage_tests {
686    use super::*;
687
688    fn run_any_map_test<M: AnyMap<i32, i32>>(any_map: &mut M) {
689        assert_eq!(any_map.len(), 0);
690        any_map.insert(1, 10);
691        assert_eq!(any_map.get(&1), Some(&10));
692        assert_eq!(any_map.get_mut(&1), Some(&mut 10));
693        assert!(any_map.contains_key(&1));
694        assert_eq!(any_map.remove(&1), Some(10));
695        any_map.insert(2, 20);
696        any_map.clear();
697        assert_eq!(any_map.len(), 0);
698    }
699
700    #[test]
701    fn test_any_map_trait_impls() {
702        let mut ord_map: OrderMap<i32, i32> = OrderMap::new();
703        run_any_map_test(&mut ord_map);
704
705        let mut hl_map: HeaplessOrderedMap<i32, i32, 2> = HeaplessOrderedMap::new();
706        run_any_map_test(&mut hl_map);
707
708        let mut small_map: SmallOrderedMap<i32, i32, 2> = SmallOrderedMap::new();
709        run_any_map_test(&mut small_map);
710    }
711
712    #[test]
713    fn test_small_ordered_map_insert_heap_branch() {
714        let mut map: SmallOrderedMap<i32, i32, 2> = SmallOrderedMap::new();
715        map.insert(1, 10);
716        map.insert(2, 20);
717        map.insert(3, 30); // spill
718
719        let old = map.insert(3, 300); // heap insert branch
720        assert_eq!(old, Some(30));
721    }
722
723    #[test]
724    fn test_small_ordered_map_partial_eq_length() {
725        let mut m1: SmallOrderedMap<i32, i32, 2> = SmallOrderedMap::new();
726        m1.insert(1, 10);
727
728        let mut m2: SmallOrderedMap<i32, i32, 2> = SmallOrderedMap::new();
729        m2.insert(1, 10);
730        m2.insert(2, 20);
731
732        assert_ne!(m1, m2);
733    }
734}
small_collections/maps/ordered_map.rs

small_collections/maps/
ordered_map.rs
