#
cargo/util/
once.rs

    1//! Extension functions for [`std::sync::OnceLock`] / [`std::cell::OnceCell`]
    2//!
    3//! This adds polyfills for functionality in `lazycell` that is not stable within `std`.
    4
    5pub trait OnceExt {
    6    type T;
    7
    8    /// This might run `f` multiple times if different threads start initializing at once.
    9    fn try_borrow_with<F, E>(&self, f: F) -> Result<&Self::T, E>
    10    where
    11        F: FnOnce() -> Result<Self::T, E>;
    12
    13    /// This might run `f` multiple times if different threads start initializing at once.
    14    fn try_borrow_mut_with<F, E>(&mut self, f: F) -> Result<&mut Self::T, E>
    15    where
    16        F: FnOnce() -> Result<Self::T, E>;
    17
    18    fn replace(&mut self, new_value: Self::T) -> Option<Self::T>;
    19
    20    fn filled(&self) -> bool;
    21}
    22
    23impl<T> OnceExt for std::sync::OnceLock<T> {
    24    type T = T;
    25
    26    fn try_borrow_with<F, E>(&self, f: F) -> Result<&T, E>
    27    where
    28        F: FnOnce() -> Result<T, E>,
    29    {
    30        if let Some(value) = self.get() {
    31            return Ok(value);
    32        }
    33
    34        // This is not how the unstable `OnceLock::get_or_try_init` works. That only starts `f` if
    35        // no other `f` is executing and the value is not initialized. However, correctly implementing that is
    36        // hard (one has properly handle panics in `f`) and not doable with the stable API of `OnceLock`.
    37        let value = f()?;
    38        // Another thread might have initialized `self` since we checked that `self.get()` returns `None`. If this is the case, `self.set()`
    39        // returns an error. We ignore it and return the value set by the other
    40        // thread.
    41        let _ = self.set(value);
    42        Ok(self.get().unwrap())
    43    }
    44
    45    fn try_borrow_mut_with<F, E>(&mut self, f: F) -> Result<&mut T, E>
    46    where
    47        F: FnOnce() -> Result<T, E>,
    48    {
    49        let value = if let Some(value) = self.take() {
    50            value
    51        } else {
    52            // This is not how the unstable `OnceLock::get_or_try_init` works. That only starts `f` if
    53            // no other `f` is executing and the value is not initialized. However, correctly implementing that is
    54            // hard (one has properly handle panics in `f`) and not doable with the stable API of `OnceLock`.
    55            f()?
    56        };
    57        // Another thread might have initialized `self` since we checked that `self.get()` returns `None`. If this is the case, `self.set()`
    58        // returns an error. We ignore it and return the value set by the other
    59        // thread.
    60        let _ = self.set(value);
    61        Ok(self.get_mut().unwrap())
    62    }
    63
    64    fn replace(&mut self, new_value: T) -> Option<T> {
    65        if let Some(value) = self.get_mut() {
    66            Some(std::mem::replace(value, new_value))
    67        } else {
    68            let result = self.set(new_value);
    69            assert!(result.is_ok());
    70            None
    71        }
    72    }
    73
    74    fn filled(&self) -> bool {
    75        self.get().is_some()
    76    }
    77}
    78
    79impl<T> OnceExt for std::cell::OnceCell<T> {
    80    type T = T;
    81
    82    fn try_borrow_with<F, E>(&self, f: F) -> Result<&T, E>
    83    where
    84        F: FnOnce() -> Result<T, E>,
    85    {
    86        if let Some(value) = self.get() {
    87            return Ok(value);
    88        }
    89
    90        // This is not how the unstable `OnceLock::get_or_try_init` works. That only starts `f` if
    91        // no other `f` is executing and the value is not initialized. However, correctly implementing that is
    92        // hard (one has properly handle panics in `f`) and not doable with the stable API of `OnceLock`.
    93        let value = f()?;
    94        // Another thread might have initialized `self` since we checked that `self.get()` returns `None`. If this is the case, `self.set()`
    95        // returns an error. We ignore it and return the value set by the other
    96        // thread.
    97        let _ = self.set(value);
    98        Ok(self.get().unwrap())
    99    }
    100
    101    fn try_borrow_mut_with<F, E>(&mut self, f: F) -> Result<&mut T, E>
    102    where
    103        F: FnOnce() -> Result<T, E>,
    104    {
    105        let value = if let Some(value) = self.take() {
    106            value
    107        } else {
    108            // This is not how the unstable `OnceLock::get_or_try_init` works. That only starts `f` if
    109            // no other `f` is executing and the value is not initialized. However, correctly implementing that is
    110            // hard (one has properly handle panics in `f`) and not doable with the stable API of `OnceLock`.
    111            f()?
    112        };
    113        // Another thread might have initialized `self` since we checked that `self.get()` returns `None`. If this is the case, `self.set()`
    114        // returns an error. We ignore it and return the value set by the other
    115        // thread.
    116        let _ = self.set(value);
    117        Ok(self.get_mut().unwrap())
    118    }
    119
    120    fn replace(&mut self, new_value: T) -> Option<T> {
    121        if let Some(value) = self.get_mut() {
    122            Some(std::mem::replace(value, new_value))
    123        } else {
    124            let result = self.set(new_value);
    125            assert!(result.is_ok());
    126            None
    127        }
    128    }
    129
    130    fn filled(&self) -> bool {
    131        self.get().is_some()
    132    }
    133}