aabel_hashmap_rs/
lib.rs

1use std::{borrow::Borrow, collections::HashMap, hash::Hash};
2
3/// Extends the [`HashMap`] functionality.
4pub trait HashmapExt<K, V> {
5    /// Updates the value of a given key in the map.
6    ///
7    /// If the map does not have this key present, [`None`] is returned.
8    ///
9    /// If the map did have this key present, the updated value is returned.
10    fn update<Q, F>(&mut self, k: &Q, f: F) -> Option<V>
11    where
12        Q: ?Sized,
13        K: Borrow<Q> + PartialEq + Eq + Hash,
14        Q: Hash + Eq,
15        V: Clone,
16        F: FnMut(&mut V);
17
18    /// Updates the value of a given key in the map.
19    ///
20    /// If the map does not have this key present, the provided value is inserted.
21    ///
22    /// If the map did have this key present, the updated value is returned.
23    fn update_or_insert<F>(&mut self, k: K, f: F, def: V) -> Option<V>
24    where
25        K: PartialEq + Eq + Hash,
26        V: Clone,
27        F: FnMut(&mut V);
28}
29
30impl<K, V> HashmapExt<K, V> for HashMap<K, V> {
31    fn update<Q, F>(&mut self, k: &Q, mut f: F) -> Option<V>
32    where
33        Q: ?Sized,
34        K: Borrow<Q> + PartialEq + Eq + Hash,
35        Q: Hash + Eq,
36        V: Clone,
37        F: FnMut(&mut V),
38    {
39        if let Some(totals) = self.get_mut(k) {
40            f(totals);
41            Some(totals.clone())
42        } else {
43            None
44        }
45    }
46
47    fn update_or_insert<F>(&mut self, k: K, f: F, def: V) -> Option<V>
48    where
49        K: PartialEq + Eq + Hash,
50        V: Clone,
51        F: FnMut(&mut V),
52    {
53        self.update(&k, f).or_else(|| {
54            self.insert(k, def.clone());
55            Some(def)
56        })
57    }
58}
59
60#[cfg(test)]
61mod tests {
62    use super::*;
63
64    #[test]
65    fn not_exist_update() {
66        let mut map: HashMap<usize, u8> = HashMap::new();
67        let res = map.update(&1usize, |u| *u += 1);
68        assert!(res.is_none())
69    }
70
71    #[test]
72    fn exist_update() {
73        let mut map: HashMap<usize, u8> = HashMap::new();
74        let res = map.insert(1usize, 10);
75        assert!(res.is_none());
76
77        let res = map.update(&1usize, |u| *u += 1);
78        assert!(res.is_some());
79        assert_eq!(Some(11), res);
80    }
81
82    #[test]
83    fn not_exist_continuation() {
84        let mut map: HashMap<usize, u8> = HashMap::new();
85
86        let res = map
87            .update(&1usize, |u| *u += 1)
88            .or_else(|| map.insert(1usize, 1).or(Some(1)));
89        assert_eq!(res, Some(1));
90    }
91
92    #[test]
93    fn exist_continuation() {
94        let mut map: HashMap<usize, u8> = HashMap::new();
95        let res = map.insert(1usize, 10);
96        assert!(res.is_none());
97
98        let res = map
99            .update(&1usize, |u| *u += 1)
100            .or_else(|| map.insert(1usize, 1).or(Some(1)));
101        assert_eq!(res, Some(11));
102    }
103
104    #[test]
105    fn not_exist_update_or_insert() {
106        let mut map: HashMap<usize, u8> = HashMap::new();
107        let res = map.update_or_insert(1usize, |u| *u += 1, 1);
108        assert_eq!(res, Some(1));
109    }
110
111    #[test]
112    fn exist_update_or_insert() {
113        let mut map: HashMap<usize, u8> = HashMap::new();
114        let res = map.insert(1usize, 10);
115        assert!(res.is_none());
116
117        let res = map.update_or_insert(1usize, |u| *u += 1, 1);
118        assert_eq!(res, Some(11));
119    }
120}
aabel_hashmap_rs/lib.rs

aabel_hashmap_rs/
lib.rs
