small_collections/maps/
btree_map.rs

1//! Sorted map that lives on the stack and spills to the heap.
2//!
3//! This module provides [`SmallBTreeMap`] — a map that stores up to `N` entries
4//! in a stack-allocated sorted vector ([`HeaplessBTreeMap`]) and transparently
5//! migrates to a `std::collections::BTreeMap` when the stack overflows.
6//!
7//! [`AnyBTreeMap`] is an object-safe trait that unifies both storage backends.
8
9use core::borrow::Borrow;
10use core::cmp::Ordering;
11use core::mem::ManuallyDrop;
12use std::collections::BTreeMap;
13use std::fmt::{self, Debug};
14use std::iter::FromIterator;
15
16use crate::maps::heapless_btree_map::{Entry, HeaplessBTreeMap};
17
18/// An object-safe abstraction over B-Tree map types.
19///
20/// Implemented by `BTreeMap<K, V>` (heap) and `SmallBTreeMap<K, V, N>` (small/stack) so
21/// that callers can write backend-agnostic code.
22pub trait AnyBTreeMap<K, V> {
23    /// Returns the number of key-value pairs.
24    fn len(&self) -> usize;
25    /// Returns `true` if the map is empty.
26    fn is_empty(&self) -> bool {
27        self.len() == 0
28    }
29    /// Inserts `(key, value)`, returning the previous value if the key existed.
30    fn insert(&mut self, key: K, value: V) -> Option<V>;
31    /// Returns a shared reference to the value for `key`, or `None`.
32    fn get<Q>(&self, key: &Q) -> Option<&V>
33    where
34        K: Borrow<Q>,
35        Q: Ord + ?Sized;
36    /// Returns an exclusive reference to the value for `key`, or `None`.
37    fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
38    where
39        K: Borrow<Q>,
40        Q: Ord + ?Sized;
41    /// Removes and returns the value for `key`, or `None` if absent.
42    fn remove<Q>(&mut self, key: &Q) -> Option<V>
43    where
44        K: Borrow<Q>,
45        Q: Ord + ?Sized;
46    /// Removes all entries.
47    fn clear(&mut self);
48}
49
50impl<K: Ord, V> AnyBTreeMap<K, V> for BTreeMap<K, V> {
51    fn len(&self) -> usize {
52        self.len()
53    }
54    fn insert(&mut self, key: K, value: V) -> Option<V> {
55        self.insert(key, value)
56    }
57    fn get<Q>(&self, key: &Q) -> Option<&V>
58    where
59        K: Borrow<Q>,
60        Q: Ord + ?Sized,
61    {
62        self.get(key)
63    }
64    fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
65    where
66        K: Borrow<Q>,
67        Q: Ord + ?Sized,
68    {
69        self.get_mut(key)
70    }
71    fn remove<Q>(&mut self, key: &Q) -> Option<V>
72    where
73        K: Borrow<Q>,
74        Q: Ord + ?Sized,
75    {
76        self.remove(key)
77    }
78    fn clear(&mut self) {
79        self.clear();
80    }
81}
82
83impl<K: Ord, V, const N: usize> AnyBTreeMap<K, V> for HeaplessBTreeMap<K, V, N> {
84    fn len(&self) -> usize {
85        self.len()
86    }
87    fn insert(&mut self, key: K, value: V) -> Option<V> {
88        self.insert(key, value).ok().flatten()
89    }
90    fn get<Q>(&self, key: &Q) -> Option<&V>
91    where
92        K: Borrow<Q>,
93        Q: Ord + ?Sized,
94    {
95        self.get(key)
96    }
97    fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
98    where
99        K: Borrow<Q>,
100        Q: Ord + ?Sized,
101    {
102        self.get_mut(key)
103    }
104    fn remove<Q>(&mut self, key: &Q) -> Option<V>
105    where
106        K: Borrow<Q>,
107        Q: Ord + ?Sized,
108    {
109        self.remove(key)
110    }
111    fn clear(&mut self) {
112        self.clear();
113    }
114}
115
116/// A sorted map that lives on the stack for up to `N` entries, then spills to the heap.
117///
118/// # Storage strategy
119/// Uses a tagged union ([`MapData`]) with:
120/// - **Stack side**: [`HeaplessBTreeMap<K, V, N>`] — a sorted `heapless::Vec`.
121/// - **Heap side**: `std::collections::BTreeMap<K, V>`.
122///
123/// Once a spill occurs the struct permanently uses the heap side.
124///
125/// # Overflow protocol
126/// When [`insert`](SmallBTreeMap::insert) is called on a full stack map with a new key,
127/// `spill_to_heap` migrates all entries to a `BTreeMap` and then re-inserts the new pair.
128///
129/// # Generic parameters
130/// | Parameter | Meaning |
131/// |-----------|--------|
132/// | `K` | Key type; must implement `Ord` |
133/// | `V` | Value type |
134/// | `N` | Stack capacity — max entries before spill |
135///
136/// # Design Considerations
137/// - **`len` tracked separately**: instead of accessing the union variant to call `.len()`,
138///   the length is cached in a plain `usize` field so `len()` never touches unsafe code.
139/// - **Sorted iteration**: the stack-side sorted order is preserved during spill, so
140///   `BTreeMap`'s sorted iteration invariant is maintained without re-sorting.
141/// - **`Entry` API**: the `entry` method (in `map.rs` / `SmallMap`) can be layered on top;
142///   this struct deliberately keeps its API minimal.
143///
144/// # Safety
145/// `on_stack` determines which variant of `MapData` is active. Only that variant
146/// may be accessed. All unsafe union accesses must first check `on_stack`.
147pub struct SmallBTreeMap<K, V, const N: usize> {
148    on_stack: bool,
149    len: usize,
150    data: MapData<K, V, N>,
151}
152
153impl<K: Ord, V, const N: usize> AnyBTreeMap<K, V> for SmallBTreeMap<K, V, N> {
154    fn len(&self) -> usize {
155        self.len()
156    }
157    fn insert(&mut self, key: K, value: V) -> Option<V> {
158        self.insert(key, value)
159    }
160    fn get<Q>(&self, key: &Q) -> Option<&V>
161    where
162        K: Borrow<Q>,
163        Q: Ord + ?Sized,
164    {
165        self.get(key)
166    }
167    fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
168    where
169        K: Borrow<Q>,
170        Q: Ord + ?Sized,
171    {
172        self.get_mut(key)
173    }
174    fn remove<Q>(&mut self, key: &Q) -> Option<V>
175    where
176        K: Borrow<Q>,
177        Q: Ord + ?Sized,
178    {
179        self.remove(key)
180    }
181    fn clear(&mut self) {
182        self.clear();
183    }
184}
185
186/// The internal storage for `SmallBTreeMap`.
187///
188/// We use `ManuallyDrop` because the compiler cannot know which field is active
189/// and therefore cannot automatically drop the correct one.
190union MapData<K, V, const N: usize> {
191    stack: ManuallyDrop<HeaplessBTreeMap<K, V, N>>,
192    heap: ManuallyDrop<BTreeMap<K, V>>,
193}
194
195impl<K, V, const N: usize> SmallBTreeMap<K, V, N>
196where
197    K: Ord,
198{
199    /// A constant parameter.
200    pub const MAX_STACK_SIZE: usize = 16 * 1024;
201
202    /// Creates a new empty map on the stack.
203    pub fn new() -> Self {
204        const {
205            assert!(
206                std::mem::size_of::<Self>() <= Self::MAX_STACK_SIZE,
207                "SmallBTreeMap is too large! The struct size exceeds the 16KB limit. Reduce N."
208            );
209        }
210        Self {
211            on_stack: true,
212            len: 0,
213            data: MapData {
214                stack: ManuallyDrop::new(HeaplessBTreeMap::new()),
215            },
216        }
217    }
218
219    /// Creates an empty map with the specified capacity.
220    /// If the capacity exceeds the stack limit `N`, it will be created directly on the heap.
221    pub fn with_capacity(cap: usize) -> Self {
222        if cap <= N {
223            Self::new()
224        } else {
225            Self {
226                on_stack: false,
227                len: 0,
228                data: MapData::<K, V, N> {
229                    heap: ManuallyDrop::new(BTreeMap::new()),
230                },
231            }
232        }
233    }
234
235    /// Returns `true` if the map is currently storing data on the stack.
236    #[inline]
237    pub fn is_on_stack(&self) -> bool {
238        self.on_stack
239    }
240
241    /// Returns the number of elements in the map.
242    pub fn len(&self) -> usize {
243        self.len
244    }
245
246    /// Returns `true` if the map contains no elements.
247    pub fn is_empty(&self) -> bool {
248        self.len == 0
249    }
250
251    /// Clears the map, removing all key-value pairs.
252    pub fn clear(&mut self) {
253        unsafe {
254            if self.on_stack {
255                (*self.data.stack).clear();
256            } else {
257                (*self.data.heap).clear();
258            }
259        }
260        self.len = 0;
261    }
262
263    /// Inserts a key-value pair into the map.
264    /// If the map is on the stack and full, this triggers a transparent spill to the heap.
265    pub fn insert(&mut self, key: K, value: V) -> Option<V> {
266        unsafe {
267            if self.on_stack {
268                let stack = &mut *self.data.stack;
269                match stack.insert(key, value) {
270                    Ok(old) => {
271                        if old.is_none() {
272                            self.len += 1;
273                        }
274                        return old;
275                    }
276                    Err((k, v)) => {
277                        // Stack is full: spill, then insert into heap with the returned k/v.
278                        self.spill_to_heap();
279                        let old = (*self.data.heap).insert(k, v);
280                        if old.is_none() {
281                            self.len += 1;
282                        }
283                        return old;
284                    }
285                }
286            }
287
288            // Heap path
289            let old = (*self.data.heap).insert(key, value);
290            if old.is_none() {
291                self.len += 1;
292            }
293            old
294        }
295    }
296
297    /// Returns a reference to the value corresponding to the key.
298    pub fn get<Q>(&self, key: &Q) -> Option<&V>
299    where
300        K: Borrow<Q>,
301        Q: Ord + ?Sized,
302    {
303        unsafe {
304            if self.on_stack {
305                self.data.stack.get(key)
306            } else {
307                (*self.data.heap).get(key)
308            }
309        }
310    }
311
312    /// Returns a mutable reference to the value corresponding to the key.
313    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
314    where
315        K: Borrow<Q>,
316        Q: Ord + ?Sized,
317    {
318        unsafe {
319            if self.on_stack {
320                (*self.data.stack).get_mut(key)
321            } else {
322                (*self.data.heap).get_mut(key)
323            }
324        }
325    }
326
327    /// Removes a key from the map, returning the value at the key if the key was previously in the map.
328    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
329    where
330        K: Borrow<Q>,
331        Q: Ord + ?Sized,
332    {
333        unsafe {
334            if self.on_stack {
335                let old = (*self.data.stack).remove(key);
336                if old.is_some() {
337                    self.len -= 1;
338                }
339                old
340            } else {
341                let old = (*self.data.heap).remove(key);
342                if old.is_some() {
343                    self.len -= 1;
344                }
345                old
346            }
347        }
348    }
349
350    /// Internal method to transition storage from stack (`HeaplessBTreeMap`) to heap (`BTreeMap`).
351    #[inline(never)]
352    unsafe fn spill_to_heap(&mut self) {
353        unsafe {
354            let mut heap = BTreeMap::new();
355            // ManuallyDrop::take copies the bits out; safe because we immediately overwrite.
356            let stack_map = ManuallyDrop::take(&mut self.data.stack);
357
358            for entry in stack_map {
359                heap.insert(entry.0, entry.1);
360            }
361
362            self.data.heap = ManuallyDrop::new(heap);
363            self.on_stack = false;
364        }
365    }
366
367    /// Returns an iterator over the elements.
368    pub fn iter(&self) -> Iter<'_, K, V> {
369        unsafe {
370            if self.on_stack {
371                Iter::Stack(self.data.stack.iter())
372            } else {
373                Iter::Heap((*self.data.heap).iter())
374            }
375        }
376    }
377}
378
379/// Types of `Iter`.
380pub enum Iter<'a, K: Ord, V> {
381    /// Automatically generated documentation for this item.
382    Stack(core::slice::Iter<'a, Entry<K, V>>),
383    /// Automatically generated documentation for this item.
384    Heap(std::collections::btree_map::Iter<'a, K, V>),
385}
386
387impl<'a, K: Ord, V> Iterator for Iter<'a, K, V> {
388    type Item = (&'a K, &'a V);
389    fn next(&mut self) -> Option<Self::Item> {
390        match self {
391            Iter::Stack(i) => i.next().map(|e| (&e.0, &e.1)),
392            Iter::Heap(i) => i.next(),
393        }
394    }
395}
396
397/// Convenience alias used by `SmallBTreeSet`.
398pub type IntoIter<K, V, const N: usize> = SmallBTreeMapIntoIter<K, V, N>;
399
400/// A structure representing `HeaplessBTreeMapIntoIter`.
401pub struct HeaplessBTreeMapIntoIter<K, V, const N: usize> {
402    inner: heapless::vec::IntoIter<Entry<K, V>, N, usize>,
403}
404
405impl<K: Ord, V, const N: usize> Iterator for HeaplessBTreeMapIntoIter<K, V, N> {
406    type Item = (K, V);
407    fn next(&mut self) -> Option<Self::Item> {
408        self.inner.next().map(|e| (e.0, e.1))
409    }
410}
411
412/// Types of `SmallBTreeMapIntoIter`.
413pub enum SmallBTreeMapIntoIter<K, V, const N: usize> {
414    /// Automatically generated documentation for this item.
415    Stack(HeaplessBTreeMapIntoIter<K, V, N>),
416    /// Automatically generated documentation for this item.
417    Heap(std::collections::btree_map::IntoIter<K, V>),
418}
419
420impl<K: Ord, V, const N: usize> Iterator for SmallBTreeMapIntoIter<K, V, N> {
421    type Item = (K, V);
422    fn next(&mut self) -> Option<Self::Item> {
423        match self {
424            SmallBTreeMapIntoIter::Stack(i) => i.next(),
425            SmallBTreeMapIntoIter::Heap(i) => i.next(),
426        }
427    }
428}
429
430impl<K, V, const N: usize> IntoIterator for SmallBTreeMap<K, V, N>
431where
432    K: Ord,
433{
434    type Item = (K, V);
435    type IntoIter = SmallBTreeMapIntoIter<K, V, N>;
436
437    fn into_iter(self) -> Self::IntoIter {
438        let mut this = ManuallyDrop::new(self);
439        unsafe {
440            if this.on_stack {
441                SmallBTreeMapIntoIter::Stack(HeaplessBTreeMapIntoIter {
442                    inner: ManuallyDrop::<HeaplessBTreeMap<K, V, N>>::take(&mut this.data.stack)
443                        .into_iter(),
444                })
445            } else {
446                SmallBTreeMapIntoIter::Heap(ManuallyDrop::take(&mut this.data.heap).into_iter())
447            }
448        }
449    }
450}
451
452impl<K, V, const N: usize> Drop for SmallBTreeMap<K, V, N> {
453    fn drop(&mut self) {
454        unsafe {
455            if self.on_stack {
456                ManuallyDrop::drop(&mut self.data.stack);
457            } else {
458                ManuallyDrop::drop(&mut self.data.heap);
459            }
460        }
461    }
462}
463
464impl<K, V, const N: usize> Clone for SmallBTreeMap<K, V, N>
465where
466    K: Ord + Clone,
467    V: Clone,
468{
469    fn clone(&self) -> Self {
470        unsafe {
471            if self.on_stack {
472                Self {
473                    on_stack: true,
474                    len: self.len,
475                    data: MapData {
476                        stack: ManuallyDrop::new((*self.data.stack).clone()),
477                    },
478                }
479            } else {
480                Self {
481                    on_stack: false,
482                    len: self.len,
483                    data: MapData {
484                        heap: ManuallyDrop::new((*self.data.heap).clone()),
485                    },
486                }
487            }
488        }
489    }
490}
491
492impl<K: PartialEq + Ord, V: PartialEq, const N: usize, M: AnyBTreeMap<K, V>> PartialEq<M>
493    for SmallBTreeMap<K, V, N>
494{
495    fn eq(&self, other: &M) -> bool {
496        if self.len() != other.len() {
497            return false;
498        }
499        self.iter().all(|(k, v)| other.get(k) == Some(v))
500    }
501}
502
503impl<K: Eq + Ord, V: Eq, const N: usize> Eq for SmallBTreeMap<K, V, N> {}
504
505impl<K: PartialOrd + Ord, V: PartialOrd, const N: usize> PartialOrd for SmallBTreeMap<K, V, N> {
506    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
507        self.iter().partial_cmp(other.iter())
508    }
509}
510
511impl<K: Ord, V: Ord, const N: usize> Ord for SmallBTreeMap<K, V, N> {
512    fn cmp(&self, other: &Self) -> Ordering {
513        self.iter().cmp(other.iter())
514    }
515}
516
517impl<K: Ord, V, const N: usize> Default for SmallBTreeMap<K, V, N> {
518    fn default() -> Self {
519        Self::new()
520    }
521}
522
523impl<K: Ord + Debug, V: Debug, const N: usize> Debug for SmallBTreeMap<K, V, N> {
524    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
525        f.debug_map().entries(self.iter()).finish()
526    }
527}
528
529impl<K, V, const N: usize> FromIterator<(K, V)> for SmallBTreeMap<K, V, N>
530where
531    K: Ord,
532{
533    fn from_iter<T: IntoIterator<Item = (K, V)>>(iter: T) -> Self {
534        let mut map = Self::new();
535        for (k, v) in iter {
536            map.insert(k, v);
537        }
538        map
539    }
540}
541
542impl<K, V, const N: usize> Extend<(K, V)> for SmallBTreeMap<K, V, N>
543where
544    K: Ord,
545{
546    fn extend<I: IntoIterator<Item = (K, V)>>(&mut self, iter: I) {
547        for (k, v) in iter {
548            self.insert(k, v);
549        }
550    }
551}
552
553// --- Index Traits ---
554use std::ops::{Index, IndexMut};
555
556impl<K, V, Q, const N: usize> Index<&Q> for SmallBTreeMap<K, V, N>
557where
558    K: Ord + Borrow<Q>,
559    Q: Ord + ?Sized,
560{
561    type Output = V;
562
563    fn index(&self, key: &Q) -> &Self::Output {
564        self.get(key).expect("no entry found for key")
565    }
566}
567
568impl<K, V, Q, const N: usize> IndexMut<&Q> for SmallBTreeMap<K, V, N>
569where
570    K: Ord + Borrow<Q>,
571    Q: Ord + ?Sized,
572{
573    fn index_mut(&mut self, key: &Q) -> &mut Self::Output {
574        self.get_mut(key).expect("no entry found for key")
575    }
576}
577
578#[cfg(test)]
579mod tests {
580    use super::*;
581    use core::cmp::Ordering;
582
583    #[test]
584    fn test_btree_stack_ops_sorted_order() {
585        let mut map: SmallBTreeMap<i32, i32, 4> = SmallBTreeMap::new();
586        map.insert(3, 30);
587        map.insert(1, 10);
588        map.insert(2, 20);
589
590        let keys: Vec<_> = map.iter().map(|(k, _)| *k).collect();
591        assert_eq!(keys, vec![1, 2, 3]);
592    }
593
594    #[test]
595    fn test_btree_spill_trigger_on_insert() {
596        let mut map: SmallBTreeMap<i32, i32, 2> = SmallBTreeMap::new();
597        map.insert(1, 10);
598        map.insert(2, 20);
599        assert!(map.is_on_stack());
600
601        map.insert(0, 0); // Spill
602        assert!(!map.is_on_stack());
603
604        let keys: Vec<_> = map.iter().map(|(k, _)| *k).collect();
605        assert_eq!(keys, vec![0, 1, 2]);
606    }
607
608    #[test]
609    fn test_btree_any_storage_get_mut() {
610        let mut map: SmallBTreeMap<i32, i32, 4> = SmallBTreeMap::new();
611        map.insert(1, 10);
612        if let Some(v) = map.get_mut(&1) {
613            *v = 20;
614        }
615        assert_eq!(map.get(&1), Some(&20));
616
617        map.insert(2, 200);
618        map.insert(3, 300);
619        map.insert(4, 400);
620        map.insert(5, 500); // Spill
621        assert!(!map.is_on_stack());
622
623        if let Some(v) = map.get_mut(&5) {
624            *v = 555;
625        }
626        assert_eq!(map.get(&5), Some(&555));
627    }
628
629    #[test]
630    fn test_btree_any_storage_remove() {
631        let mut map: SmallBTreeMap<i32, i32, 4> = SmallBTreeMap::new();
632        map.insert(1, 10);
633        map.insert(2, 20);
634        assert_eq!(map.remove(&1), Some(10));
635        assert_eq!(map.len(), 1);
636        assert!(map.get(&1).is_none());
637
638        map.insert(3, 30);
639        map.insert(4, 40);
640        map.insert(5, 50); // Spill
641        assert_eq!(map.remove(&5), Some(50));
642        assert_eq!(map.len(), 3);
643    }
644
645    #[test]
646    fn test_btree_any_storage_clear() {
647        let mut map: SmallBTreeMap<i32, i32, 4> = SmallBTreeMap::new();
648        map.insert(1, 10);
649        map.clear();
650        assert!(map.is_empty());
651        assert_eq!(map.len(), 0);
652
653        for i in 0..5 {
654            map.insert(i, i * 10);
655        }
656        assert!(!map.is_on_stack());
657        map.clear();
658        assert!(map.is_empty());
659        assert!(!map.is_on_stack()); // Stays allocated on heap (conceptually empty)
660    }
661
662    #[test]
663    fn test_btree_traits_exhaustive() {
664        let mut map: SmallBTreeMap<i32, i32, 4> =
665            SmallBTreeMap::from_iter([(1, 10), (3, 30), (2, 20)]);
666        assert_eq!(map.len(), 3);
667        assert_eq!(map.get(&2), Some(&20));
668
669        // Update
670        map.insert(2, 22);
671        assert_eq!(map.get(&2), Some(&22));
672
673        // Remove
674        assert_eq!(map.remove(&2), Some(22));
675        assert_eq!(map.len(), 2);
676        assert!(map.get(&2).is_none());
677
678        // Clone
679        let cloned = map.clone();
680        assert_eq!(cloned.len(), 2);
681
682        // Debug
683        let debug = format!("{:?}", map);
684        assert!(debug.contains("1: 10"));
685        assert!(debug.contains("3: 30"));
686
687        // FromIterator
688        let collected: SmallBTreeMap<i32, i32, 4> = vec![(1, 10), (2, 20)].into_iter().collect();
689        assert_eq!(collected.len(), 2);
690
691        // Extend
692        let mut map2 = SmallBTreeMap::<i32, i32, 4>::new();
693        map2.extend(vec![(1, 10), (2, 20)]);
694        assert_eq!(map2.len(), 2);
695
696        // IntoIterator (Stack)
697        let vec: Vec<_> = map2.into_iter().collect();
698        assert_eq!(vec.len(), 2);
699        assert_eq!(vec[0], (1, 10));
700        assert_eq!(vec[1], (2, 20));
701
702        // Spill and traits
703        let mut map_spill: SmallBTreeMap<i32, i32, 2> = SmallBTreeMap::new();
704        map_spill.insert(1, 10);
705        map_spill.insert(2, 20);
706        map_spill.insert(3, 30); // Spill
707        assert!(!map_spill.is_on_stack());
708
709        let cloned_heap = map_spill.clone();
710        assert_eq!(cloned_heap.len(), 3);
711
712        let debug_heap = format!("{:?}", map_spill);
713        assert!(debug_heap.contains("1: 10"));
714
715        // IntoIterator (Heap)
716        let vec_heap: Vec<_> = map_spill.into_iter().collect();
717        assert_eq!(vec_heap.len(), 3);
718    }
719
720    #[test]
721    fn test_btree_entry_edge_cases() {
722        // Entry comparisons
723        let e1 = Entry(1, 10);
724        let e2 = Entry(1, 20);
725        let e3 = Entry(2, 10);
726        assert_eq!(e1, e2);
727        assert!(e1 < e3);
728        assert_eq!(e1.cmp(&e2), Ordering::Equal);
729        assert_eq!(e1.partial_cmp(&e3), Some(Ordering::Less));
730    }
731
732    #[test]
733    fn test_btree_any_storage_remove_non_existent() {
734        let mut map: SmallBTreeMap<i32, i32, 2> = SmallBTreeMap::new();
735        assert_eq!(map.remove(&1), None);
736        map.insert(1, 10);
737        map.insert(2, 20);
738        map.insert(3, 30); // Spill
739        assert_eq!(map.remove(&4), None);
740        assert_eq!(map.remove(&3), Some(30));
741    }
742
743    #[test]
744    fn test_btree_traits_into_iter_empty() {
745        let map_empty: SmallBTreeMap<i32, i32, 4> = SmallBTreeMap::new();
746        let mut it = map_empty.into_iter();
747        assert_eq!(it.next(), None);
748    }
749
750    #[test]
751    fn test_btree_any_storage_heap_manipulation() {
752        let mut map: SmallBTreeMap<i32, i32, 2> =
753            vec![(1, 10), (2, 20), (3, 30)].into_iter().collect();
754        assert!(!map.is_on_stack());
755
756        // get_mut heap
757        if let Some(v) = map.get_mut(&1) {
758            *v = 11;
759        }
760        assert_eq!(map[&1], 11);
761
762        // clear heap
763        map.clear();
764        assert!(map.is_empty());
765        assert!(!map.is_on_stack());
766    }
767
768    #[test]
769    fn test_btree_stack_binary_search_order() {
770        let mut s: SmallBTreeMap<i32, i32, 4> = SmallBTreeMap::new();
771        s.insert(2, 2);
772        s.insert(1, 1);
773        assert_eq!(s.get(&1), Some(&1));
774    }
775
776    #[test]
777    fn test_btree_map_traits_comparison() {
778        let map1: SmallBTreeMap<i32, i32, 2> = vec![(1, 10), (2, 20)].into_iter().collect();
779        let map2: SmallBTreeMap<i32, i32, 2> = vec![(1, 10), (2, 20)].into_iter().collect();
780        let map3: SmallBTreeMap<i32, i32, 2> = vec![(1, 10), (3, 30)].into_iter().collect();
781
782        // PartialEq
783        assert_eq!(map1, map2);
784        assert_ne!(map1, map3);
785
786        // PartialOrd / Ord
787        assert!(map1 < map3);
788        assert!(map3 > map1);
789
790        // Spill vs Stack Comparison
791        let mut map4: SmallBTreeMap<i32, i32, 2> = vec![(1, 10), (2, 20)].into_iter().collect();
792        map4.insert(3, 30); // Spill
793        assert_ne!(map1, map4);
794        assert!(map1 < map4);
795    }
796
797    #[test]
798    fn test_btree_map_traits_interop() {
799        let mut map: SmallBTreeMap<i32, i32, 2> = SmallBTreeMap::new();
800        map.insert(1, 10);
801        map.insert(2, 20);
802
803        let mut std_map = std::collections::BTreeMap::new();
804        std_map.insert(1, 10);
805        std_map.insert(2, 20);
806
807        assert_eq!(map, std_map);
808    }
809
810    #[test]
811    fn test_btree_map_traits_any_btreemap_interop() {
812        fn check_any<M: AnyBTreeMap<i32, i32>>(map: &M) {
813            assert_eq!(map.len(), 2);
814            assert_eq!(map.get(&1), Some(&10));
815        }
816
817        let map: SmallBTreeMap<i32, i32, 2> = vec![(1, 10), (2, 20)].into_iter().collect();
818        check_any(&map);
819    }
820}
821
822#[cfg(test)]
823mod btree_map_coverage_tests {
824    use super::*;
825
826    fn run_any_btree_map_test<M: AnyBTreeMap<i32, i32>>(any_map: &mut M) {
827        assert_eq!(any_map.len(), 0);
828        assert!(any_map.is_empty());
829        any_map.insert(1, 10);
830        assert_eq!(any_map.get(&1), Some(&10));
831        assert_eq!(any_map.get_mut(&1), Some(&mut 10));
832        assert_eq!(any_map.remove(&1), Some(10));
833        any_map.insert(2, 20);
834        any_map.clear();
835        assert_eq!(any_map.len(), 0);
836    }
837
838    #[test]
839    fn test_any_btree_map_trait_impls() {
840        let mut std_map: std::collections::BTreeMap<i32, i32> = std::collections::BTreeMap::new();
841        run_any_btree_map_test(&mut std_map);
842
843        let mut hl_map: HeaplessBTreeMap<i32, i32, 2> = HeaplessBTreeMap::new();
844        run_any_btree_map_test(&mut hl_map);
845
846        let mut small_map: SmallBTreeMap<i32, i32, 2> = SmallBTreeMap::new();
847        run_any_btree_map_test(&mut small_map);
848    }
849
850    #[test]
851    fn test_small_btree_map_with_capacity_heap() {
852        // cap > N creates heap directly
853        let map: SmallBTreeMap<i32, i32, 2> = SmallBTreeMap::with_capacity(3);
854        assert!(!map.is_on_stack());
855        assert_eq!(map.len(), 0);
856    }
857
858    #[test]
859    fn test_small_btree_map_insert_heap_replace() {
860        let mut map: SmallBTreeMap<i32, i32, 2> = SmallBTreeMap::new();
861        map.insert(1, 10);
862        map.insert(2, 20);
863        map.insert(3, 30); // spill
864
865        // replace existing key on heap
866        let old = map.insert(2, 200);
867        assert_eq!(old, Some(20));
868        assert_eq!(map.len(), 3);
869    }
870
871    #[test]
872    #[should_panic(expected = "no entry found for key")]
873    fn test_small_btree_map_index_mut_panic() {
874        let mut map: SmallBTreeMap<i32, i32, 2> = SmallBTreeMap::new();
875        let _val = &mut map[&1];
876    }
877}
small_collections/maps/btree_map.rs

small_collections/maps/
btree_map.rs
