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 iter_mut(&mut self) -> IterMut<'_, K, V> {
106        match self {
107            Self::One(o) => IterMut::One(o.iter_mut()),
108            Self::Multi(map) => IterMut::Multi(map.iter_mut()),
109        }
110    }
111
112    #[inline]
113    pub fn keys(&self) -> Keys<'_, K, V> {
114        match self {
115            Self::One(o) => Keys::One(o.iter()),
116            Self::Multi(map) => Keys::Multi(map.keys()),
117        }
118    }
119
120    #[inline]
121    pub fn values(&self) -> Values<'_, K, V> {
122        match self {
123            Self::One(o) => Values::One(o.iter()),
124            Self::Multi(map) => Values::Multi(map.values()),
125        }
126    }
127
128    #[inline]
129    pub fn values_mut(&mut self) -> ValuesMut<'_, K, V> {
130        match self {
131            Self::One(o) => ValuesMut::One(o.iter_mut()),
132            Self::Multi(map) => ValuesMut::Multi(map.values_mut()),
133        }
134    }
135
136    #[inline]
137    pub fn entry(&mut self, key: K) -> Entry<'_, K, V> {
138        let mut exists = false;
139        match self {
140            Self::One(Some(item)) => {
141                if item.0 == key {
142                    exists = true;
143                }
144            }
145            Self::Multi(map) => match map.entry(key) {
146                btree_map::Entry::Occupied(ent) => {
147                    return Entry::Occupied(OccupiedEntry::Multi(ent));
148                }
149                btree_map::Entry::Vacant(ent) => return Entry::Vacant(VacantEntry::Multi(ent)),
150            },
151            _ => {}
152        }
153        if !exists {
154            return Entry::Vacant(VacantEntry::One(VacantEntryOne { key, map: self }));
155        }
156        // in order to resolve borrow issue,
157        // return field ref only after returning self ref.
158        if let Self::One(item) = self {
159            return Entry::Occupied(OccupiedEntry::One(item));
160        }
161        unreachable!();
162    }
163
164    #[inline]
165    pub fn len(&self) -> usize {
166        match self {
167            Self::One(Some(_item)) => 1,
168            Self::Multi(map) => map.len(),
169            _ => 0,
170        }
171    }
172
173    #[inline]
174    pub fn is_empty(&self) -> bool {
175        match self {
176            Self::One(item) => item.is_none(),
177            Self::Multi(map) => map.is_empty(),
178        }
179    }
180
181    #[inline]
182    pub fn contains_key<Q>(&self, key: &Q) -> bool
183    where
184        K: Borrow<Q> + Ord,
185        Q: Ord + ?Sized,
186    {
187        match self {
188            Self::One(Some(item)) => item.0.borrow() == key,
189            Self::Multi(map) => map.contains_key(key),
190            _ => false,
191        }
192    }
193}
194
195impl<K, V> IntoIterator for VariousMap<K, V> {
196    type Item = (K, V);
197    type IntoIter = IntoIter<K, V>;
198
199    #[inline]
200    fn into_iter(self) -> Self::IntoIter {
201        match self {
202            Self::One(o) => IntoIter::One(o),
203            Self::Multi(map) => IntoIter::Multi(map.into_iter()),
204        }
205    }
206}
207
208impl<'a, K: Ord, V> IntoIterator for &'a VariousMap<K, V> {
209    type Item = (&'a K, &'a V);
210    type IntoIter = Iter<'a, K, V>;
211
212    #[inline]
213    fn into_iter(self) -> Self::IntoIter {
214        self.iter()
215    }
216}
217
218pub enum Iter<'a, K, V> {
219    One(option::Iter<'a, (K, V)>),
220    Multi(btree_map::Iter<'a, K, V>),
221}
222
223impl<'a, K, V> Iterator for Iter<'a, K, V> {
224    type Item = (&'a K, &'a V);
225
226    #[inline]
227    fn next(&mut self) -> Option<Self::Item> {
228        match self {
229            Self::One(iter) => {
230                if let Some(item) = iter.next() {
231                    Some((&item.0, &item.1))
232                } else {
233                    None
234                }
235            }
236            Self::Multi(iter) => iter.next(),
237        }
238    }
239
240    #[inline]
241    fn size_hint(&self) -> (usize, Option<usize>) {
242        match self {
243            Self::One(iter) => {
244                let l = iter.len();
245                (l, Some(l))
246            }
247            Self::Multi(iter) => {
248                let l = iter.len();
249                (l, Some(l))
250            }
251        }
252    }
253}
254
255impl<'a, K, V> ExactSizeIterator for Iter<'a, K, V> {
256    #[inline]
257    fn len(&self) -> usize {
258        match self {
259            Self::One(iter) => iter.len(),
260            Self::Multi(iter) => iter.len(),
261        }
262    }
263}
264
265impl<'a, K, V> DoubleEndedIterator for Iter<'a, K, V> {
266    #[inline]
267    fn next_back(&mut self) -> Option<Self::Item> {
268        match self {
269            Self::One(iter) => {
270                if let Some(item) = iter.next_back() {
271                    Some((&item.0, &item.1))
272                } else {
273                    None
274                }
275            }
276            Self::Multi(iter) => iter.next_back(),
277        }
278    }
279}
280
281pub enum IterMut<'a, K, V> {
282    One(option::IterMut<'a, (K, V)>),
283    Multi(btree_map::IterMut<'a, K, V>),
284}
285
286impl<'a, K, V> Iterator for IterMut<'a, K, V> {
287    type Item = (&'a K, &'a mut V);
288
289    #[inline]
290    fn next(&mut self) -> Option<Self::Item> {
291        match self {
292            Self::One(iter) => {
293                if let Some(item) = iter.next() {
294                    Some((&item.0, &mut item.1))
295                } else {
296                    None
297                }
298            }
299            Self::Multi(iter) => iter.next(),
300        }
301    }
302
303    #[inline]
304    fn size_hint(&self) -> (usize, Option<usize>) {
305        match self {
306            Self::One(iter) => {
307                let l = iter.len();
308                (l, Some(l))
309            }
310            Self::Multi(iter) => {
311                let l = iter.len();
312                (l, Some(l))
313            }
314        }
315    }
316}
317
318impl<'a, K, V> ExactSizeIterator for IterMut<'a, K, V> {
319    #[inline]
320    fn len(&self) -> usize {
321        match self {
322            Self::One(iter) => iter.len(),
323            Self::Multi(iter) => iter.len(),
324        }
325    }
326}
327
328impl<'a, K, V> DoubleEndedIterator for IterMut<'a, K, V> {
329    #[inline]
330    fn next_back(&mut self) -> Option<Self::Item> {
331        match self {
332            Self::One(iter) => {
333                if let Some(item) = iter.next_back() {
334                    Some((&item.0, &mut item.1))
335                } else {
336                    None
337                }
338            }
339            Self::Multi(iter) => iter.next_back(),
340        }
341    }
342}
343
344pub enum IntoIter<K, V> {
345    One(Option<(K, V)>),
346    Multi(btree_map::IntoIter<K, V>),
347}
348
349impl<K, V> Iterator for IntoIter<K, V> {
350    type Item = (K, V);
351
352    #[inline]
353    fn next(&mut self) -> Option<Self::Item> {
354        match self {
355            Self::One(iter) => iter.take(),
356            Self::Multi(iter) => iter.next(),
357        }
358    }
359}
360
361impl<K, V> ExactSizeIterator for IntoIter<K, V> {
362    #[inline]
363    fn len(&self) -> usize {
364        match self {
365            Self::One(iter) => {
366                if iter.is_some() {
367                    1
368                } else {
369                    0
370                }
371            }
372            Self::Multi(iter) => iter.len(),
373        }
374    }
375}
376
377impl<K, V> DoubleEndedIterator for IntoIter<K, V> {
378    #[inline]
379    fn next_back(&mut self) -> Option<Self::Item> {
380        match self {
381            Self::One(iter) => iter.take(),
382            Self::Multi(iter) => iter.next_back(),
383        }
384    }
385}
386
387pub enum Keys<'a, K, V> {
388    One(option::Iter<'a, (K, V)>),
389    Multi(btree_map::Keys<'a, K, V>),
390}
391
392impl<'a, K, V> Clone for Keys<'a, K, V> {
393    #[inline]
394    fn clone(&self) -> Self {
395        match self {
396            Self::One(iter) => Self::One(iter.clone()),
397            Self::Multi(iter) => Self::Multi(iter.clone()),
398        }
399    }
400}
401
402impl<'a, K, V> Iterator for Keys<'a, K, V> {
403    type Item = &'a K;
404
405    #[inline]
406    fn next(&mut self) -> Option<Self::Item> {
407        match self {
408            Self::One(iter) => iter.next().map(|item| &item.0),
409            Self::Multi(iter) => iter.next(),
410        }
411    }
412
413    #[inline]
414    fn size_hint(&self) -> (usize, Option<usize>) {
415        match self {
416            Self::One(iter) => {
417                let l = iter.len();
418                (l, Some(l))
419            }
420            Self::Multi(iter) => {
421                let l = iter.len();
422                (l, Some(l))
423            }
424        }
425    }
426}
427
428impl<'a, K, V> ExactSizeIterator for Keys<'a, K, V> {
429    #[inline]
430    fn len(&self) -> usize {
431        match self {
432            Self::One(iter) => iter.len(),
433            Self::Multi(iter) => iter.len(),
434        }
435    }
436}
437
438impl<'a, K, V> DoubleEndedIterator for Keys<'a, K, V> {
439    #[inline]
440    fn next_back(&mut self) -> Option<Self::Item> {
441        match self {
442            Self::One(iter) => iter.next_back().map(|item| &item.0),
443            Self::Multi(iter) => iter.next_back(),
444        }
445    }
446}
447
448pub enum Values<'a, K, V> {
449    One(option::Iter<'a, (K, V)>),
450    Multi(btree_map::Values<'a, K, V>),
451}
452
453impl<'a, K, V> Clone for Values<'a, K, V> {
454    #[inline]
455    fn clone(&self) -> Self {
456        match self {
457            Self::One(iter) => Self::One(iter.clone()),
458            Self::Multi(iter) => Self::Multi(iter.clone()),
459        }
460    }
461}
462
463impl<'a, K, V> Iterator for Values<'a, K, V> {
464    type Item = &'a V;
465
466    #[inline]
467    fn next(&mut self) -> Option<Self::Item> {
468        match self {
469            Self::One(iter) => iter.next().map(|item| &item.1),
470            Self::Multi(iter) => iter.next(),
471        }
472    }
473
474    #[inline]
475    fn size_hint(&self) -> (usize, Option<usize>) {
476        match self {
477            Self::One(iter) => {
478                let l = iter.len();
479                (l, Some(l))
480            }
481            Self::Multi(iter) => {
482                let l = iter.len();
483                (l, Some(l))
484            }
485        }
486    }
487}
488
489impl<'a, K, V> ExactSizeIterator for Values<'a, K, V> {
490    #[inline]
491    fn len(&self) -> usize {
492        match self {
493            Self::One(iter) => iter.len(),
494            Self::Multi(iter) => iter.len(),
495        }
496    }
497}
498
499impl<'a, K, V> DoubleEndedIterator for Values<'a, K, V> {
500    #[inline]
501    fn next_back(&mut self) -> Option<Self::Item> {
502        match self {
503            Self::One(iter) => iter.next_back().map(|item| &item.1),
504            Self::Multi(iter) => iter.next_back(),
505        }
506    }
507}
508
509pub enum ValuesMut<'a, K, V> {
510    One(option::IterMut<'a, (K, V)>),
511    Multi(btree_map::ValuesMut<'a, K, V>),
512}
513
514impl<'a, K, V> Iterator for ValuesMut<'a, K, V> {
515    type Item = &'a mut V;
516
517    #[inline]
518    fn next(&mut self) -> Option<Self::Item> {
519        match self {
520            Self::One(iter) => iter.next().map(|item| &mut item.1),
521            Self::Multi(iter) => iter.next(),
522        }
523    }
524
525    #[inline]
526    fn size_hint(&self) -> (usize, Option<usize>) {
527        match self {
528            Self::One(iter) => {
529                let l = iter.len();
530                (l, Some(l))
531            }
532            Self::Multi(iter) => {
533                let l = iter.len();
534                (l, Some(l))
535            }
536        }
537    }
538}
539
540impl<'a, K, V> ExactSizeIterator for ValuesMut<'a, K, V> {
541    #[inline]
542    fn len(&self) -> usize {
543        match self {
544            Self::One(iter) => iter.len(),
545            Self::Multi(iter) => iter.len(),
546        }
547    }
548}
549
550impl<'a, K, V> DoubleEndedIterator for ValuesMut<'a, K, V> {
551    #[inline]
552    fn next_back(&mut self) -> Option<Self::Item> {
553        match self {
554            Self::One(iter) => iter.next_back().map(|item| &mut item.1),
555            Self::Multi(iter) => iter.next_back(),
556        }
557    }
558}
559
560pub enum Entry<'a, K: 'a, V: 'a> {
561    Occupied(OccupiedEntry<'a, K, V>),
562    Vacant(VacantEntry<'a, K, V>),
563}
564
565pub enum OccupiedEntry<'a, K: 'a, V: 'a> {
566    One(&'a mut Option<(K, V)>),
567    Multi(btree_map::OccupiedEntry<'a, K, V>),
568}
569
570impl<'a, K, V> OccupiedEntry<'a, K, V>
571where
572    K: Ord,
573{
574    #[inline]
575    pub fn get(&self) -> &V {
576        match self {
577            Self::One(o) => &o.as_ref().unwrap().1,
578            Self::Multi(ent) => ent.get(),
579        }
580    }
581
582    #[inline]
583    pub fn get_mut(&mut self) -> &mut V {
584        match self {
585            Self::One(o) => &mut o.as_mut().unwrap().1,
586            Self::Multi(ent) => ent.get_mut(),
587        }
588    }
589
590    #[inline]
591    pub fn key(&self) -> &K {
592        match self {
593            Self::One(o) => &o.as_ref().unwrap().0,
594            Self::Multi(ent) => ent.key(),
595        }
596    }
597
598    #[inline]
599    pub fn remove(self) -> V {
600        match self {
601            Self::One(o) => {
602                let (_k, v) = o.take().unwrap();
603                v
604            }
605            Self::Multi(ent) => ent.remove(),
606        }
607    }
608
609    #[inline]
610    pub fn remove_entry(self) -> (K, V) {
611        match self {
612            Self::One(o) => o.take().unwrap(),
613            Self::Multi(ent) => ent.remove_entry(),
614        }
615    }
616
617    #[inline]
618    pub fn insert(self, value: V) -> V {
619        match self {
620            Self::One(o) => {
621                let (k, old_v) = o.take().unwrap();
622                o.replace((k, value));
623                old_v
624            }
625            Self::Multi(mut ent) => ent.insert(value),
626        }
627    }
628
629    #[inline]
630    pub fn into_mut(self) -> &'a mut V {
631        match self {
632            Self::One(o) => &mut o.as_mut().unwrap().1,
633            Self::Multi(ent) => ent.into_mut(),
634        }
635    }
636}
637
638pub enum VacantEntry<'a, K, V> {
639    One(VacantEntryOne<'a, K, V>),
640    Multi(btree_map::VacantEntry<'a, K, V>),
641}
642
643struct VacantEntryOne<'a, K: 'a, V: 'a> {
644    pub(crate) key: K,                        // Owned key to insert
645    pub(crate) map: &'a mut VariousMap<K, V>, // Reference to the VariousMap
646}
647
648impl<'a, K, V> VacantEntry<'a, K, V>
649where
650    K: Ord,
651{
652    #[inline]
653    pub fn key(&self) -> &K {
654        match self {
655            Self::One(ent) => &ent.key,
656            Self::Multi(ent) => ent.key(),
657        }
658    }
659
660    #[inline]
661    pub fn into_key(self) -> K {
662        match self {
663            Self::One(ent) => ent.key,
664            Self::Multi(ent) => ent.into_key(),
665        }
666    }
667
668    #[inline]
669    pub fn insert(self, value: V) -> &'a mut V {
670        match self {
671            Self::One(ent) => {
672                let mut _value = MaybeUninit::new(value);
673                let mut _key = MaybeUninit::new(ent.key);
674                let map = ent.map;
675                if let VariousMap::One(item) = map {
676                    if item.is_none() {
677                        unsafe {
678                            // we should have return here, but don't because of the borrow checker
679                            item.replace((_key.assume_init_read(), _value.assume_init_read()));
680                        }
681                    } else {
682                        let (old_k, old_v) = item.take().unwrap();
683                        unsafe {
684                            if &old_k == _key.assume_init_ref() {
685                                // we should have return here, but don't because of the borrow checker
686                                item.replace((_key.assume_init_read(), _value.assume_init_read()));
687                            } else {
688                                let _ = item;
689                                let mut _map = BTreeMap::new();
690                                _map.insert(old_k, old_v);
691                                *map = VariousMap::Multi(_map);
692                            }
693                        }
694                    }
695                }
696                match map {
697                    VariousMap::One(Some(item)) => &mut item.1,
698                    VariousMap::Multi(map) => unsafe {
699                        map.entry(_key.assume_init_read()).or_insert(_value.assume_init_read())
700                    },
701                    _ => unreachable!(),
702                }
703            }
704            Self::Multi(ent) => ent.insert(value),
705        }
706    }
707}
708
709impl<'a, K, V> Entry<'a, K, V>
710where
711    K: Ord,
712{
713    #[inline]
714    pub fn or_insert(self, default: V) -> &'a mut V {
715        match self {
716            Entry::Occupied(entry) => entry.into_mut(),
717            Entry::Vacant(entry) => entry.insert(default),
718        }
719    }
720
721    #[inline]
722    pub fn or_insert_with<F: FnOnce() -> V>(self, default_fn: F) -> &'a mut V {
723        match self {
724            Entry::Occupied(entry) => entry.into_mut(),
725            Entry::Vacant(entry) => entry.insert(default_fn()),
726        }
727    }
728
729    #[inline]
730    pub fn and_modify<F>(mut self, f: F) -> Self
731    where
732        F: FnOnce(&mut V),
733    {
734        if let Entry::Occupied(ref mut entry) = self {
735            f(entry.get_mut());
736        }
737        self
738    }
739
740    #[inline]
741    pub fn key(&self) -> &K {
742        match self {
743            Entry::Occupied(entry) => entry.key(),
744            Entry::Vacant(entry) => entry.key(),
745        }
746    }
747
748    #[inline]
749    pub fn insert_entry(self, value: V) -> OccupiedEntry<'a, K, V> {
750        match self {
751            Entry::Occupied(mut ent) => {
752                match ent {
753                    OccupiedEntry::One(ref mut o) => {
754                        let (k, _old_v) = o.take().unwrap();
755                        o.replace((k, value));
756                    }
757                    OccupiedEntry::Multi(ref mut ent) => {
758                        ent.insert(value);
759                    }
760                }
761                ent
762            }
763            Entry::Vacant(VacantEntry::One(entry)) => {
764                let mut _value = MaybeUninit::new(value);
765                let mut _key = MaybeUninit::new(entry.key);
766                let map = entry.map;
767                if let VariousMap::One(item) = map {
768                    if item.is_none() {
769                        unsafe {
770                            item.replace((_key.assume_init_read(), _value.assume_init_read()));
771                        }
772                    } else {
773                        let (old_k, old_v) = item.take().unwrap();
774                        unsafe {
775                            if &old_k == _key.assume_init_ref() {
776                                item.replace((_key.assume_init_read(), _value.assume_init_read()));
777                            } else {
778                                let _ = item;
779                                let mut _map = BTreeMap::new();
780                                _map.insert(old_k, old_v);
781                                *map = VariousMap::Multi(_map);
782                            }
783                        }
784                    }
785                }
786                match map {
787                    VariousMap::One(o) => OccupiedEntry::One(o),
788                    VariousMap::Multi(map) => {
789                        let ent = unsafe {
790                            map.entry(_key.assume_init_read())
791                                .insert_entry(_value.assume_init_read())
792                        };
793                        OccupiedEntry::Multi(ent)
794                    }
795                }
796            }
797            Entry::Vacant(VacantEntry::Multi(ent)) => OccupiedEntry::Multi(ent.insert_entry(value)),
798        }
799    }
800
801    #[inline]
802    pub fn or_default(self) -> &'a mut V
803    where
804        V: Default,
805    {
806        self.or_insert_with(Default::default)
807    }
808}
809
810#[cfg(test)]
811mod tests {
812    use super::*;
813
814    #[test]
815    fn test_new() {
816        let map: VariousMap<i32, String> = VariousMap::new();
817        assert!(matches!(map, VariousMap::One(None)));
818    }
819
820    #[test]
821    fn test_insert_and_get_one() {
822        let mut map = VariousMap::<usize, usize>::new();
823        assert_eq!(map.insert(1, 1), None);
824        assert!(matches!(map, VariousMap::One(Some((ref k, ref v))) if *k == 1 && *v == 1));
825        assert_eq!(map.get(&1), Some(&1));
826        assert_eq!(map.get(&2), None);
827        **(map.get_mut(&1).as_mut().unwrap()) = 2;
828        assert_eq!(map.get(&1), Some(&2));
829        assert_eq!(map.insert(1, 3), Some(2));
830        assert_eq!(map.get(&1), Some(&3));
831    }
832
833    #[test]
834    fn test_insert_transition_to_multi() {
835        let mut map = VariousMap::new();
836        map.insert(1, "one".to_string());
837        assert_eq!(map.insert(2, "two".to_string()), None);
838
839        // After inserting a second element, it should transition to BTreeMap
840        match map {
841            VariousMap::Multi(ref btree_map) => {
842                assert_eq!(btree_map.len(), 2);
843                assert_eq!(btree_map.get(&1), Some(&"one".to_string()));
844                assert_eq!(btree_map.get(&2), Some(&"two".to_string()));
845            }
846            _ => panic!("Expected Multi variant after inserting two elements"),
847        }
848
849        assert_eq!(map.get(&1), Some(&"one".to_string()));
850        assert_eq!(map.get(&2), Some(&"two".to_string()));
851        assert_eq!(map.get(&3), None);
852
853        **map.get_mut(&2).as_mut().unwrap() = "two_update".to_string();
854        assert_eq!(map.get(&2), Some(&"two_update".to_string()));
855    }
856
857    #[test]
858    fn test_iter_one() {
859        let mut map = VariousMap::new();
860        map.insert(1, "one".to_string());
861        let mut collected: Vec<_> = map.iter().collect();
862        collected.sort_by_key(|(k, _)| *k);
863        assert_eq!(collected, vec![(&1, &"one".to_string())]);
864    }
865
866    #[test]
867    fn test_iter_multi() {
868        let mut map = VariousMap::new();
869        map.insert(1, "one".to_string());
870        map.insert(2, "two".to_string()); // Transition to Multi
871        map.insert(3, "three".to_string());
872
873        let mut collected: Vec<_> = map.iter().collect();
874        collected.sort_by_key(|(k, _)| *k);
875        assert_eq!(
876            collected,
877            vec![(&1, &"one".to_string()), (&2, &"two".to_string()), (&3, &"three".to_string())]
878        );
879    }
880
881    #[test]
882    fn test_into_iter_one() {
883        let mut map = VariousMap::new();
884        map.insert(1, "one".to_string());
885        let mut collected: Vec<_> = map.into_iter().collect();
886        collected.sort_by_key(|(k, _)| *k);
887        assert_eq!(collected, vec![(1, "one".to_string())]);
888    }
889
890    #[test]
891    fn test_contains_key() {
892        let mut map = VariousMap::new();
893        map.insert(1, "one".to_string());
894        assert!(map.contains_key(&1));
895        assert!(!map.contains_key(&2));
896
897        map.insert(2, "two".to_string()); // Transition to Multi
898        assert!(map.contains_key(&1));
899        assert!(map.contains_key(&2));
900        assert!(!map.contains_key(&3));
901    }
902
903    #[test]
904    fn test_entry_api_one() {
905        let mut map = VariousMap::new();
906        map.insert(1, "one".to_string());
907
908        // Occupied
909        match map.entry(1) {
910            Entry::Occupied(ent) => {
911                assert_eq!(ent.get(), &"one".to_string());
912                ent.insert("one_updated".to_string());
913            }
914            Entry::Vacant(_) => panic!("Should be occupied"),
915        }
916        assert_eq!(map.get(&1), Some(&"one_updated".to_string()));
917
918        // Vacant
919        match map.entry(2) {
920            Entry::Occupied(_) => panic!("Should be vacant"),
921            Entry::Vacant(ent) => {
922                ent.insert("two".to_string());
923            }
924        }
925        assert!(map.contains_key(&2));
926    }
927
928    #[test]
929    fn test_entry_api_or_insert() {
930        let mut map = VariousMap::new();
931        map.entry(1).or_insert("one".to_string());
932
933        // Occupied
934        let v = map
935            .entry(1)
936            .and_modify(|v| {
937                *v = "one_3".to_string();
938            })
939            .or_insert("one_2".to_string());
940        assert_eq!(v, &"one_3".to_string());
941
942        // Vacant
943        map.entry(2).or_insert("two".to_string());
944        assert_eq!(map.get(&2), Some(&"two".to_string()));
945    }
946
947    #[test]
948    fn test_entry_api_insert_entry() {
949        let mut map = VariousMap::new();
950
951        // Vacant -> insert_entry
952        let occupied = map.entry(1).insert_entry("one".to_string());
953        assert_eq!(occupied.get(), &"one".to_string());
954
955        // Occupied -> insert_entry
956        let occupied = map.entry(1).insert_entry("one_updated".to_string());
957        assert_eq!(occupied.get(), &"one_updated".to_string());
958        assert_eq!(map.get(&1), Some(&"one_updated".to_string()));
959    }
960
961    #[test]
962    fn test_keys() {
963        let mut map = VariousMap::new();
964        map.insert(1, "one".to_string());
965        assert_eq!(map.keys().collect::<Vec<_>>(), vec![&1]);
966
967        map.insert(2, "two".to_string());
968        let mut keys = map.keys().collect::<Vec<_>>();
969        keys.sort();
970        assert_eq!(keys, vec![&1, &2]);
971    }
972
973    #[test]
974    fn test_values() {
975        let mut map = VariousMap::new();
976        map.insert(1, "one".to_string());
977        assert_eq!(map.values().collect::<Vec<_>>(), vec![&"one".to_string()]);
978
979        map.insert(2, "two".to_string());
980        let mut values = map.values().collect::<Vec<_>>();
981        values.sort();
982        assert_eq!(values, vec![&"one".to_string(), &"two".to_string()]);
983    }
984
985    #[test]
986    fn test_values_mut() {
987        let mut map = VariousMap::new();
988        map.insert(1, "one".to_string());
989        for v in map.values_mut() {
990            *v = "one_updated".to_string();
991        }
992        assert_eq!(map.get(&1), Some(&"one_updated".to_string()));
993
994        map.insert(2, "two".to_string());
995        for v in map.values_mut() {
996            v.push_str("_mut");
997        }
998        assert_eq!(map.get(&1), Some(&"one_updated_mut".to_string()));
999        assert_eq!(map.get(&2), Some(&"two_mut".to_string()));
1000    }
1001
1002    #[test]
1003    fn test_iter_mut() {
1004        let mut map = VariousMap::new();
1005        map.insert(1, "one".to_string());
1006        for (k, v) in map.iter_mut() {
1007            assert_eq!(k, &1);
1008            *v = "one_updated".to_string();
1009        }
1010        assert_eq!(map.get(&1), Some(&"one_updated".to_string()));
1011
1012        map.insert(2, "two".to_string());
1013        for (k, v) in map.iter_mut() {
1014            if *k == 1 {
1015                *v = "one_final".to_string();
1016            } else if *k == 2 {
1017                *v = "two_final".to_string();
1018            }
1019        }
1020        assert_eq!(map.get(&1), Some(&"one_final".to_string()));
1021        assert_eq!(map.get(&2), Some(&"two_final".to_string()));
1022    }
1023}
embed_collections/various_map.rs

embed_collections/
various_map.rs
