try_insert_ext/
entry.rs

1#[cfg(feature = "std")]
2use std::collections::{btree_map, hash_map};
3
4/// Extends map entries with `or_try_insert_with` and `or_try_insert_with_key`.
5pub trait EntryInsertExt<'a, K, V> {
6    /// If empty, computes the value from the default function. If the function
7    /// returns `Ok`, inserts the value. If `f` returns `Err`, returns the
8    /// error. If there is no error, returns a mutable reference to the value in
9    /// the entry.
10    ///
11    /// # Examples
12    ///
13    /// ```
14    /// use std::collections::HashMap;
15    ///
16    /// use try_insert_ext::EntryInsertExt;
17    ///
18    /// let mut map: HashMap<&str, String> = HashMap::new();
19    /// let s = "hoho".to_string();
20    ///
21    /// let e: Result<&mut String, ()> = map
22    ///     .entry("poneyland")
23    ///     .or_try_insert_with(|| Err(()));
24    /// assert!(e.is_err());
25    /// map.entry("poneyland").or_try_insert_with::<_, ()>(|| Ok(s));
26    ///
27    /// assert_eq!(map["poneyland"], "hoho".to_string());
28    /// ```
29    fn or_try_insert_with<F: FnOnce() -> Result<V, E>, E>(self, default: F)
30        -> Result<&'a mut V, E>;
31
32    /// If empty, computes the value from the default function. If the function
33    /// returns `Ok`, inserts the value. If `f` returns `Err`, returns the
34    /// error. If there is no error, returns a mutable reference to the value in
35    /// the entry. This method allows for generating key-derived values for
36    /// insertion by providing the default function a reference to the key
37    /// that was moved during the `.entry(key)` method call.
38    ///
39    /// # Examples
40    ///
41    /// ```
42    /// use std::collections::HashMap;
43    ///
44    /// use try_insert_ext::EntryInsertExt;
45    ///
46    /// let mut map: HashMap<&str, usize> = HashMap::new();
47    ///
48    /// let e: Result<&mut usize, ()> = map
49    ///     .entry("poneyland")
50    ///     .or_try_insert_with_key(|_| Err(()));
51    /// assert!(e.is_err());
52    /// map
53    ///     .entry("poneyland")
54    ///     .or_try_insert_with_key::<_, ()>(|key| Ok(key.chars().count()));
55    ///
56    /// assert_eq!(map["poneyland"], 9);
57    /// ```
58    fn or_try_insert_with_key<F: FnOnce(&K) -> Result<V, E>, E>(
59        self,
60        default: F,
61    ) -> Result<&'a mut V, E>;
62}
63
64#[cfg(feature = "std")]
65impl<'a, K, V> EntryInsertExt<'a, K, V> for btree_map::Entry<'a, K, V>
66where
67    K: Ord,
68{
69    #[inline]
70    fn or_try_insert_with<F: FnOnce() -> Result<V, E>, E>(
71        self,
72        default: F,
73    ) -> Result<&'a mut V, E> {
74        match self {
75            Self::Occupied(entry) => Ok(entry.into_mut()),
76            Self::Vacant(entry) => default().map(|default| entry.insert(default)),
77        }
78    }
79
80    #[inline]
81    fn or_try_insert_with_key<F: FnOnce(&K) -> Result<V, E>, E>(
82        self,
83        default: F,
84    ) -> Result<&'a mut V, E> {
85        match self {
86            Self::Occupied(entry) => Ok(entry.into_mut()),
87            Self::Vacant(entry) => default(entry.key()).map(|default| entry.insert(default)),
88        }
89    }
90}
91
92#[cfg(feature = "std")]
93impl<'a, K, V> EntryInsertExt<'a, K, V> for hash_map::Entry<'a, K, V> {
94    #[inline]
95    fn or_try_insert_with<F: FnOnce() -> Result<V, E>, E>(
96        self,
97        default: F,
98    ) -> Result<&'a mut V, E> {
99        match self {
100            Self::Occupied(entry) => Ok(entry.into_mut()),
101            Self::Vacant(entry) => default().map(|default| entry.insert(default)),
102        }
103    }
104
105    #[inline]
106    fn or_try_insert_with_key<F: FnOnce(&K) -> Result<V, E>, E>(
107        self,
108        default: F,
109    ) -> Result<&'a mut V, E> {
110        match self {
111            Self::Occupied(entry) => Ok(entry.into_mut()),
112            Self::Vacant(entry) => default(entry.key()).map(|default| entry.insert(default)),
113        }
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    #[cfg(feature = "std")]
120    mod btree_map {
121        use std::collections::BTreeMap;
122
123        use crate::EntryInsertExt;
124
125        #[test]
126        fn it_works_when_occupied() {
127            let mut map = BTreeMap::new();
128            map.insert(0, 0);
129            *map.entry(0).or_try_insert_with::<_, ()>(|| Ok(3)).unwrap() += 1;
130            *map.entry(0).or_try_insert_with(|| Err(())).unwrap() += 1;
131            assert_eq!(map.get(&0), Some(&2));
132        }
133
134        #[test]
135        fn it_works_when_vacant() {
136            let mut map = BTreeMap::new();
137            map.entry(0).or_try_insert_with::<_, ()>(|| Ok(1)).unwrap();
138            assert_eq!(map.get(&0), Some(&1));
139        }
140
141        #[test]
142        fn it_errors() {
143            let mut map = BTreeMap::<i32, i32>::new();
144            assert_eq!(map.entry(0).or_try_insert_with(|| Err(())), Err(()));
145        }
146    }
147
148    #[cfg(feature = "std")]
149    mod hash_map {
150        use std::collections::HashMap;
151
152        use crate::EntryInsertExt;
153
154        #[test]
155        fn it_works_when_occupied() {
156            let mut map = HashMap::new();
157            map.insert(0, 0);
158            *map.entry(0).or_try_insert_with::<_, ()>(|| Ok(3)).unwrap() += 1;
159            *map.entry(0).or_try_insert_with(|| Err(())).unwrap() += 1;
160            assert_eq!(map.get(&0), Some(&2));
161        }
162
163        #[test]
164        fn it_works_when_vacant() {
165            let mut map = HashMap::new();
166            map.entry(0).or_try_insert_with::<_, ()>(|| Ok(1)).unwrap();
167            assert_eq!(map.get(&0), Some(&1));
168        }
169
170        #[test]
171        fn it_errors() {
172            let mut map = HashMap::<i32, i32>::new();
173            assert_eq!(map.entry(0).or_try_insert_with(|| Err(())), Err(()));
174        }
175    }
176}
try_insert_ext/entry.rs

try_insert_ext/
entry.rs
