/// Extends `Option` with `get_or_try_insert_with`.
pub trait OptionInsertExt<T> {
    /// If the option is `None`, computes the value from `f`. If `f` returns
    /// `Ok`, inserts the value. If `f` returns `Err`, returns the error. If
    /// there is no error, returns a mutable reference to the contained value.
    ///
    /// # Examples
    ///
    /// ```
    /// use try_insert_ext::OptionInsertExt;
    ///
    /// let mut x = None;
    ///
    /// {
    ///     let e: Result<&mut u32, ()> = x.get_or_try_insert_with(|| Err(()));
    ///     assert!(e.is_err());
    ///     let y: Result<&mut u32, ()> = x.get_or_try_insert_with(|| Ok(5));
    ///     assert_eq!(y, Ok(&mut 5));
    ///
    ///     *y.unwrap() = 7;
    /// }
    ///
    /// assert_eq!(x, Some(7));
    /// ```
    fn get_or_try_insert_with<F, E>(&mut self, f: F) -> Result<&mut T, E>
    where
        F: FnOnce() -> Result<T, E>;
}

impl<T> OptionInsertExt<T> for Option<T> {
    #[inline]
    fn get_or_try_insert_with<F, E>(&mut self, f: F) -> Result<&mut T, E>
    where
        F: FnOnce() -> Result<T, E>,
    {
        if let None = *self {
            *self = Some(f()?);
        }
        match self {
            Some(v) => Ok(v),
            None => unsafe { core::hint::unreachable_unchecked() },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::OptionInsertExt;

    #[test]
    fn it_works_when_some() {
        let mut opt = Some(0);
        *opt.get_or_try_insert_with::<_, ()>(|| Ok(3)).unwrap() += 1;
        *opt.get_or_try_insert_with::<_, ()>(|| Err(())).unwrap() += 1;
        assert_eq!(opt, Some(2));
    }

    #[test]
    fn it_works_when_none() {
        let mut opt = None;
        opt.get_or_try_insert_with::<_, ()>(|| Ok(1)).unwrap();
        assert_eq!(opt, Some(1));
    }

    #[test]
    fn it_errors() {
        let mut opt: Option<i32> = None;
        assert_eq!(opt.get_or_try_insert_with(|| Err(())), Err(()));
    }
}