oncelock 0.1.0-alpha.0

//! Implements the [`OnceCell`] and [`LazyCell`] types.

use crate::polyfill::{FinishNonExhaustive, UncheckedUnwrap};
use core::cell::Cell;
use core::fmt::{Debug, Formatter};
use core::ops::Deref;

enum Void {}

/// A cell which can only be initialized once.
///
/// Like [`crate::sync::OnceLock`], but not thread-safe.
pub struct OnceCell<T>(Cell<Option<T>>);
impl<T> OnceCell<T> {
    /// Create an uninitialized cell.
    #[inline]
    pub const fn new() -> OnceCell<T> {
        OnceCell(Cell::new(None))
    }

    /// Get a reference to the current value, or `None` if uninitialized.
    #[inline]
    pub fn get(&self) -> Option<&T> {
        // SAFETY:  Writing to tag is atomic. If initialized, will not be written to again
        unsafe { (&*self.0.as_ptr()).as_ref() }
    }

    /// Get a reference to the current value if initialed,
    /// invoking the specified callback otherwise.
    #[inline]
    #[cfg_attr(has_track_caller, track_caller)]
    pub fn get_or_init(&self, func: impl FnOnce() -> T) -> &'_ T {
        self.get()
            .unwrap_or_else(|| match self.initialize::<Void, _>(|| Ok(func())) {
                Ok(value) => value,
                Err(e) => match e {},
            })
    }

    /// Get a reference to the current value if initialized,
    /// invoking the specified callback otherwise.
    ///
    /// # Errors
    /// Any errors from callback function will be propagated upwards,
    /// and in that case the cell will remain unchanged.
    #[inline]
    #[cfg_attr(has_track_caller, track_caller)]
    pub fn get_or_try_init<E>(&self, func: impl FnOnce() -> Result<T, E>) -> Result<&'_ T, E> {
        match self.get() {
            Some(existing) => Ok(existing),
            None => self.initialize::<E, _>(func),
        }
    }

    /// Actually perform initialization.
    ///
    /// Should only be called if the value is uninitialized,
    /// although this is a correctness requirement not a safety requirement.
    ///
    /// # Panics
    /// Will panic on double initialization.
    #[cold]
    #[cfg_attr(has_track_caller, track_caller)]
    // NOTE: `impl FnOnce()` triggers error on 1.31 when mixed with explicit type params
    fn initialize<E, F: FnOnce () -> Result<T, E>>(&self, func: F) -> Result<&T, E> {
        debug_assert!(self.get().is_none(), "only call when uninitialized");
        let res = func()?;
        self.set(res)
            .unwrap_or_else(|_| panic!("double initialization"));
        // SAFETY: At this point, we are guaranteed to be uninitialized
        unsafe { Ok(self.get().polyfill_unwrap_unchecked()) }
    }

    /// Initializes the contents of the cell to value,
    /// returning an error if already initialized to some other value.
    ///
    /// # Errors
    /// Returns `Ok(())` if successfully initialized,
    /// or `Err(given_Value)` if already initialized.
    #[inline]
    pub fn set(&self, value: T) -> Result<(), T> {
        if self.get().is_none() {
            // SAFETY: Just checked we are uninitialized,
            // so there are no outstanding references to the value
            unsafe { core::ptr::write(self.0.as_ptr() as *mut T, value) }
            Ok(())
        } else {
            Err(value)
        }
    }

    /// Takes the value out of this `OnceCell`,
    /// moving it back to an uninitialized state.
    ///
    /// Returns `None` if not currently initialized.
    ///
    /// This is safe because it takes a mutable reference.
    #[inline]
    pub fn take(&mut self) -> Option<T> {
        self.0.get_mut().take()
    }

    /// Consumes the cell, returning the wrapped value (if any).
    pub fn into_inner(self) -> Option<T> {
        self.0.into_inner()
    }
}
// SAFETY: As long as there are no references, this is safe
unsafe impl<T: Send> Send for OnceCell<T> {}
impl<T> From<T> for OnceCell<T> {
    fn from(value: T) -> Self {
        OnceCell(Cell::new(Some(value)))
    }
}
impl<T: Clone> Clone for OnceCell<T> {
    fn clone(&self) -> Self {
        match self.get() {
            Some(existing) => OnceCell::from(existing.clone()),
            None => OnceCell::new(),
        }
    }
}
impl<T> Default for OnceCell<T> {
    #[inline]
    fn default() -> Self {
        Self::new()
    }
}
impl<T: Debug> Debug for OnceCell<T> {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        f.debug_tuple("OnceCell").field(&self.get()).finish()
    }
}

enum LazyState<T, F> {
    Initialized(T),
    Poisoned,
    InProgress,
    Uninitialized(F),
}
impl<T, F> Debug for LazyState<T, F> {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        match self {
            LazyState::Initialized(_) => f.write_str("Initialized"),
            LazyState::Poisoned => f.write_str("Poisoned"),
            LazyState::InProgress => f.write_str("InProgress"),
            LazyState::Uninitialized(_) => f.write_str("Uninitialized"),
        }
    }
}

/// A value which is initialized on the first access.
///
/// Wraps a [`OnceCell`] with a pre-fixed initialization function.
pub struct LazyCell<T, F = fn() -> T> {
    state: Cell<LazyState<T, F>>,
}
impl<T, F> LazyCell<T, F> {
    /// Create a new [`LazyCell`] using the specified initialization function.
    ///
    /// Logically, this should require [`F: FnOnce() -> T`](FnOnce),
    /// but it cannot since that is not supported on the MSRV.
    #[inline]
    pub const fn new(func: F) -> Self {
        LazyCell {
            state: Cell::new(LazyState::Uninitialized(func)),
        }
    }

    /// Get the value stored in this cell,
    /// or `None` if not yet initialized.
    #[inline]
    pub fn get(this: &Self) -> Option<&T> {
        // SAFETY: Reading the tag is atomic,
        // and once initialized will never be modified further
        match unsafe { &*this.state.as_ptr() } {
            LazyState::Initialized(ref result) => Some(result),
            LazyState::Poisoned | LazyState::InProgress | LazyState::Uninitialized(_) => None,
        }
    }
}
impl<T, F: FnOnce() -> T> LazyCell<T, F> {
    /// Get this value or invoke the initialization function.
    ///
    /// Equivalent to the [`Deref`] impl but more explicit.
    #[inline]
    #[cfg_attr(has_track_caller, track_caller)]
    pub fn force(this: &Self) -> &T {
        Self::get(this).unwrap_or_else(|| Self::force_fallback(this))
    }
    #[cold]
    #[cfg_attr(has_track_caller, track_caller)]
    fn force_fallback(this: &Self) -> &T {
        let func = {
            // SAFETY: Reading tag is atomic and not borrowed from unless fully initialized
            // The &mut reference is transient and not held while invoking the callback
            // because of the possibility of recursive init
            let state = unsafe { &mut *this.state.as_ptr() };
            match *state {
                LazyState::Initialized(_) => {
                    // SAFETY: Just verified we are fully initialized
                    return unsafe { Self::get(this).polyfill_unwrap_unchecked() };
                }
                LazyState::Poisoned => panic_poisoned(),
                LazyState::InProgress => panic_recurisve_init(),
                LazyState::Uninitialized(_) => {
                    match core::mem::replace(state, LazyState::InProgress) {
                        LazyState::Uninitialized(func) => func,
                        _ => {
                            // SAFETY: Just verified we are in the `unint` state
                            unsafe { core::hint::unreachable_unchecked() }
                        }
                    }
                }
            }
        };
        /// Sets the state to [`LazyState::Poisoned`] if a panic occurs.
        ///
        /// Not necessary for safety because
        struct PanicGuard<'a, T, F>(&'a Cell<LazyState<T, F>>);
        impl<T, F> Drop for PanicGuard<'_, T, F> {
            fn drop(&mut self) {
                // We call core::mem::forget on success,
                // so only executed on panic
                self.0.set(LazyState::Poisoned);
            }
        }
        let guard = PanicGuard(&this.state);
        let result = func();
        core::mem::forget(guard);
        {
            // SAFETY: Reading the tag is atomic,
            // and we only borrow the result if fully initialized
            let state = unsafe { &mut *this.state.as_ptr() };
            match &*state {
                LazyState::Uninitialized(_) | LazyState::Poisoned | LazyState::Initialized(_) => {
                    if cfg!(debug_assertions) {
                        unreachable!("bad state {:?}", &*state)
                    } else {
                        // SAFETY: Can only successfully initialize if previously `InProgress`
                        unsafe { core::hint::unreachable_unchecked() }
                    }
                }
                LazyState::InProgress => {
                    // NOTE: This is slightly faster than an unconditional override
                    // because it avoids calling the Drop implementation in release mode
                    *state = LazyState::Initialized(result);
                }
            }
        }
        // SAFETY: Just finished initalization
        unsafe { Self::get(this).polyfill_unwrap_unchecked() }
    }
}
impl<T, F: FnOnce() -> T> Deref for LazyCell<T, F> {
    type Target = T;

    #[inline]
    fn deref(&self) -> &Self::Target {
        Self::force(self)
    }
}
impl<T: Debug, F> Debug for LazyCell<T, F> {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("LazyCell")
            .field("value", &Self::get(self))
            .polyfill_finish_non_exhaustive()
    }
}
// SAFETY: Valid to Send if our data is
unsafe impl<T: Send, F: Send> Send for LazyCell<T, F> {}

#[cold]
#[inline(never)]
#[cfg_attr(has_track_caller, track_caller)]
fn panic_poisoned() -> ! {
    panic!("LazyCell instance has previously been poisoned")
}

#[cold]
#[inline(never)]
#[cfg_attr(has_track_caller, track_caller)]
fn panic_recurisve_init() -> ! {
    panic!("LazyCell detected recursive initialization")
}