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    V: Default,
435{
436    #[inline]
437    pub fn or_insert(self, default: V) -> &'a mut V {
438        match self {
439            Entry::Occupied(entry) => entry.into_mut(),
440            Entry::Vacant(entry) => entry.insert(default),
441        }
442    }
443
444    #[inline]
445    pub fn or_insert_with<F: FnOnce() -> V>(self, default_fn: F) -> &'a mut V {
446        match self {
447            Entry::Occupied(entry) => entry.into_mut(),
448            Entry::Vacant(entry) => entry.insert(default_fn()),
449        }
450    }
451
452    #[inline]
453    pub fn and_modify<F>(mut self, f: F) -> Self
454    where
455        F: FnOnce(&mut V),
456    {
457        if let Entry::Occupied(ref mut entry) = self {
458            f(entry.get_mut());
459        }
460        self
461    }
462
463    #[inline]
464    pub fn key(&self) -> &K {
465        match self {
466            Entry::Occupied(entry) => entry.key(),
467            Entry::Vacant(entry) => entry.key(),
468        }
469    }
470
471    #[inline]
472    pub fn insert_entry(self, value: V) -> OccupiedEntry<'a, K, V> {
473        match self {
474            Entry::Occupied(mut ent) => {
475                match ent {
476                    OccupiedEntry::One(ref mut o) => {
477                        let (k, _old_v) = o.take().unwrap();
478                        o.replace((k, value));
479                    }
480                    OccupiedEntry::Multi(ref mut ent) => {
481                        ent.insert(value);
482                    }
483                }
484                ent
485            }
486            Entry::Vacant(VacantEntry::One(entry)) => {
487                let mut _value = MaybeUninit::new(value);
488                let mut _key = MaybeUninit::new(entry.key);
489                let map = entry.map;
490                if let VariousMap::One(item) = map {
491                    if item.is_none() {
492                        unsafe {
493                            item.replace((_key.assume_init_read(), _value.assume_init_read()));
494                        }
495                    } else {
496                        let (old_k, old_v) = item.take().unwrap();
497                        unsafe {
498                            if &old_k == _key.assume_init_ref() {
499                                item.replace((_key.assume_init_read(), _value.assume_init_read()));
500                            } else {
501                                let _ = item;
502                                let mut _map = BTreeMap::new();
503                                _map.insert(old_k, old_v);
504                                *map = VariousMap::Multi(_map);
505                            }
506                        }
507                    }
508                }
509                match map {
510                    VariousMap::One(o) => OccupiedEntry::One(o),
511                    VariousMap::Multi(map) => {
512                        let ent = unsafe {
513                            map.entry(_key.assume_init_read())
514                                .insert_entry(_value.assume_init_read())
515                        };
516                        OccupiedEntry::Multi(ent)
517                    }
518                }
519            }
520            Entry::Vacant(VacantEntry::Multi(ent)) => OccupiedEntry::Multi(ent.insert_entry(value)),
521        }
522    }
523
524    #[inline]
525    pub fn or_default(self) -> &'a mut V
526    where
527        V: Default,
528    {
529        self.or_insert_with(Default::default)
530    }
531}
532
533#[cfg(test)]
534mod tests {
535    use super::*;
536
537    #[test]
538    fn test_new() {
539        let map: VariousMap<i32, String> = VariousMap::new();
540        assert!(matches!(map, VariousMap::One(None)));
541    }
542
543    #[test]
544    fn test_insert_and_get_one() {
545        let mut map = VariousMap::<usize, usize>::new();
546        assert_eq!(map.insert(1, 1), None);
547        assert!(matches!(map, VariousMap::One(Some((ref k, ref v))) if *k == 1 && *v == 1));
548        assert_eq!(map.get(&1), Some(&1));
549        assert_eq!(map.get(&2), None);
550        **(map.get_mut(&1).as_mut().unwrap()) = 2;
551        assert_eq!(map.get(&1), Some(&2));
552        assert_eq!(map.insert(1, 3), Some(2));
553        assert_eq!(map.get(&1), Some(&3));
554    }
555
556    #[test]
557    fn test_insert_transition_to_multi() {
558        let mut map = VariousMap::new();
559        map.insert(1, "one".to_string());
560        assert_eq!(map.insert(2, "two".to_string()), None);
561
562        // After inserting a second element, it should transition to BTreeMap
563        match map {
564            VariousMap::Multi(ref btree_map) => {
565                assert_eq!(btree_map.len(), 2);
566                assert_eq!(btree_map.get(&1), Some(&"one".to_string()));
567                assert_eq!(btree_map.get(&2), Some(&"two".to_string()));
568            }
569            _ => panic!("Expected Multi variant after inserting two elements"),
570        }
571
572        assert_eq!(map.get(&1), Some(&"one".to_string()));
573        assert_eq!(map.get(&2), Some(&"two".to_string()));
574        assert_eq!(map.get(&3), None);
575
576        **map.get_mut(&2).as_mut().unwrap() = "two_update".to_string();
577        assert_eq!(map.get(&2), Some(&"two_update".to_string()));
578    }
579
580    #[test]
581    fn test_iter_one() {
582        let mut map = VariousMap::new();
583        map.insert(1, "one".to_string());
584        let mut collected: Vec<_> = map.iter().collect();
585        collected.sort_by_key(|(k, _)| *k);
586        assert_eq!(collected, vec![(&1, &"one".to_string())]);
587    }
588
589    #[test]
590    fn test_iter_multi() {
591        let mut map = VariousMap::new();
592        map.insert(1, "one".to_string());
593        map.insert(2, "two".to_string()); // Transition to Multi
594        map.insert(3, "three".to_string());
595
596        let mut collected: Vec<_> = map.iter().collect();
597        collected.sort_by_key(|(k, _)| *k);
598        assert_eq!(
599            collected,
600            vec![(&1, &"one".to_string()), (&2, &"two".to_string()), (&3, &"three".to_string())]
601        );
602    }
603
604    #[test]
605    fn test_into_iter_one() {
606        let mut map = VariousMap::new();
607        map.insert(1, "one".to_string());
608        let mut collected: Vec<_> = map.into_iter().collect();
609        collected.sort_by_key(|(k, _)| *k);
610        assert_eq!(collected, vec![(1, "one".to_string())]);
611    }
612
613    #[test]
614    fn test_contains_key() {
615        let mut map = VariousMap::new();
616        map.insert(1, "one".to_string());
617        assert!(map.contains_key(&1));
618        assert!(!map.contains_key(&2));
619
620        map.insert(2, "two".to_string()); // Transition to Multi
621        assert!(map.contains_key(&1));
622        assert!(map.contains_key(&2));
623        assert!(!map.contains_key(&3));
624    }
625
626    #[test]
627    fn test_entry_api_one() {
628        let mut map = VariousMap::new();
629        map.insert(1, "one".to_string());
630
631        // Occupied
632        match map.entry(1) {
633            Entry::Occupied(ent) => {
634                assert_eq!(ent.get(), &"one".to_string());
635                ent.insert("one_updated".to_string());
636            }
637            Entry::Vacant(_) => panic!("Should be occupied"),
638        }
639        assert_eq!(map.get(&1), Some(&"one_updated".to_string()));
640
641        // Vacant
642        match map.entry(2) {
643            Entry::Occupied(_) => panic!("Should be vacant"),
644            Entry::Vacant(ent) => {
645                ent.insert("two".to_string());
646            }
647        }
648        assert!(map.contains_key(&2));
649    }
650
651    #[test]
652    fn test_entry_api_or_insert() {
653        let mut map = VariousMap::new();
654        map.entry(1).or_insert("one".to_string());
655
656        // Occupied
657        let v = map
658            .entry(1)
659            .and_modify(|v| {
660                *v = "one_3".to_string();
661            })
662            .or_insert("one_2".to_string());
663        assert_eq!(v, &"one_3".to_string());
664
665        // Vacant
666        map.entry(2).or_insert("two".to_string());
667        assert_eq!(map.get(&2), Some(&"two".to_string()));
668    }
669
670    #[test]
671    fn test_entry_api_insert_entry() {
672        let mut map = VariousMap::new();
673
674        // Vacant -> insert_entry
675        let occupied = map.entry(1).insert_entry("one".to_string());
676        assert_eq!(occupied.get(), &"one".to_string());
677
678        // Occupied -> insert_entry
679        let occupied = map.entry(1).insert_entry("one_updated".to_string());
680        assert_eq!(occupied.get(), &"one_updated".to_string());
681        assert_eq!(map.get(&1), Some(&"one_updated".to_string()));
682    }
683}
embed_collections/various_map.rs

embed_collections/
various_map.rs
