defaultmap/
btreemap.rs

1use derive_more::Debug;
2use std::borrow::Borrow;
3use std::collections::btree_map::*;
4use std::collections::BTreeMap;
5use std::iter::{FromIterator, IntoIterator};
6use std::ops::RangeBounds;
7use std::ops::{Index, IndexMut};
8
9use crate::DefaultFn;
10
11/// A `BTreeMap` that returns a default when keys are accessed that are not present.
12#[derive(Clone, Debug)]
13#[cfg_attr(feature = "with-serde", derive(serde::Serialize, serde::Deserialize))]
14pub struct DefaultBTreeMap<K: Eq + Ord, V> {
15    map: BTreeMap<K, V>,
16    default: V,
17    #[debug(skip)]
18    #[cfg_attr(feature = "with-serde", serde(skip))]
19    default_fn: Box<dyn DefaultFn<V>>,
20}
21
22impl<K: Eq + Ord, V: PartialEq> PartialEq for DefaultBTreeMap<K, V> {
23    fn eq(&self, other: &Self) -> bool {
24        self.map == other.map && self.default == other.default
25    }
26}
27
28impl<K: Eq + Ord, V: Eq> Eq for DefaultBTreeMap<K, V> {}
29
30impl<K: Eq + Ord, V: Default> DefaultBTreeMap<K, V> {
31    /// The `new()` constructor creates an empty DefaultBTreeMap with the default of `V`
32    /// as the default for missing keys.
33    /// This is desired default for most use cases, if your case requires a
34    /// different default you should use the `with_default()` constructor.
35    pub fn new() -> DefaultBTreeMap<K, V> {
36        DefaultBTreeMap {
37            map: BTreeMap::default(),
38            default_fn: Box::new(|| V::default()),
39            default: V::default(),
40        }
41    }
42}
43
44impl<K: Eq + Ord, V: Default> Default for DefaultBTreeMap<K, V> {
45    /// The `default()` method is equivalent to `DefaultBTreeMap::new()`.
46    fn default() -> DefaultBTreeMap<K, V> {
47        DefaultBTreeMap::new()
48    }
49}
50
51impl<K: Eq + Ord, V: Default> From<BTreeMap<K, V>> for DefaultBTreeMap<K, V> {
52    /// If you already have a `BTreeMap` that you would like to convert to a
53    /// `DefaultBTreeMap` you can use the `into()` method on the `BTreeMap` or the
54    /// `from()` constructor of `DefaultBTreeMap`.
55    /// The default value for missing keys will be `V::default()`,
56    /// if this is not desired `DefaultBTreeMap::from_map_with_default()` should be used.
57    fn from(map: BTreeMap<K, V>) -> DefaultBTreeMap<K, V> {
58        DefaultBTreeMap {
59            map,
60            default_fn: Box::new(|| V::default()),
61            default: V::default(),
62        }
63    }
64}
65
66impl<K: Eq + Ord, V> From<DefaultBTreeMap<K, V>> for BTreeMap<K, V> {
67    /// The into method can be used to convert a `DefaultBTreeMap` back into a
68    /// `BTreeMap`.
69    fn from(default_map: DefaultBTreeMap<K, V>) -> BTreeMap<K, V> {
70        default_map.map
71    }
72}
73
74impl<K: Eq + Ord, V: Clone + 'static> DefaultBTreeMap<K, V> {
75    /// Creates an empty `DefaultBTreeMap` with `default` as the default for missing keys.
76    /// When the provided `default` is equivalent to `V::default()` it is preferred to use
77    /// `DefaultBTreeMap::default()` instead.
78    pub fn with_default(default: V) -> DefaultBTreeMap<K, V> {
79        DefaultBTreeMap {
80            map: BTreeMap::new(),
81            default: default.clone(),
82            default_fn: Box::new(move || default.clone()),
83        }
84    }
85
86    /// Creates a `DefaultBTreeMap` based on a default and an already existing `BTreeMap`.
87    /// If `V::default()` is the supplied default, usage of the `from()` constructor or the
88    /// `into()` method on the original `BTreeMap` is preferred.
89    pub fn from_map_with_default(map: BTreeMap<K, V>, default: V) -> DefaultBTreeMap<K, V> {
90        DefaultBTreeMap {
91            map,
92            default: default.clone(),
93            default_fn: Box::new(move || default.clone()),
94        }
95    }
96
97    /// Changes the default value permanently or until `set_default()` is called again.
98    pub fn set_default(&mut self, new_default: V) {
99        self.default = new_default.clone();
100        self.default_fn = Box::new(move || new_default.clone());
101    }
102}
103
104impl<K: Eq + Ord, V> DefaultBTreeMap<K, V> {
105    /// Returns a reference to the value stored for the provided key.
106    /// If the key is not in the `DefaultBTreeMap` a reference to the default value is returned.
107    /// Usually the `map[key]` method of retrieving keys is preferred over using `get` directly.
108    /// This method accepts both references and owned values as the key.
109    pub fn get<Q, QB: Borrow<Q>>(&self, key: QB) -> &V
110    where
111        K: Borrow<Q>,
112        Q: ?Sized + Ord + Eq,
113    {
114        self.map.get(key.borrow()).unwrap_or(&self.default)
115    }
116
117    /// Returns the an owned version of the default value
118    /// ```
119    /// use defaultmap::DefaultBTreeMap;
120    /// assert_eq!(DefaultBTreeMap::<String, i32>::new().get_default(), 0);
121    /// ```
122    pub fn get_default(&self) -> V {
123        self.default_fn.call()
124    }
125
126    /// Creates an empty `DefaultBTreeMap` with `default_fn` as the default value generation
127    /// function for missing keys. When the provided `default_fn` only calls clone on a value,
128    /// using `DefaultBTreeMap::new` is preferred.
129    pub fn with_fn(default_fn: impl DefaultFn<V> + 'static) -> DefaultBTreeMap<K, V> {
130        DefaultBTreeMap {
131            map: BTreeMap::new(),
132            default: default_fn.call(),
133            default_fn: Box::new(default_fn),
134        }
135    }
136
137    /// Creates a `DefaultBTreeMap` based on an existing map and using `default_fn` as the default
138    /// value generation function for missing keys. When the provided `default_fn` is equivalent to
139    /// V::default(), then using `DefaultBTreeMap::from(map)` is preferred.
140    pub fn from_map_with_fn(
141        map: BTreeMap<K, V>,
142        default_fn: impl DefaultFn<V> + 'static,
143    ) -> DefaultBTreeMap<K, V> {
144        DefaultBTreeMap {
145            map,
146            default: default_fn.call(),
147            default_fn: Box::new(default_fn),
148        }
149    }
150}
151
152impl<K: Eq + Ord, V> DefaultBTreeMap<K, V> {
153    /// Returns a mutable reference to the value stored for the provided key.
154    /// If there is no value stored for the key the default value is first inserted for this
155    /// key before returning the reference.
156    /// Usually the `map[key] = new_val` is prefered over using `get_mut` directly.
157    /// This method only accepts owned values as the key.
158    pub fn get_mut(&mut self, key: K) -> &mut V {
159        let entry = self.map.entry(key);
160        match entry {
161            Entry::Occupied(occupied) => occupied.into_mut(),
162            Entry::Vacant(vacant) => vacant.insert(self.default_fn.call()),
163        }
164    }
165}
166
167/// Implements the `Index` trait so you can do `map[key]`.
168/// Nonmutable indexing can be done both by passing a reference or an owned value as the key.
169impl<K: Eq + Ord, KB: Borrow<K>, V> Index<KB> for DefaultBTreeMap<K, V> {
170    type Output = V;
171
172    fn index(&self, index: KB) -> &V {
173        self.get(index)
174    }
175}
176
177/// Implements the `IndexMut` trait so you can do `map[key] = val`.
178/// Mutably indexing can only be done when passing an owned value as the key.
179impl<K: Eq + Ord, V> IndexMut<K> for DefaultBTreeMap<K, V> {
180    #[inline]
181    fn index_mut(&mut self, index: K) -> &mut V {
182        self.get_mut(index)
183    }
184}
185
186// grcov-excl-start
187/// These methods simply forward to the underlying `BTreeMap`, see that
188/// [documentation](https://doc.rust-lang.org/std/collections/struct.BTreeMap.html) for
189/// the usage of these methods.
190impl<K: Eq + Ord, V> DefaultBTreeMap<K, V> {
191    #[inline]
192    pub fn clear(&mut self) {
193        self.map.clear()
194    }
195    #[inline]
196    pub fn first_key_value(&self) -> Option<(&K, &V)>
197    where
198        K: Ord,
199    {
200        self.map.first_key_value()
201    }
202    #[inline]
203    pub fn first_entry(&mut self) -> Option<OccupiedEntry<'_, K, V>>
204    where
205        K: Ord,
206    {
207        self.map.first_entry()
208    }
209    #[inline]
210    pub fn pop_first(&mut self) -> Option<(K, V)>
211    where
212        K: Ord,
213    {
214        self.map.pop_first()
215    }
216    #[inline]
217    pub fn last_key_value(&self) -> Option<(&K, &V)>
218    where
219        K: Ord,
220    {
221        self.map.last_key_value()
222    }
223    #[inline]
224    pub fn last_entry(&mut self) -> Option<OccupiedEntry<'_, K, V>>
225    where
226        K: Ord,
227    {
228        self.map.last_entry()
229    }
230    #[inline]
231    pub fn pop_last(&mut self) -> Option<(K, V)>
232    where
233        K: Ord,
234    {
235        self.map.pop_last()
236    }
237    #[inline]
238    pub fn contains_key<Q>(&self, k: &Q) -> bool
239    where
240        K: Borrow<Q>,
241        Q: ?Sized + Ord,
242    {
243        self.map.contains_key(k)
244    }
245    #[inline]
246    pub fn insert(&mut self, k: K, v: V) -> Option<V> {
247        self.map.insert(k, v)
248    }
249    #[inline]
250    pub fn remove<Q>(&mut self, k: &Q) -> Option<V>
251    where
252        K: Borrow<Q>,
253        Q: ?Sized + Ord,
254    {
255        self.map.remove(k)
256    }
257    #[inline]
258    pub fn retain<RF>(&mut self, f: RF)
259    where
260        RF: FnMut(&K, &mut V) -> bool,
261    {
262        self.map.retain(f)
263    }
264    #[inline]
265    pub fn append(&mut self, other: &mut Self) {
266        self.map.append(&mut other.map)
267    }
268    #[inline]
269    pub fn range<T: ?Sized, R>(&self, range: R) -> Range<'_, K, V>
270    where
271        T: Ord,
272        K: Borrow<T> + Ord,
273        R: RangeBounds<T>,
274    {
275        self.map.range(range)
276    }
277    #[inline]
278    pub fn range_mut<T: ?Sized, R>(&mut self, range: R) -> RangeMut<'_, K, V>
279    where
280        T: Ord,
281        K: Borrow<T> + Ord,
282        R: RangeBounds<T>,
283    {
284        self.map.range_mut(range)
285    }
286    #[inline]
287    pub fn entry(&mut self, key: K) -> Entry<K, V> {
288        self.map.entry(key)
289    }
290    #[inline]
291    pub fn into_keys(self) -> IntoKeys<K, V> {
292        self.map.into_keys()
293    }
294    #[inline]
295    pub fn into_values(self) -> IntoValues<K, V> {
296        self.map.into_values()
297    }
298
299    #[inline]
300    pub fn iter(&self) -> Iter<K, V> {
301        self.map.iter()
302    }
303    #[inline]
304    pub fn iter_mut(&mut self) -> IterMut<K, V> {
305        self.map.iter_mut()
306    }
307    #[inline]
308    pub fn keys(&self) -> Keys<K, V> {
309        self.map.keys()
310    }
311    #[inline]
312    pub fn values(&self) -> Values<K, V> {
313        self.map.values()
314    }
315    #[inline]
316    pub fn values_mut(&mut self) -> ValuesMut<K, V> {
317        self.map.values_mut()
318    }
319    #[inline]
320    pub fn len(&self) -> usize {
321        self.map.len()
322    }
323    #[inline]
324    pub fn is_empty(&self) -> bool {
325        self.map.is_empty()
326    }
327}
328// grcov-excl-stop
329
330impl<K: Eq + Ord, V: Default> FromIterator<(K, V)> for DefaultBTreeMap<K, V> {
331    fn from_iter<I>(iter: I) -> Self
332    where
333        I: IntoIterator<Item = (K, V)>,
334    {
335        Self {
336            map: BTreeMap::from_iter(iter),
337            default: V::default(),
338            default_fn: Box::new(|| V::default()),
339        }
340    }
341}
342
343/// The `defaultbtreemap!` macro can be used to easily initialize a `DefaultBTreeMap` in the
344/// following ways:
345///
346/// ```
347/// # #[macro_use] extern crate defaultmap;
348/// # use defaultmap::*;
349/// // An empty map with the default as default
350/// let _: DefaultBTreeMap<i32, i32> = defaultbtreemap!{};
351///
352/// // An empty map with a specified default
353/// let _: DefaultBTreeMap<i32, i32> = defaultbtreemap!{5};
354///
355/// // A prefilled map with the default as the default
356/// let _: DefaultBTreeMap<i32, i32> = defaultbtreemap!{
357///     1 => 10,
358///     5 => 20,
359///     6 => 30,
360/// };
361///
362/// // A prefilled map with a custom default
363/// let _: DefaultBTreeMap<i32, i32> = defaultbtreemap!{
364///     5,
365///     1 => 10,
366///     5 => 20,
367///     6 => 30,
368/// };
369/// ```
370#[macro_export]
371macro_rules! defaultbtreemap {
372    (@single $($x:tt)*) => (());
373    (@count $($rest:expr),*) => (<[()]>::len(&[$(defaultbtreemap!(@single $rest)),*]));
374    // Copied almost verbatim from maplit
375    (@btreemap $($key:expr => $value:expr),*) => {
376        {
377            let mut _map = ::std::collections::BTreeMap::new();
378            $(
379                _map.insert($key, $value);
380            )*
381            _map
382        }
383    };
384
385    ($($key:expr => $value:expr,)+) => { defaultbtreemap!($($key => $value),+) };
386    ($($key:expr => $value:expr),*) => {
387        {
388            let _map = defaultbtreemap!(@btreemap $($key => $value),*);
389            $crate::DefaultBTreeMap::from(_map)
390        }
391    };
392
393    ($default:expr$(, $key:expr => $value:expr)+ ,) => { defaultbtreemap!($default, $($key => $value),+) };
394    ($default:expr$(, $key:expr => $value:expr)*) => {
395        {
396            let _map = defaultbtreemap!(@btreemap $($key => $value),*);
397            $crate::DefaultBTreeMap::from_map_with_default(_map, $default)
398        }
399    };
400}
401
402#[cfg(test)]
403mod tests {
404    use super::DefaultBTreeMap;
405    use std::collections::BTreeMap;
406
407    #[test]
408    fn macro_test() {
409        // empty default
410        let macro_map: DefaultBTreeMap<i32, i32> = defaultbtreemap! {};
411        let normal_map = DefaultBTreeMap::<i32, i32>::default();
412        assert_eq!(macro_map, normal_map);
413
414        // with content
415        let macro_map: DefaultBTreeMap<_, _> = defaultbtreemap! {
416            1 => 2,
417            2 => 3,
418        };
419        let mut normal_map = DefaultBTreeMap::<_, _>::default();
420        normal_map[1] = 2;
421        normal_map[2] = 3;
422        assert_eq!(macro_map, normal_map);
423
424        // empty with custom default
425        let macro_map: DefaultBTreeMap<i32, i32> = defaultbtreemap! {5};
426        let normal_map = DefaultBTreeMap::<i32, i32>::with_default(5);
427        assert_eq!(macro_map, normal_map);
428
429        // filled btreemap with custom default
430        let macro_map: DefaultBTreeMap<_, _> = defaultbtreemap! {
431            5,
432            1 => 2,
433            2 => 3,
434        };
435        let mut normal_map = DefaultBTreeMap::<_, _>::with_default(5);
436        normal_map[1] = 2;
437        normal_map[2] = 3;
438        assert_eq!(macro_map, normal_map);
439    }
440
441    #[test]
442    fn add() {
443        let mut map: DefaultBTreeMap<i32, i32> = DefaultBTreeMap::default();
444        *map.get_mut(0) += 1;
445        map[1] += 4;
446        map[2] = map[0] + map.get(&1);
447        assert_eq!(*map.get(0), 1);
448        assert_eq!(*map.get(&0), 1);
449        assert_eq!(map[0], 1);
450        assert_eq!(map[&0], 1);
451        assert_eq!(*map.get(&1), 4);
452        assert_eq!(*map.get(&2), 5);
453        assert_eq!(*map.get(999), 0);
454        assert_eq!(*map.get(&999), 0);
455        assert_eq!(map[999], 0);
456        assert_eq!(map[&999], 0);
457    }
458
459    #[test]
460    fn counter() {
461        let nums = [1, 4, 3, 3, 4, 2, 4];
462        let mut counts: DefaultBTreeMap<i32, i32> = DefaultBTreeMap::default();
463        for num in nums.iter() {
464            counts[*num] += 1;
465        }
466
467        assert_eq!(1, counts[1]);
468        assert_eq!(1, counts[2]);
469        assert_eq!(2, counts[3]);
470        assert_eq!(3, counts[4]);
471        assert_eq!(0, counts[5]);
472    }
473
474    #[test]
475    fn change_default() {
476        let mut numbers: DefaultBTreeMap<i32, String> =
477            DefaultBTreeMap::with_default("Mexico".to_string());
478
479        assert_eq!("Mexico", numbers.get_mut(1));
480        assert_eq!("Mexico", numbers.get_mut(2));
481        assert_eq!("Mexico", numbers[3]);
482
483        numbers.set_default("Cybernetics".to_string());
484        assert_eq!("Mexico", numbers[1]);
485        assert_eq!("Mexico", numbers[2]);
486        assert_eq!("Cybernetics", numbers[3]);
487        assert_eq!("Cybernetics", numbers[4]);
488        assert_eq!("Cybernetics", numbers[5]);
489    }
490
491    #[test]
492    fn synonyms() {
493        let synonym_tuples = [
494            ("nice", "sweet"),
495            ("sweet", "candy"),
496            ("nice", "entertaining"),
497            ("nice", "good"),
498            ("entertaining", "absorbing"),
499        ];
500
501        let mut synonym_map: DefaultBTreeMap<&str, Vec<&str>> = DefaultBTreeMap::default();
502
503        for &(l, r) in synonym_tuples.iter() {
504            synonym_map[l].push(r);
505            synonym_map[r].push(l);
506        }
507
508        println!("{:#?}", synonym_map);
509        assert_eq!(synonym_map["good"], vec!["nice"]);
510        assert_eq!(synonym_map["nice"], vec!["sweet", "entertaining", "good"]);
511        assert_eq!(synonym_map["evil"], Vec::<&str>::new());
512    }
513
514    #[derive(Clone)]
515    struct Clonable;
516
517    #[derive(Default, Clone)]
518    struct DefaultableValue;
519
520    #[derive(Ord, PartialOrd, Eq, PartialEq)]
521    struct Orderable(i32);
522
523    #[test]
524    fn minimal_derives() {
525        let _: DefaultBTreeMap<Orderable, Clonable> = DefaultBTreeMap::with_default(Clonable);
526        let _: DefaultBTreeMap<Orderable, DefaultableValue> = DefaultBTreeMap::default();
527    }
528
529    #[test]
530    fn from() {
531        let normal: BTreeMap<i32, i32> = vec![(0, 1), (2, 3)].into_iter().collect();
532        let mut default: DefaultBTreeMap<_, _> = normal.into();
533        default.get_mut(4);
534        assert_eq!(default[0], 1);
535        assert_eq!(default[2], 3);
536        assert_eq!(default[1], 0);
537        assert_eq!(default[4], 0);
538        let expected: BTreeMap<i32, i32> = vec![(0, 1), (2, 3), (4, 0)].into_iter().collect();
539        assert_eq!(expected, default.into());
540    }
541
542    #[test]
543    fn with_fn() {
544        let i: i32 = 20;
545        let mut map = DefaultBTreeMap::with_fn(move || i);
546        map[0] += 1;
547        assert_eq!(21, map[0]);
548        assert_eq!(20, map[1]);
549    }
550
551    #[test]
552    fn from_map_with_fn() {
553        let i: i32 = 20;
554        let normal: BTreeMap<i32, i32> = vec![(0, 1), (2, 3)].into_iter().collect();
555        let mut map = DefaultBTreeMap::from_map_with_fn(normal, move || i);
556        map[0] += 1;
557        assert_eq!(map[0], 2);
558        assert_eq!(map[1], 20);
559        assert_eq!(map[2], 3);
560    }
561
562    #[cfg(feature = "with-serde")]
563    mod serde_tests {
564        use super::*;
565
566        #[test]
567        fn deserialize_static() {
568            let s = "{ 
569                        \"map\" : 
570                            {   \"foo\": 3, 
571                                \"bar\": 5 
572                            }, 
573                        \"default\":15 
574                    }";
575            let h: Result<DefaultBTreeMap<&str, i32>, _> = serde_json::from_str(&s);
576            let h = h.unwrap();
577            assert_eq!(h["foo"] * h["bar"], h["foobar"])
578        }
579
580        #[test]
581        fn serialize_and_back() {
582            let h1: DefaultBTreeMap<i32, u64> = defaultbtreemap!(1 => 10, 2 => 20, 3 => 30);
583            let s = serde_json::to_string(&h1).unwrap();
584            let h2: DefaultBTreeMap<i32, u64> = serde_json::from_str(&s).unwrap();
585            assert_eq!(h2, h2);
586            assert_eq!(h2[3], 30);
587        }
588
589        #[test]
590        fn serialize_default() {
591            let h1: DefaultBTreeMap<&str, u64> = DefaultBTreeMap::with_default(42);
592            let s = serde_json::to_string(&h1).unwrap();
593            let h2: DefaultBTreeMap<&str, u64> = serde_json::from_str(&s).unwrap();
594            assert_eq!(h2["answer"], 42);
595        }
596
597        #[test]
598        fn std_btreemap() {
599            let h1: DefaultBTreeMap<i32, i32> = defaultbtreemap!(1=> 10, 2=> 20);
600            let stdhm: std::collections::BTreeMap<i32, i32> = h1.clone().into();
601            let s = serde_json::to_string(&stdhm).unwrap();
602            let h2: DefaultBTreeMap<i32, i32> = DefaultBTreeMap::from_map_with_default(
603                serde_json::from_str(&s).unwrap(),
604                i32::default(),
605            );
606            assert_eq!(h1, h2);
607        }
608    }
609}
defaultmap/btreemap.rs

defaultmap/
btreemap.rs
