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

embed_btree/
various_map.rs
