embed_collections/
various_map.rs

1//! VariousMap: Provide a tempoary map optimise for empty or one condition,
2//! to delay allocation.
3//!
4//! Initial to be Option<(K, V)>.
5//! If multi item inserted, transit from Option to std::collections::BTreeMap.
6
7use alloc::collections::BTreeMap;
8use alloc::collections::btree_map;
9use core::borrow::Borrow;
10use core::mem::MaybeUninit;
11use core::option;
12
13/// A tempoary map optimise for empty or one condition,
14/// to delay allocation.
15///
16/// Initial to be Option<(K, V)>.
17/// If multi item inserted, transit from Option to std::collections::BTreeMap.
18pub enum VariousMap<K, V> {
19    One(Option<(K, V)>),
20    Multi(BTreeMap<K, V>),
21}
22
23impl<K: Ord, V> VariousMap<K, V> {
24    #[inline]
25    pub fn new() -> Self {
26        Self::One(None)
27    }
28
29    #[inline]
30    pub fn get<Q>(&self, key: &Q) -> Option<&V>
31    where
32        K: Borrow<Q> + Ord,
33        Q: Ord + ?Sized,
34    {
35        match self {
36            Self::One(Some(item)) => {
37                if item.0.borrow() == key {
38                    Some(&item.1)
39                } else {
40                    None
41                }
42            }
43            Self::Multi(map) => map.get(key),
44            _ => None,
45        }
46    }
47
48    #[inline]
49    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
50    where
51        K: Borrow<Q> + Ord,
52        Q: Ord + ?Sized,
53    {
54        match self {
55            Self::One(Some(item)) => {
56                if item.0.borrow() == key {
57                    Some(&mut item.1)
58                } else {
59                    None
60                }
61            }
62            Self::Multi(map) => map.get_mut(key),
63            _ => None,
64        }
65    }
66
67    #[inline]
68    pub fn insert(&mut self, key: K, value: V) -> Option<V>
69    where
70        K: Ord,
71    {
72        let (old_k, old_v) = match self {
73            Self::One(item) => {
74                if item.is_none() {
75                    item.replace((key, value));
76                    return None;
77                }
78                let (old_k, old_v) = item.take().unwrap();
79                if old_k == key {
80                    item.replace((key, value));
81                    return Some(old_v);
82                }
83                (old_k, old_v)
84            }
85            Self::Multi(map) => {
86                return map.insert(key, value);
87            }
88        };
89        let mut map = BTreeMap::new();
90        map.insert(key, value);
91        map.insert(old_k, old_v);
92        *self = Self::Multi(map);
93        None
94    }
95
96    #[inline]
97    pub fn iter(&self) -> Iter<'_, K, V> {
98        match self {
99            Self::One(o) => Iter::One(o.iter()),
100            Self::Multi(map) => Iter::Multi(map.iter()),
101        }
102    }
103
104    #[inline]
105    pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
106        let mut exists = false;
107        match self {
108            Self::One(Some(item)) => {
109                if item.0 == key {
110                    exists = true;
111                }
112            }
113            Self::Multi(map) => match map.entry(key) {
114                btree_map::Entry::Occupied(ent) => {
115                    return Entry::Occupied(OccupiedEntry::Multi(ent));
116                }
117                btree_map::Entry::Vacant(ent) => return Entry::Vacant(VacantEntry::Multi(ent)),
118            },
119            _ => {}
120        }
121        if !exists {
122            return Entry::Vacant(VacantEntry::One(VacantEntryOne { key, map: self }));
123        }
124        // in order to resolve borrow issue,
125        // return field ref only after returning self ref.
126        if let Self::One(item) = self {
127            return Entry::Occupied(OccupiedEntry::One(item));
128        }
129        unreachable!();
130    }
131
132    #[inline]
133    pub fn len(&self) -> usize {
134        match self {
135            Self::One(Some(_item)) => 1,
136            Self::Multi(map) => map.len(),
137            _ => 0,
138        }
139    }
140
141    #[inline]
142    pub fn is_empty(&self) -> bool {
143        match self {
144            Self::One(item) => item.is_none(),
145            Self::Multi(map) => map.is_empty(),
146        }
147    }
148
149    #[inline]
150    pub fn contains_key<Q>(&self, key: &Q) -> bool
151    where
152        K: Borrow<Q> + Ord,
153        Q: Ord + ?Sized,
154    {
155        match self {
156            Self::One(Some(item)) => item.0.borrow() == key,
157            Self::Multi(map) => map.contains_key(key),
158            _ => false,
159        }
160    }
161}
162
163impl<K, V> IntoIterator for VariousMap<K, V> {
164    type Item = (K, V);
165    type IntoIter = IntoIter<K, V>;
166
167    #[inline]
168    fn into_iter(self) -> Self::IntoIter {
169        match self {
170            Self::One(o) => IntoIter::One(o),
171            Self::Multi(map) => IntoIter::Multi(map.into_iter()),
172        }
173    }
174}
175
176pub enum Iter<'a, K, V> {
177    One(option::Iter<'a, (K, V)>),
178    Multi(btree_map::Iter<'a, K, V>),
179}
180
181impl<'a, K, V> Iterator for Iter<'a, K, V> {
182    type Item = (&'a K, &'a V);
183
184    #[inline]
185    fn next(&mut self) -> Option<Self::Item> {
186        match self {
187            Self::One(iter) => {
188                if let Some(item) = iter.next() {
189                    Some((&item.0, &item.1))
190                } else {
191                    None
192                }
193            }
194            Self::Multi(iter) => iter.next(),
195        }
196    }
197
198    #[inline]
199    fn size_hint(&self) -> (usize, Option<usize>) {
200        match self {
201            Self::One(iter) => {
202                let l = iter.len();
203                (l, Some(l))
204            }
205            Self::Multi(iter) => {
206                let l = iter.len();
207                (l, Some(l))
208            }
209        }
210    }
211}
212
213impl<'a, K, V> ExactSizeIterator for Iter<'a, K, V> {
214    #[inline]
215    fn len(&self) -> usize {
216        match self {
217            Self::One(iter) => iter.len(),
218            Self::Multi(iter) => iter.len(),
219        }
220    }
221}
222
223impl<'a, K, V> DoubleEndedIterator for Iter<'a, K, V> {
224    #[inline]
225    fn next_back(&mut self) -> Option<Self::Item> {
226        match self {
227            Self::One(iter) => {
228                if let Some(item) = iter.next_back() {
229                    Some((&item.0, &item.1))
230                } else {
231                    None
232                }
233            }
234            Self::Multi(iter) => iter.next_back(),
235        }
236    }
237}
238
239pub enum IntoIter<K, V> {
240    One(Option<(K, V)>),
241    Multi(btree_map::IntoIter<K, V>),
242}
243
244impl<K, V> Iterator for IntoIter<K, V> {
245    type Item = (K, V);
246
247    #[inline]
248    fn next(&mut self) -> Option<Self::Item> {
249        match self {
250            Self::One(iter) => iter.take(),
251            Self::Multi(iter) => iter.next(),
252        }
253    }
254}
255
256impl<K, V> ExactSizeIterator for IntoIter<K, V> {
257    #[inline]
258    fn len(&self) -> usize {
259        match self {
260            Self::One(iter) => {
261                if iter.is_some() {
262                    1
263                } else {
264                    0
265                }
266            }
267            Self::Multi(iter) => iter.len(),
268        }
269    }
270}
271
272impl<K, V> DoubleEndedIterator for IntoIter<K, V> {
273    #[inline]
274    fn next_back(&mut self) -> Option<Self::Item> {
275        match self {
276            Self::One(iter) => iter.take(),
277            Self::Multi(iter) => iter.next_back(),
278        }
279    }
280}
281
282pub enum Entry<'a, K: 'a, V: 'a> {
283    Occupied(OccupiedEntry<'a, K, V>),
284    Vacant(VacantEntry<'a, K, V>),
285}
286
287pub enum OccupiedEntry<'a, K: 'a, V: 'a> {
288    One(&'a mut Option<(K, V)>),
289    Multi(btree_map::OccupiedEntry<'a, K, V>),
290}
291
292impl<'a, K, V> OccupiedEntry<'a, K, V>
293where
294    K: Ord,
295{
296    #[inline]
297    pub fn get(&self) -> &V {
298        match self {
299            Self::One(o) => &o.as_ref().unwrap().1,
300            Self::Multi(ent) => ent.get(),
301        }
302    }
303
304    #[inline]
305    pub fn get_mut(&mut self) -> &mut V {
306        match self {
307            Self::One(o) => &mut o.as_mut().unwrap().1,
308            Self::Multi(ent) => ent.get_mut(),
309        }
310    }
311
312    #[inline]
313    pub fn key(&self) -> &K {
314        match self {
315            Self::One(o) => &o.as_ref().unwrap().0,
316            Self::Multi(ent) => ent.key(),
317        }
318    }
319
320    #[inline]
321    pub fn remove(self) -> V {
322        match self {
323            Self::One(o) => {
324                let (_k, v) = o.take().unwrap();
325                v
326            }
327            Self::Multi(ent) => ent.remove(),
328        }
329    }
330
331    #[inline]
332    pub fn remove_entry(self) -> (K, V) {
333        match self {
334            Self::One(o) => o.take().unwrap(),
335            Self::Multi(ent) => ent.remove_entry(),
336        }
337    }
338
339    #[inline]
340    pub fn insert(self, value: V) -> V {
341        match self {
342            Self::One(o) => {
343                let (k, old_v) = o.take().unwrap();
344                o.replace((k, value));
345                old_v
346            }
347            Self::Multi(mut ent) => ent.insert(value),
348        }
349    }
350
351    #[inline]
352    pub fn into_mut(self) -> &'a mut V {
353        match self {
354            Self::One(o) => &mut o.as_mut().unwrap().1,
355            Self::Multi(ent) => ent.into_mut(),
356        }
357    }
358}
359
360pub enum VacantEntry<'a, K, V> {
361    One(VacantEntryOne<'a, K, V>),
362    Multi(btree_map::VacantEntry<'a, K, V>),
363}
364
365struct VacantEntryOne<'a, K: 'a, V: 'a> {
366    pub(crate) key: K,                        // Owned key to insert
367    pub(crate) map: &'a mut VariousMap<K, V>, // Reference to the VariousMap
368}
369
370impl<'a, K, V> VacantEntry<'a, K, V>
371where
372    K: Ord,
373{
374    #[inline]
375    pub fn key(&self) -> &K {
376        match self {
377            Self::One(ent) => &ent.key,
378            Self::Multi(ent) => ent.key(),
379        }
380    }
381
382    #[inline]
383    pub fn into_key(self) -> K {
384        match self {
385            Self::One(ent) => ent.key,
386            Self::Multi(ent) => ent.into_key(),
387        }
388    }
389
390    #[inline]
391    pub fn insert(self, value: V) -> &'a mut V {
392        match self {
393            Self::One(ent) => {
394                let mut _value = MaybeUninit::new(value);
395                let mut _key = MaybeUninit::new(ent.key);
396                let map = ent.map;
397                if let VariousMap::One(item) = map {
398                    if item.is_none() {
399                        unsafe {
400                            // we should have return here, but don't because of the borrow checker
401                            item.replace((_key.assume_init_read(), _value.assume_init_read()));
402                        }
403                    } else {
404                        let (old_k, old_v) = item.take().unwrap();
405                        unsafe {
406                            if &old_k == _key.assume_init_ref() {
407                                // we should have return here, but don't because of the borrow checker
408                                item.replace((_key.assume_init_read(), _value.assume_init_read()));
409                            } else {
410                                let _ = item;
411                                let mut _map = BTreeMap::new();
412                                _map.insert(old_k, old_v);
413                                *map = VariousMap::Multi(_map);
414                            }
415                        }
416                    }
417                }
418                match map {
419                    VariousMap::One(Some(item)) => &mut item.1,
420                    VariousMap::Multi(map) => unsafe {
421                        map.entry(_key.assume_init_read()).or_insert(_value.assume_init_read())
422                    },
423                    _ => unreachable!(),
424                }
425            }
426            Self::Multi(ent) => ent.insert(value),
427        }
428    }
429}
430
431impl<'a, K, V> Entry<'a, K, V>
432where
433    K: Ord,
434{
435    #[inline]
436    pub fn or_insert(self, default: V) -> &'a mut V {
437        match self {
438            Entry::Occupied(entry) => entry.into_mut(),
439            Entry::Vacant(entry) => entry.insert(default),
440        }
441    }
442
443    #[inline]
444    pub fn or_insert_with<F: FnOnce() -> V>(self, default_fn: F) -> &'a mut V {
445        match self {
446            Entry::Occupied(entry) => entry.into_mut(),
447            Entry::Vacant(entry) => entry.insert(default_fn()),
448        }
449    }
450
451    #[inline]
452    pub fn and_modify<F>(mut self, f: F) -> Self
453    where
454        F: FnOnce(&mut V),
455    {
456        if let Entry::Occupied(ref mut entry) = self {
457            f(entry.get_mut());
458        }
459        self
460    }
461
462    #[inline]
463    pub fn key(&self) -> &K {
464        match self {
465            Entry::Occupied(entry) => entry.key(),
466            Entry::Vacant(entry) => entry.key(),
467        }
468    }
469
470    #[inline]
471    pub fn insert_entry(self, value: V) -> OccupiedEntry<'a, K, V> {
472        match self {
473            Entry::Occupied(mut ent) => {
474                match ent {
475                    OccupiedEntry::One(ref mut o) => {
476                        let (k, _old_v) = o.take().unwrap();
477                        o.replace((k, value));
478                    }
479                    OccupiedEntry::Multi(ref mut ent) => {
480                        ent.insert(value);
481                    }
482                }
483                ent
484            }
485            Entry::Vacant(VacantEntry::One(entry)) => {
486                let mut _value = MaybeUninit::new(value);
487                let mut _key = MaybeUninit::new(entry.key);
488                let map = entry.map;
489                if let VariousMap::One(item) = map {
490                    if item.is_none() {
491                        unsafe {
492                            item.replace((_key.assume_init_read(), _value.assume_init_read()));
493                        }
494                    } else {
495                        let (old_k, old_v) = item.take().unwrap();
496                        unsafe {
497                            if &old_k == _key.assume_init_ref() {
498                                item.replace((_key.assume_init_read(), _value.assume_init_read()));
499                            } else {
500                                let _ = item;
501                                let mut _map = BTreeMap::new();
502                                _map.insert(old_k, old_v);
503                                *map = VariousMap::Multi(_map);
504                            }
505                        }
506                    }
507                }
508                match map {
509                    VariousMap::One(o) => OccupiedEntry::One(o),
510                    VariousMap::Multi(map) => {
511                        let ent = unsafe {
512                            map.entry(_key.assume_init_read())
513                                .insert_entry(_value.assume_init_read())
514                        };
515                        OccupiedEntry::Multi(ent)
516                    }
517                }
518            }
519            Entry::Vacant(VacantEntry::Multi(ent)) => OccupiedEntry::Multi(ent.insert_entry(value)),
520        }
521    }
522
523    #[inline]
524    pub fn or_default(self) -> &'a mut V
525    where
526        V: Default,
527    {
528        self.or_insert_with(Default::default)
529    }
530}
531
532#[cfg(test)]
533mod tests {
534    use super::*;
535
536    #[test]
537    fn test_new() {
538        let map: VariousMap<i32, String> = VariousMap::new();
539        assert!(matches!(map, VariousMap::One(None)));
540    }
541
542    #[test]
543    fn test_insert_and_get_one() {
544        let mut map = VariousMap::<usize, usize>::new();
545        assert_eq!(map.insert(1, 1), None);
546        assert!(matches!(map, VariousMap::One(Some((ref k, ref v))) if *k == 1 && *v == 1));
547        assert_eq!(map.get(&1), Some(&1));
548        assert_eq!(map.get(&2), None);
549        **(map.get_mut(&1).as_mut().unwrap()) = 2;
550        assert_eq!(map.get(&1), Some(&2));
551        assert_eq!(map.insert(1, 3), Some(2));
552        assert_eq!(map.get(&1), Some(&3));
553    }
554
555    #[test]
556    fn test_insert_transition_to_multi() {
557        let mut map = VariousMap::new();
558        map.insert(1, "one".to_string());
559        assert_eq!(map.insert(2, "two".to_string()), None);
560
561        // After inserting a second element, it should transition to BTreeMap
562        match map {
563            VariousMap::Multi(ref btree_map) => {
564                assert_eq!(btree_map.len(), 2);
565                assert_eq!(btree_map.get(&1), Some(&"one".to_string()));
566                assert_eq!(btree_map.get(&2), Some(&"two".to_string()));
567            }
568            _ => panic!("Expected Multi variant after inserting two elements"),
569        }
570
571        assert_eq!(map.get(&1), Some(&"one".to_string()));
572        assert_eq!(map.get(&2), Some(&"two".to_string()));
573        assert_eq!(map.get(&3), None);
574
575        **map.get_mut(&2).as_mut().unwrap() = "two_update".to_string();
576        assert_eq!(map.get(&2), Some(&"two_update".to_string()));
577    }
578
579    #[test]
580    fn test_iter_one() {
581        let mut map = VariousMap::new();
582        map.insert(1, "one".to_string());
583        let mut collected: Vec<_> = map.iter().collect();
584        collected.sort_by_key(|(k, _)| *k);
585        assert_eq!(collected, vec![(&1, &"one".to_string())]);
586    }
587
588    #[test]
589    fn test_iter_multi() {
590        let mut map = VariousMap::new();
591        map.insert(1, "one".to_string());
592        map.insert(2, "two".to_string()); // Transition to Multi
593        map.insert(3, "three".to_string());
594
595        let mut collected: Vec<_> = map.iter().collect();
596        collected.sort_by_key(|(k, _)| *k);
597        assert_eq!(
598            collected,
599            vec![(&1, &"one".to_string()), (&2, &"two".to_string()), (&3, &"three".to_string())]
600        );
601    }
602
603    #[test]
604    fn test_into_iter_one() {
605        let mut map = VariousMap::new();
606        map.insert(1, "one".to_string());
607        let mut collected: Vec<_> = map.into_iter().collect();
608        collected.sort_by_key(|(k, _)| *k);
609        assert_eq!(collected, vec![(1, "one".to_string())]);
610    }
611
612    #[test]
613    fn test_contains_key() {
614        let mut map = VariousMap::new();
615        map.insert(1, "one".to_string());
616        assert!(map.contains_key(&1));
617        assert!(!map.contains_key(&2));
618
619        map.insert(2, "two".to_string()); // Transition to Multi
620        assert!(map.contains_key(&1));
621        assert!(map.contains_key(&2));
622        assert!(!map.contains_key(&3));
623    }
624
625    #[test]
626    fn test_entry_api_one() {
627        let mut map = VariousMap::new();
628        map.insert(1, "one".to_string());
629
630        // Occupied
631        match map.entry(1) {
632            Entry::Occupied(ent) => {
633                assert_eq!(ent.get(), &"one".to_string());
634                ent.insert("one_updated".to_string());
635            }
636            Entry::Vacant(_) => panic!("Should be occupied"),
637        }
638        assert_eq!(map.get(&1), Some(&"one_updated".to_string()));
639
640        // Vacant
641        match map.entry(2) {
642            Entry::Occupied(_) => panic!("Should be vacant"),
643            Entry::Vacant(ent) => {
644                ent.insert("two".to_string());
645            }
646        }
647        assert!(map.contains_key(&2));
648    }
649
650    #[test]
651    fn test_entry_api_or_insert() {
652        let mut map = VariousMap::new();
653        map.entry(1).or_insert("one".to_string());
654
655        // Occupied
656        let v = map
657            .entry(1)
658            .and_modify(|v| {
659                *v = "one_3".to_string();
660            })
661            .or_insert("one_2".to_string());
662        assert_eq!(v, &"one_3".to_string());
663
664        // Vacant
665        map.entry(2).or_insert("two".to_string());
666        assert_eq!(map.get(&2), Some(&"two".to_string()));
667    }
668
669    #[test]
670    fn test_entry_api_insert_entry() {
671        let mut map = VariousMap::new();
672
673        // Vacant -> insert_entry
674        let occupied = map.entry(1).insert_entry("one".to_string());
675        assert_eq!(occupied.get(), &"one".to_string());
676
677        // Occupied -> insert_entry
678        let occupied = map.entry(1).insert_entry("one_updated".to_string());
679        assert_eq!(occupied.get(), &"one_updated".to_string());
680        assert_eq!(map.get(&1), Some(&"one_updated".to_string()));
681    }
682}
embed_collections/various_map.rs

embed_collections/
various_map.rs
