#![doc(html_root_url = "https://docs.rs/singleton-cell/0.3.1")]
#![no_std]

use core::cell::UnsafeCell;
use core::marker::PhantomData;
use core::mem;
pub use singleton_trait::{Singleton, Exists, Erased};

/**
 * Macro-only code which must be public
 */
pub mod internal {
    pub use paste;
}

/**
 * SCell, or SingletonCell provides an interface of a
 * Ghost Cell, where the Key is allowed to be any singleton,
 * rather than a particular token type.
 *
 * As a result, the uniqueness for the key can be provided
 * by any means, and the key type / resulting cell may also
 * be 'static for long-lived data structures
 */
#[repr(transparent)]
pub struct SCell<Key: ?Sized, T: ?Sized> {
    _phantom: PhantomData<fn(Key)>,
    inner: UnsafeCell<T>,
}
/**
 * A more verbose alias for SCell
 */
pub type SingletonCell<Key, T> = SCell<Key, T>;
/**
 * An owned SCell is equivalent to its underlying value,
 * and can be converted between them, so SCell<Key, T> is Send
 * if and only if T is Send
 */
unsafe impl<Key, T: Send> Send for SCell<Key, T> {}

/**
 * A shared reference to SCell can access the underlying value
 * both via shared reference, and via mutable reference,
 * so it can be Sync only if T is both Sync and Send.
 *
 * SCell does not otherwise put any constraints on Sync
 * since all shared usages must use references to the Key,
 * which must be sent between threads as normal.
 */
unsafe impl<Key, T: Send + Sync> Sync for SCell<Key, T> {}

impl<Key: ?Sized, T: ?Sized> SCell<Key, T> {
    /**
     * Convert a unique reference to a value to a unique reference
     * to this cell type. These two types are equivalent when
     * accessed uniquely.
     */
    #[inline(always)]
    pub fn from_mut(t: &mut T) -> &mut Self {
        // Safety:
        // * The representations are transparent down to T
        // * unique access to T enables unique access to self
        unsafe { &mut *(t as *mut T as *mut Self) }
    }

    /**
     * Uniquely borrow this cell in order to access T mutably.
     *
     * This requires unique access to self, rather than using
     * a key. For shared access, see the `borrow_mut` method
     * instead.
     *
     * Because this requires unique access to self, and all
     * other borrows require at least a shared reference,
     * unique access to the underlying data is safe.
     */
    #[inline(always)]
    pub fn get_mut(&mut self) -> &mut T {
        self.inner.get_mut()
    }

    /**
     * Get a raw pointer to this cell's data.
     */
    #[inline(always)]
    pub fn as_ptr(&self) -> *mut T {
        self.inner.get()
    }
}

impl<Key: ?Sized, T> SCell<Key, T> {
    /**
     * Construct a new SCell from underlying data.
     *
     * See also `SCell::from_mut` which can be used when the
     * usage of SCell is only within the scope of mutable
     * borrows of the data.
     */
    #[inline(always)]
    pub fn new(t: T) -> Self {
        SCell {
            inner: UnsafeCell::new(t),
            _phantom: PhantomData,
        }
    }

    /**
     * Destruct the SCell and access the inner data.
     *
     * This requires ownership of the SCell and hence guarantees
     * no simultaneous access.
     */
    #[inline(always)]
    pub fn into_inner(self) -> T {
        self.inner.into_inner()
    }
}

impl<Key: ?Sized, T> SCell<Key, [T]> {
    /**
     * Returns a slice of cells from a cell of a slice.
     */
    #[inline(always)]
    pub fn as_slice_of_cells(&self) -> &[SCell<Key, T>] {
        // Safety:
        // * The representations are identical
        // * the output lifetime is constrained to the input lifetime
        unsafe { &*(self.as_ptr() as *mut [SCell<Key, T>]) }
    }
}

impl<Key: Singleton + ?Sized, T: ?Sized> SCell<Key, T> {
    /**
     * Borrow the data underlying this cell, using a reference to
     * the singleton Key type.
     *
     * This is safe because any mutable borrow must either come
     * from a mutable reference to this cell, or a mutable reference
     * to the key, which cannot exist since this takes a shared
     * borrow of both this cell and the key.
     */
    #[inline(always)]
    pub fn borrow<'a>(&'a self, _: impl Exists<&'a Key>) -> &'a T
    where
        Key: 'a,
    {
        // Safety:
        // By the contract for Singleton & Erased
        // the unique value of Key is borrowed immutably for 'a
        //
        // * Every borrow to T requires either &mut self or a &Key/&mut Key
        //   so the only other borrow available would be this method taking &Key,
        //   which is compatible with the shared reference output
        unsafe { &*self.inner.get() }
    }

    /**
     * Mutably borrow the data underlying this cell,
     * using a mutable reference to the singleton Key type.
     *
     * This is safe because it requires a unique borrow of the key,
     * and a shared borrow of self, which prevents all other borrows
     * into the data.
     */
    #[allow(clippy::mut_from_ref)] // Clippy can't see through the Exists
    #[inline(always)]
    pub fn borrow_mut<'a>(&'a self, _: impl Exists<&'a mut Key>) -> &'a mut T
    where
        Key: 'a,
    {
        // Safety:
        // By the contract for Singleton & Erased
        // the unique value of Key is borrowed mutably for 'a
        //
        // * Every borrow to T requires either &mut self or a &Key/&mut Key
        //   which are excluded by the borrows in the signature
        unsafe { &mut *self.inner.get() }
    }
}

impl<Key: Singleton + ?Sized, T> SCell<Key, T> {
    /**
     * Replace the value behind this cell with a new one.
     *
     * This mutates the data and hence requires unique access to the key
     * to ensure that no borrows are invalidated.
     */
    #[inline(always)]
    pub fn replace<'a>(&'a self, value: T, token: impl Exists<&'a mut Key>) -> T
    where
        Key: 'a,
    {
        mem::replace(self.borrow_mut(token), value)
    }

    /**
     * Replace the value behind this cell with the default one.
     *
     * This mutates the data and so requires unique access to the key
     * to ensure that no borrows are invalidated.
     */
    #[inline(always)]
    pub fn take<'a>(&'a self, token: impl Exists<&'a mut Key>) -> T
    where
        Key: 'a,
        T: Default,
    {
        mem::take(self.borrow_mut(token))
    }
}

/**
 * A temporary singleton token created by the with_token function.
 *
 * See the with_token function for details.
 */
pub struct Token<'brand>(PhantomData<fn(&'brand ()) -> &'brand ()>);
// SAFETY:
// We only construct this with a unique, invariant brand in with_token
// For everything else, we just defer to ghost cell's GhostToken
unsafe impl<'brand> Singleton for Token<'brand> {}
unsafe impl<'brand> Send for Token<'brand> {}
unsafe impl<'brand> Sync for Token<'brand> {}

/**
 * Generate a local token type which is guaranteed to be a singleton via
 * a unique lifetime brand. SCells can be created within this scope via
 * SCell::new or SCell::from_mut and used via this token.
 *
 * Because the 'brand lifetime is unique to this scope, neither this
 * token, nor any SCell's created with it as the key may escape the scope.
 *
 * It is most likely useful when cells are created with SCell::from_mut,
 * which means that cells can be used for data structures that don't
 * mention SCell or the 'brand.
 *
 * For a token type which does not have this restriction, see new_singleton
 * for one implementation.
 *
 * ```
 *  use singleton_cell::*;
 *
 *  let mut x = 0;
 *  with_token(|mut tok| {
 *      let cell_borrow = &*SCell::from_mut(&mut x);
 *      *cell_borrow.borrow_mut(&mut tok) += 2;
 *
 *      let cell_one = SCell::new(0);
 *      *cell_one.borrow_mut(&mut tok) += 1;
 *      cell_one.into_inner()
 *  });
 *  assert_eq!(2, x);
 * ```
 */
#[inline(always)]
pub fn with_token<R>(f: impl for<'brand> FnOnce(Token<'brand>) -> R) -> R {
    // SAFETY: 'brand can't leak as f is polymorphic over it
    let token = Token(PhantomData);
    f(token)
}

///
/// Generates a new global singleton type with a fallible constructor.
/// This is intended for cases where a single instance is created early
/// in the program and passed throughout.
///
/// This is one way to create a Singleton which is 'static and Send/Sync
///
/// ```
/// #[macro_use] extern crate singleton_cell;
/// new_singleton!(pub Token);
/// fn main() {
///     let tok = Token::new().expect("Main called only once");
///
///     // possibly combine with Erased borrows
///     use singleton_trait::Erased;
///     let tok = Erased::new(tok);
///
/// }
///
/// ```
///
/// This method will return Some the first time it is called, and None
/// any time thereafter. This does not establish any ordering between
/// multiple threads if they access at the same time.
///
// Credit to Yandros for significant help with the idea / designing
// the original macro using std::sync::Once instead of AtomicBool.
//
// We use AtomicBool here instead to remain no-std and simplify the
// resulting code given that we don't need most features of Once.
#[macro_export]
macro_rules! new_singleton {
    ( $pub:vis $Name:ident ) => {
        $crate::internal::paste::paste! {
            $pub use [< __ $Name _help__ >]::$Name;
            #[allow(nonstandard_style)]
            mod [< __ $Name _help__ >] {

                pub
                struct $Name(());

                impl $Name {

                    pub
                    fn new () -> Option<Self> {
                        let first_run = {
                            static ONCE: ::core::sync::atomic::AtomicBool =
                                ::core::sync::atomic::AtomicBool::new(false);
                            &ONCE
                        };
                        if first_run
                            // We don't need any ordering guarantees between accesses, only atomicity
                            .compare_exchange(false, true,
                                ::core::sync::atomic::Ordering::Relaxed,
                                ::core::sync::atomic::Ordering::Relaxed)
                            .is_ok()
                        {
                            // Safety:
                            // * ONCE is a static, hence it is a single memory location which is false at
                            // program start.
                            // * The program start happens before all reads
                            // * Only one CAS from false to true can succeed since it will never become
                            // false again
                            Some( $Name(()) )
                        } else {
                            None
                        }
                    }

                    #[inline(always)]
                    pub
                    fn new_erased() -> Option<$crate::Erased<Self>> {
                        Self::new().map(|x| $crate::Erased::new(x))
                    }
                }

                // Safety: construction is guarded by the CAS on a static above
                // So only one instance can be created per program
                unsafe
                impl $crate::Singleton for $Name {}
            }
        }
    };
}

#[cfg(test)]
mod tests {
    extern crate alloc;
    extern crate std;
    use crate as singleton_cell;
    use super::*;
    use alloc::string::String;

    #[test]
    fn ghost_trivial() {
        let mut i = String::new();

        with_token(|mut tok| {
            let i = SCell::from_mut(&mut i);
            i.get_mut().push_str("A");

            i.borrow_mut(&mut tok).push_str("B");
            let first = i.borrow(&tok).clone();

            i.borrow_mut(&mut tok).push_str("C");
            let last = i.borrow(&tok);

            assert_eq!("AB", &*first);
            assert_eq!("ABC", &*last);
        });

        assert_eq!("ABC", &*i);
    }

    #[inline(never)]
    fn write_to<Key: Singleton>(sc: &SCell<Key, String>, k: &mut Key) {
        sc.borrow_mut(k).push_str("!");
    }

    #[test]
    fn ghost_read_after_write() {
        let result = with_token(|mut tok| {
            let mut s = String::from("A");
            let s = SCell::from_mut(&mut s);

            let borrow = &*s;
            write_to(&*s, &mut tok);
            write_to(&*s, &mut tok);

            borrow.borrow(&tok).clone()
        });
        assert_eq!("A!!", &*result);
    }

    #[test]
    fn local_type_round_trip() {
        let local = ();
        let mut tok = {
            struct Local<'a>(&'a ());
            unsafe impl<'a> Singleton for Local<'a> {}
            Local(&local)
        };

        let mut s = String::from("A");
        let s_mut = &mut s;
        let s_cell = SCell::from_mut(s_mut);
        let s_cell2 = &*s_cell;

        write_to(s_cell2, &mut tok);

        assert_eq!("A!", s_cell.borrow_mut(&mut tok));
        write_to(s_cell2, &mut tok);
        assert_eq!("A!!", s_cell.borrow(&tok));
        assert_eq!("A!!", s_cell.get_mut());
        s_cell.get_mut().push_str("C");
        assert_eq!("A!!C", s_mut);
    }

    #[test]
    fn global_singleton_unique() {
        new_singleton!(pub A);
        let _ = A::new();
        assert!(A::new().is_none(), "Second call to A::new() should be None");
    }

    #[test]
    fn global_singleton_distinct() {
        new_singleton!(pub A);
        new_singleton!(pub B);
        let a: Option<A> = A::new();
        let b: Option<B> = B::new();

        assert!(a.is_some(), "A should be Some");
        assert!(b.is_some(), "B should be Some");
    }
}