try_insert_ext/
option.rs

1/// Extends `Option` with `get_or_try_insert_with`.
2pub trait OptionInsertExt<T> {
3    /// If the option is `None`, computes the value from `f`. If `f` returns
4    /// `Ok`, inserts the value. If `f` returns `Err`, returns the error. If
5    /// there is no error, returns a mutable reference to the contained value.
6    ///
7    /// # Examples
8    ///
9    /// ```
10    /// use try_insert_ext::OptionInsertExt;
11    ///
12    /// let mut x = None;
13    ///
14    /// {
15    ///     let e: Result<&mut u32, ()> = x.get_or_try_insert_with(|| Err(()));
16    ///     assert!(e.is_err());
17    ///     let y: Result<&mut u32, ()> = x.get_or_try_insert_with(|| Ok(5));
18    ///     assert_eq!(y, Ok(&mut 5));
19    ///
20    ///     *y.unwrap() = 7;
21    /// }
22    ///
23    /// assert_eq!(x, Some(7));
24    /// ```
25    fn get_or_try_insert_with<F, E>(&mut self, f: F) -> Result<&mut T, E>
26    where
27        F: FnOnce() -> Result<T, E>;
28}
29
30impl<T> OptionInsertExt<T> for Option<T> {
31    #[inline]
32    fn get_or_try_insert_with<F, E>(&mut self, f: F) -> Result<&mut T, E>
33    where
34        F: FnOnce() -> Result<T, E>,
35    {
36        if let None = *self {
37            *self = Some(f()?);
38        }
39        match self {
40            Some(v) => Ok(v),
41            None => unsafe { core::hint::unreachable_unchecked() },
42        }
43    }
44}
45
46#[cfg(test)]
47mod tests {
48    use super::OptionInsertExt;
49
50    #[test]
51    fn it_works_when_some() {
52        let mut opt = Some(0);
53        *opt.get_or_try_insert_with::<_, ()>(|| Ok(3)).unwrap() += 1;
54        *opt.get_or_try_insert_with::<_, ()>(|| Err(())).unwrap() += 1;
55        assert_eq!(opt, Some(2));
56    }
57
58    #[test]
59    fn it_works_when_none() {
60        let mut opt = None;
61        opt.get_or_try_insert_with::<_, ()>(|| Ok(1)).unwrap();
62        assert_eq!(opt, Some(1));
63    }
64
65    #[test]
66    fn it_errors() {
67        let mut opt: Option<i32> = None;
68        assert_eq!(opt.get_or_try_insert_with(|| Err(())), Err(()));
69    }
70}