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

embed_collections/
various_map.rs
