eager_static.rs

use core::sync::atomic::{*, Ordering::*};
use core::mem::MaybeUninit;
use core::cell::UnsafeCell;

/// A macro similar to `lazy_static::lazy_static!`, but that will attempt to
/// eagerly initialize the static value using [`on_startup!`](crate::on_startup).
#[macro_export]
macro_rules! eager_static {
    ($(
        $(#[$attr:meta])*
        $v:vis static ref $N:ident : $T:ty = $init:expr;
    )*) => {
        $(
            $(#[$attr])*
            $v static $N: $crate::_private::EagerStatic<$T> =
                $crate::_private::EagerStatic::<$T>::__new(|| $init);
        )*
        $crate::on_startup!{
            $(let _: &$T = &*$N;)*
        }
    };
}

/// `EagerStatic` is a wrapper type that `Deref`s to the target type, and can be
/// used as such. It's mostly private, but is generated by the `eager_static!`
/// macro.
///
/// When you do:
///
/// ```
/// startup::eager_static! {
///     static ref FOO: Vec<i32> = vec![1, 2, 3, 4];
/// }
/// ```
///
/// The actual item that is produced is something like:
///
/// ```
/// # stringify! {
/// static FOO: startup::_private::EagerStatic<T> = ...;
/// # }
/// ```
pub struct EagerStatic<T, F = fn() -> T> {
    state: AtomicUsize,
    inner: UnsafeCell<MaybeUninit<T>>,
    init_fn: F,
}

const UNINIT: usize = 0;
const INITIALIZING: usize = 1;
const INITIALIZED: usize = 2;
// FIXME: poisoning sucks.
const PANICED: usize = 3;

impl<T, F> EagerStatic<T, F> {
    #[inline]
    #[doc(hidden)]
    pub const fn __new(init_fn: F) -> Self {
        EagerStatic {
            state: AtomicUsize::new(0),
            inner: UnsafeCell::new(MaybeUninit::uninit()),
            init_fn
        }
    }
}

impl<T> core::ops::Deref for EagerStatic<T> {
    type Target = T;
    #[inline]
    fn deref(&self) -> &T {
        // use a double-check sorta scheme to avoid needing to RMW in the common case.
        let v = self.state.load(Acquire);
        if v != INITIALIZED {
            self.try_init();
            debug_assert_eq!(self.state.load(Acquire), INITIALIZED);
        }
        unsafe { &*self.inner.get().cast::<T>() }
    }
}

struct StateResetOnDrop<'a>(&'a AtomicUsize);
impl<'a> Drop for StateResetOnDrop<'a> {
    fn drop(&mut self) {
        self.0.store(PANICED, Release);
    }
}

#[inline(never)]
#[cold]
fn poisoned() -> ! {
    panic!("I don't really want to commit to supporting this behavior, but for now we poison on panic");
}

impl<T> EagerStatic<T> {
    // separated to discourage inlining
    fn try_init(&self) {
        match self.state.compare_exchange(UNINIT, INITIALIZING, Acquire, Relaxed) {
            Err(PANICED) => poisoned(),
            Err(INITIALIZED) => {},
            Err(INITIALIZING) => self.wait_for_initialized(),
            Ok(_) => unsafe {
                let panic_guard = StateResetOnDrop(&self.state);
                self.inner.get().write(MaybeUninit::new((self.init_fn)()));
                core::mem::forget(panic_guard);
                self.state.store(INITIALIZED, Release);
            }
            // could use unchecked unreachable in release code...
            Err(state) if cfg!(debug_assertions) => unreachable!("bug: {}", state),
            _ => {}
        }
    }

    // Realistically this shouldn't get called outside of init cycles or when
    // running on a platform with no support for static init. We just spin
    #[cold]
    #[inline(never)]
    fn wait_for_initialized(&self) {
        loop {
            // if we cared this could be load(Relaxed) w/ a fence(Acquire) before
            // returning. In practice this function should never be called
            // though.
            match self.state.load(Acquire) {
                INITIALIZED => break,
                INITIALIZING => core::hint::spin_loop(),
                PANICED => poisoned(),
                // could use unchecked unreachable in release code...
                state if cfg!(debug_assertions) => unreachable!("bug: {}", state),
                _ => {}
            }
        }
    }
}

unsafe impl<T: Send> Send for EagerStatic<T> {}
unsafe impl<T: Send + Sync> Sync for EagerStatic<T> {}