realloc 0.1.1

use core::cell::UnsafeCell;
use core::mem::MaybeUninit;
use core::ops::Deref;
use core::sync::atomic::{AtomicU8, Ordering};

/// A thread-safe `const`-compatible lazy initializing container, intended for
/// creating allocators in `static` items.
///
/// Implements <code>[Deref]\<T\></code>, which delegates to [`get`](Static::get).
pub struct Static<A> {
    /// The function to be used during initialization.
    ///
    /// If changed post-init, does nothing.
    pub init: fn() -> A,
    allocator: UnsafeCell<MaybeUninit<A>>,
    state: AtomicU8,
}

// SAFETY: the locking ensure no memory unsafety in regards to the inner
// allocator itself, and no mutable references to it are given out
unsafe impl<A> Sync for Static<A> {}

const UNINIT: u8 = 0;
const LOCKED: u8 = 1;
const INIT: u8 = 2;

impl<A> Static<A> {
    /// Create a new `Static`.
    ///
    /// Note that `init` will not be called until [`get`](Self::get) is called,
    /// or equivalently, the `Static` is dereferenced.
    pub const fn new(init: fn() -> A) -> Self
    where
        A: crate::Allocator,
    {
        Self {
            init,
            allocator: UnsafeCell::new(MaybeUninit::uninit()),
            state: AtomicU8::new(UNINIT),
        }
    }

    /// Return a reference to the inner value, initializing it if it hasn't
    /// already been.
    pub fn get(&self) -> &A {
        while self.state.load(Ordering::SeqCst) == LOCKED {
            core::hint::spin_loop();
        }

        if self.state.swap(LOCKED, Ordering::SeqCst) == INIT {
            self.state.store(INIT, Ordering::SeqCst);
            // SAFETY: this is guaranteed to never be mutated again
            return unsafe { &*self.allocator.get().cast() };
        }

        // SAFETY: guaranteed to be no outstanding references to it, since if anybody
        // get called this earlier, the above condition would've triggered and early-returned
        let alloc = unsafe { &mut *self.allocator.get() };
        alloc.write((self.init)());
        self.state.store(INIT, Ordering::SeqCst);

        // SAFETY: it was just initialized
        unsafe { alloc.assume_init_ref() }
    }
}

impl<A> Deref for Static<A> {
    type Target = A;

    fn deref(&self) -> &Self::Target {
        self.get()
    }
}