oncelock 0.1.0-alpha.0

//! Implements the [`OnceLock`] and [`LazyLock`] types.

use crate::polyfill::{FinishNonExhaustive, UntaggedOptionCell};
use core::cell::UnsafeCell;
use core::fmt::{Debug, Formatter};
use core::mem::ManuallyDrop;
use core::ops::Deref;

#[cfg(sync_backend = "parking_lot")]
mod parking_lot;
#[cfg(sync_backend = "parking_lot")]
use self::parking_lot as backend;
#[cfg(sync_backend = "std")]
mod std_sync;
#[cfg(sync_backend = "std")]
use self::std_sync as backend;
#[cfg(sync_backend = "spin")]
mod spin;
#[cfg(sync_backend = "spin")]
use self::spin as backend;

/// If this is `true`, then [`OnceLock::get_or_init`] may be poisoned on panic.
///
/// This is necessary when using versions of `std` before 1.51 or versions of `spin` before v0.10,
/// which don't support [`std::sync::Once::call_once_force`].
pub const LEGACY_POISON_BEHAVIOR: bool = backend::LEGACY_POISON_BEHAVIOR;

/// A thread-safe value which is initialized at most once.
///
/// Emulates the [`std::sync::OnceLock`] type on old versions of rust.
///
/// Reading a non-None value out of `OnceCell` establishes a happens-before relationship with a corresponding write,
/// just like use of the [`std::sync::Once::call_once`] method.
pub struct OnceLock<T> {
    once: backend::Once,
    value: UntaggedOptionCell<T>,
}
impl<T> Default for OnceLock<T> {
    #[inline]
    fn default() -> Self {
        Self::new()
    }
}
impl<T> OnceLock<T> {
    /// Create a new `OnceLock` with no associated data.
    #[inline]
    #[rustversion::attr(since(1.32), const)] // const UnsafeCell::new
    pub fn new() -> Self {
        OnceLock {
            once: backend::Once::new(),
            value: UntaggedOptionCell::uninit(),
        }
    }

    /// Initializes the contents of this cell to the specified value.
    ///
    /// May block if initialization is currently in progress.
    ///
    /// # Errors
    /// Returns an error containing the passed value if the cell is already initialized,
    /// or `Ok(())` if the value
    #[inline]
    pub fn set(&self, value: T) -> Result<(), T> {
        let mut value = Some(value);
        self.get_or_init(|| value.take().unwrap());
        if let Some(failed) = value {
            Err(failed)
        } else {
            Ok(())
        }
    }

    /// Get the reference to the underlying value,
    /// triggering undefined behavior if uninitialized.
    ///
    /// # Safety
    /// Caller is responsible for ensuring the data is fully initialized,
    /// and for avoiding data races with concurrent initialization.
    ///
    /// # Panics
    /// In debug mode, this may panic if uninitialized.
    #[inline]
    pub unsafe fn get_unchecked(&self) -> &T {
        debug_assert!(self.get().is_some(), "UB: access uninitialized data");
        // SAFETY: Guaranteed by caller
        unsafe { &*self.value.get_ptr_unchecked() }
    }

    /// Get the value associated with this lock,
    /// or `None` if uninitialized.
    #[inline]
    pub fn get(&self) -> Option<&T> {
        if self.once.is_completed() {
            // SAFETY: Guaranteed to be initialized
            Some(unsafe { self.get_unchecked() })
        } else {
            None
        }
    }

    /// Get the contents of the cell if initialized,
    /// otherwise using the specified callback to initialize the value.
    ///
    /// Even in the presence of multiple threads, only one initialization function is ever invoked.
    ///
    /// # Panics
    /// If the [`LEGACY_POISON_BEHAVIOR`] flag is false, initialization will be retried if a thread panics.
    /// If [`LEGACY_POISON_BEHAVIOR`] is true, this will be poisoned if an initialization function panics.
    /// In that case, all feature calls will unconditionally panic (just like [`std::sync::Once::call_once`]).
    #[inline]
    pub fn get_or_init<F: FnOnce() -> T>(&self, f: F) -> &T {
        self.get().unwrap_or_else(|| self.init_fallback(f))
    }

    #[cold]
    fn init_fallback(&self, func: impl FnOnce() -> T) -> &T {
        self.once.call_once_force(|| {
            let value = func();
            // SAFETY: Guaranteed to exclude all other threads,
            // reads can only come after this `call_once` completes
            unsafe {
                self.value.write(value);
            }
        });
        // SAFETY: Completion of call_once means we are initialized
        unsafe { &*self.value.get_ptr_unchecked() }
    }

    /// Consumes the `OnceLock`, returning the wrapped value.
    ///
    /// Returns `None` if the cell was uninitialized.
    #[inline]
    pub fn into_inner(self) -> Option<T> {
        let mut this = ManuallyDrop::new(self);
        this.take()
    }

    /// Takes the value out of this `OnceLock`, moving it back to an uninitialized state.
    ///
    /// This is safe because it borrows the lock mutably.
    #[inline]
    pub fn take(&mut self) -> Option<T> {
        if self.get().is_some() {
            // A panic here would be UB
            // Neither of these methods can possibly panic,
            // but use a guard to make extra sure
            let guard = AbortGuard;
            struct AbortGuard;
            impl Drop for AbortGuard {
                fn drop(&mut self) {
                    // double panic due to abort
                    panic!("Panic here could cause invalid state to be seen");
                }
            }
            // SAFETY: Value is guaranteed to be initialized because `get` returns Some
            // Guard ensures that there is no panic before we reset our state
            let value = unsafe { core::ptr::read(self.value.get_ptr_unchecked()) };
            self.once = backend::Once::new();
            core::mem::forget(guard);
            Some(value)
        } else {
            None
        }
    }
}
impl<T> Drop for OnceLock<T> {
    fn drop(&mut self) {
        drop(self.take());
    }
}
impl<T: Clone> Clone for OnceLock<T> {
    fn clone(&self) -> Self {
        match self.get() {
            None => OnceLock::new(),
            Some(inner) => OnceLock::from(inner.clone()),
        }
    }
}
// SAFETY: Valid since `T: Send`
unsafe impl<T: Send> Send for OnceLock<T> {}
// SAFETY: Valid since `T: Sync`, all mutation is otherwise synchronized
unsafe impl<T: Sync> Sync for OnceLock<T> {}
impl<T> From<T> for OnceLock<T> {
    fn from(value: T) -> Self {
        let res = OnceLock::new();
        match res.set(value) {
            Ok(()) => res,
            Err(_) => {
                // SAFETY: Not possible since we are the only user
                unsafe { core::hint::unreachable_unchecked() }
            }
        }
    }
}

enum LazyState<T, F> {
    Initialized(T),
    InProgress,
    Uninitialized(F),
}
/// A lazily initialized value using a particular initialization function.
///
/// Unlike [`OnceLock`], this type is irrevocably poisoned if the initialization function panics.
pub struct LazyLock<T, F = fn() -> T> {
    once: backend::Once,
    state: UnsafeCell<LazyState<T, F>>,
}
impl<T, F> LazyLock<T, F> {
    /// Creates a new lazy value with the given initializing function.
    ///
    /// Logically, this should require [`F: FnOnce() -> T`](FnOnce),
    /// but it cannot since that is not supported on the MSRV.
    #[inline]
    #[rustversion::attr(since(1.32), const)] // const UnsafeCell::new
    pub fn new(func: F) -> Self {
        LazyLock {
            once: backend::Once::new(),
            state: UnsafeCell::new(LazyState::Uninitialized(func)),
        }
    }

    /// Get the value if the cell is initialized,
    /// or `None` if not initialized yet.
    #[inline]
    pub fn get(this: &Self) -> Option<&T> {
        if this.once.is_completed() {
            // SAFETY: The `is_completed` method only returns true
            // once initialization has finished.
            Some(unsafe { Self::get_unchecked(this) })
        } else {
            None
        }
    }

    #[inline]
    unsafe fn get_unchecked(this: &Self) -> &T {
        // SAFETY: Caller guarantees type is fully initialized,
        // so not possible to observe these states
        match unsafe { &*this.state.get() } {
            LazyState::Uninitialized(_) | LazyState::InProgress => {
                // SAFETY: Caller guarantees fully initialized, so not possible to encounter these
                unsafe { core::hint::unreachable_unchecked() }
            }
            LazyState::Initialized(ref value) => value,
        }
    }
}
impl<T, F: FnOnce() -> T> LazyLock<T, F> {
    /// Forces the evaluation of this lazy value and returns a reference to result.
    /// This is equivalent to the Deref impl, but is explicit.
    ///
    /// This method will block the calling thread if another initialization routine is currently running.
    ///
    /// This function is implemented in terms of [`OnceLock::get_or_init`].
    #[inline]
    #[cfg_attr(has_track_caller, track_caller)]
    pub fn force(this: &Self) -> &T {
        Self::get(this).unwrap_or_else(|| Self::init_fallback(this))
    }
    #[cold]
    #[cfg_attr(has_track_caller, track_caller)]
    fn init_fallback(this: &Self) -> &T {
        this.once.call_once(|| {
            // SAFETY: Have exclusive access
            // (recursive initialization will either deadlock or panic)
            let this = unsafe { &mut *this.state.get() };
            let func = match core::mem::replace(this, LazyState::InProgress) {
                LazyState::Initialized(_) | LazyState::InProgress => unreachable!(),
                LazyState::Uninitialized(func) => func,
            };
            *this = LazyState::Initialized(func());
        });
        // SAFETY: After call_once completes, we are fully initialized
        unsafe { Self::get_unchecked(this) }
    }
}
impl<T, F: FnOnce() -> T> Deref for LazyLock<T, F> {
    type Target = T;

    #[inline]
    fn deref(&self) -> &Self::Target {
        LazyLock::force(self)
    }
}

impl<T: Debug, F> Debug for LazyLock<T, F> {
    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
        f.debug_struct("LazyLock")
            .field("value", &Self::get(self))
            .polyfill_finish_non_exhaustive()
    }
}