//! Thread-safe, non-blocking, "first one wins" flavor of `OnceCell`.

//!

//! If two threads race to initialize a type from the `race` module, they

//! don't block, execute initialization function together, but only one of

//! them stores the result.

//!

//! This module does not require `std` feature.

//!

//! # Atomic orderings

//!

//! All types in this module use `Acquire` and `Release`

//! [atomic orderings](Ordering) for all their operations. While this is not

//! strictly necessary for types other than `OnceBox`, it is useful for users as

//! it allows them to be certain that after `get` or `get_or_init` returns on

//! one thread, any side-effects caused by the setter thread prior to them

//! calling `set` or `get_or_init` will be made visible to that thread; without

//! it, it's possible for it to appear as if they haven't happened yet from the

//! getter thread's perspective. This is an acceptable tradeoff to make since

//! `Acquire` and `Release` have very little performance overhead on most

//! architectures versus `Relaxed`.

#[cfg(not(feature = "portable-atomic"))]

use core::sync::atomic;

#[cfg(feature = "portable-atomic")]

use portable_atomic as atomic;

use atomic::{AtomicPtr, AtomicUsize, Ordering};

use core::cell::UnsafeCell;

use core::marker::PhantomData;

use core::num::NonZeroUsize;

use core::ptr;

/// A thread-safe cell which can be written to only once.

#[derive(Default, Debug)]

pub struct OnceNonZeroUsize {

    inner: AtomicUsize,

}

impl OnceNonZeroUsize {

    /// Creates a new empty cell.

    #[inline]

    pub const fn new() -> OnceNonZeroUsize {

        OnceNonZeroUsize { inner: AtomicUsize::new(0) }

    }

    /// Gets the underlying value.

    #[inline]

    pub fn get(&self) -> Option<NonZeroUsize> {

        let val = self.inner.load(Ordering::Acquire);

        NonZeroUsize::new(val)

    }

    /// Sets the contents of this cell to `value`.

    ///

    /// Returns `Ok(())` if the cell was empty and `Err(())` if it was

    /// full.

    #[inline]

    pub fn set(&self, value: NonZeroUsize) -> Result<(), ()> {

        let exchange =

            self.inner.compare_exchange(0, value.get(), Ordering::AcqRel, Ordering::Acquire);

        match exchange {

            Ok(_) => Ok(()),

            Err(_) => Err(()),

        }

    }

    /// Gets the contents of the cell, initializing it with `f` if the cell was

    /// empty.

    ///

    /// If several threads concurrently run `get_or_init`, more than one `f` can

    /// be called. However, all threads will return the same value, produced by

    /// some `f`.

    pub fn get_or_init<F>(&self, f: F) -> NonZeroUsize

    where

        F: FnOnce() -> NonZeroUsize,

    {

        enum Void {}

        match self.get_or_try_init(|| Ok::<NonZeroUsize, Void>(f())) {

            Ok(val) => val,

            Err(void) => match void {},

        }

    }

    /// Gets the contents of the cell, initializing it with `f` if

    /// the cell was empty. If the cell was empty and `f` failed, an

    /// error is returned.

    ///

    /// If several threads concurrently run `get_or_init`, more than one `f` can

    /// be called. However, all threads will return the same value, produced by

    /// some `f`.

    pub fn get_or_try_init<F, E>(&self, f: F) -> Result<NonZeroUsize, E>

    where

        F: FnOnce() -> Result<NonZeroUsize, E>,

    {

        let val = self.inner.load(Ordering::Acquire);

        match NonZeroUsize::new(val) {

            Some(it) => Ok(it),

            None => self.init(f),

        }

    }

    #[cold]

    #[inline(never)]

    fn init<E>(&self, f: impl FnOnce() -> Result<NonZeroUsize, E>) -> Result<NonZeroUsize, E> {

        let mut val = f()?.get();

        let exchange = self.inner.compare_exchange(0, val, Ordering::AcqRel, Ordering::Acquire);

        if let Err(old) = exchange {

            val = old;

        }

        Ok(unsafe { NonZeroUsize::new_unchecked(val) })

    }

}

/// A thread-safe cell which can be written to only once.

#[derive(Default, Debug)]

pub struct OnceBool {

    inner: OnceNonZeroUsize,

}

impl OnceBool {

    /// Creates a new empty cell.

    #[inline]

    pub const fn new() -> OnceBool {

        OnceBool { inner: OnceNonZeroUsize::new() }

    }

    /// Gets the underlying value.

    #[inline]

    pub fn get(&self) -> Option<bool> {

        self.inner.get().map(OnceBool::from_usize)

    }

    /// Sets the contents of this cell to `value`.

    ///

    /// Returns `Ok(())` if the cell was empty and `Err(())` if it was

    /// full.

    #[inline]

    pub fn set(&self, value: bool) -> Result<(), ()> {

        self.inner.set(OnceBool::to_usize(value))

    }

    /// Gets the contents of the cell, initializing it with `f` if the cell was

    /// empty.

    ///

    /// If several threads concurrently run `get_or_init`, more than one `f` can

    /// be called. However, all threads will return the same value, produced by

    /// some `f`.

    pub fn get_or_init<F>(&self, f: F) -> bool

    where

        F: FnOnce() -> bool,

    {

        OnceBool::from_usize(self.inner.get_or_init(|| OnceBool::to_usize(f())))

    }

    /// Gets the contents of the cell, initializing it with `f` if

    /// the cell was empty. If the cell was empty and `f` failed, an

    /// error is returned.

    ///

    /// If several threads concurrently run `get_or_init`, more than one `f` can

    /// be called. However, all threads will return the same value, produced by

    /// some `f`.

    pub fn get_or_try_init<F, E>(&self, f: F) -> Result<bool, E>

    where

        F: FnOnce() -> Result<bool, E>,

    {

        self.inner.get_or_try_init(|| f().map(OnceBool::to_usize)).map(OnceBool::from_usize)

    }

    #[inline]

    fn from_usize(value: NonZeroUsize) -> bool {

        value.get() == 1

    }

    #[inline]

    fn to_usize(value: bool) -> NonZeroUsize {

        unsafe { NonZeroUsize::new_unchecked(if value { 1 } else { 2 }) }

    }

}

/// A thread-safe cell which can be written to only once.

pub struct OnceRef<'a, T> {

    inner: AtomicPtr<T>,

    ghost: PhantomData<UnsafeCell<&'a T>>,

}

// TODO: Replace UnsafeCell with SyncUnsafeCell once stabilized

unsafe impl<'a, T: Sync> Sync for OnceRef<'a, T> {}

impl<'a, T> core::fmt::Debug for OnceRef<'a, T> {

    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {

        write!(f, "OnceRef({:?})", self.inner)

    }

}

impl<'a, T> Default for OnceRef<'a, T> {

    fn default() -> Self {

        Self::new()

    }

}

impl<'a, T> OnceRef<'a, T> {

    /// Creates a new empty cell.

    pub const fn new() -> OnceRef<'a, T> {

        OnceRef { inner: AtomicPtr::new(ptr::null_mut()), ghost: PhantomData }

    }

    /// Gets a reference to the underlying value.

    pub fn get(&self) -> Option<&'a T> {

        let ptr = self.inner.load(Ordering::Acquire);

        unsafe { ptr.as_ref() }

    }

    /// Sets the contents of this cell to `value`.

    ///

    /// Returns `Ok(())` if the cell was empty and `Err(value)` if it was

    /// full.

    pub fn set(&self, value: &'a T) -> Result<(), ()> {

        let ptr = value as *const T as *mut T;

        let exchange =

            self.inner.compare_exchange(ptr::null_mut(), ptr, Ordering::AcqRel, Ordering::Acquire);

        match exchange {

            Ok(_) => Ok(()),

            Err(_) => Err(()),

        }

    }

    /// Gets the contents of the cell, initializing it with `f` if the cell was

    /// empty.

    ///

    /// If several threads concurrently run `get_or_init`, more than one `f` can

    /// be called. However, all threads will return the same value, produced by

    /// some `f`.

    pub fn get_or_init<F>(&self, f: F) -> &'a T

    where

        F: FnOnce() -> &'a T,

    {

        enum Void {}

        match self.get_or_try_init(|| Ok::<&'a T, Void>(f())) {

            Ok(val) => val,

            Err(void) => match void {},

        }

    }

    /// Gets the contents of the cell, initializing it with `f` if

    /// the cell was empty. If the cell was empty and `f` failed, an

    /// error is returned.

    ///

    /// If several threads concurrently run `get_or_init`, more than one `f` can

    /// be called. However, all threads will return the same value, produced by

    /// some `f`.

    pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&'a T, E>

    where

        F: FnOnce() -> Result<&'a T, E>,

    {

        let mut ptr = self.inner.load(Ordering::Acquire);

        if ptr.is_null() {

            // TODO replace with `cast_mut` when MSRV reaches 1.65.0 (also in `set`)

            ptr = f()? as *const T as *mut T;

            let exchange = self.inner.compare_exchange(

                ptr::null_mut(),

                ptr,

                Ordering::AcqRel,

                Ordering::Acquire,

            );

            if let Err(old) = exchange {

                ptr = old;

            }

        }

        Ok(unsafe { &*ptr })

    }

    /// ```compile_fail

    /// use once_cell::race::OnceRef;

    ///

    /// let mut l = OnceRef::new();

    ///

    /// {

    ///     let y = 2;

    ///     let mut r = OnceRef::new();

    ///     r.set(&y).unwrap();

    ///     core::mem::swap(&mut l, &mut r);

    /// }

    ///

    /// // l now contains a dangling reference to y

    /// eprintln!("uaf: {}", l.get().unwrap());

    /// ```

    fn _dummy() {}

}

#[cfg(feature = "alloc")]

pub use self::once_box::OnceBox;

#[cfg(feature = "alloc")]

mod once_box {

    use super::atomic::{AtomicPtr, Ordering};

    use core::{marker::PhantomData, ptr};

    use alloc::boxed::Box;

    /// A thread-safe cell which can be written to only once.

    pub struct OnceBox<T> {

        inner: AtomicPtr<T>,

        ghost: PhantomData<Option<Box<T>>>,

    }

    impl<T> core::fmt::Debug for OnceBox<T> {

        fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {

            write!(f, "OnceBox({:?})", self.inner.load(Ordering::Relaxed))

        }

    }

    impl<T> Default for OnceBox<T> {

        fn default() -> Self {

            Self::new()

        }

    }

    impl<T> Drop for OnceBox<T> {

        fn drop(&mut self) {

            let ptr = *self.inner.get_mut();

            if !ptr.is_null() {

                drop(unsafe { Box::from_raw(ptr) })

            }

        }

    }

    impl<T> OnceBox<T> {

        /// Creates a new empty cell.

        pub const fn new() -> OnceBox<T> {

            OnceBox { inner: AtomicPtr::new(ptr::null_mut()), ghost: PhantomData }

        }

        /// Gets a reference to the underlying value.

        pub fn get(&self) -> Option<&T> {

            let ptr = self.inner.load(Ordering::Acquire);

            if ptr.is_null() {

                return None;

            }

            Some(unsafe { &*ptr })

        }

        /// Sets the contents of this cell to `value`.

        ///

        /// Returns `Ok(())` if the cell was empty and `Err(value)` if it was

        /// full.

        pub fn set(&self, value: Box<T>) -> Result<(), Box<T>> {

            let ptr = Box::into_raw(value);

            let exchange = self.inner.compare_exchange(

                ptr::null_mut(),

                ptr,

                Ordering::AcqRel,

                Ordering::Acquire,

            );

            if exchange.is_err() {

                let value = unsafe { Box::from_raw(ptr) };

                return Err(value);

            }

            Ok(())

        }

        /// Gets the contents of the cell, initializing it with `f` if the cell was

        /// empty.

        ///

        /// If several threads concurrently run `get_or_init`, more than one `f` can

        /// be called. However, all threads will return the same value, produced by

        /// some `f`.

        pub fn get_or_init<F>(&self, f: F) -> &T

        where

            F: FnOnce() -> Box<T>,

        {

            enum Void {}

            match self.get_or_try_init(|| Ok::<Box<T>, Void>(f())) {

                Ok(val) => val,

                Err(void) => match void {},

            }

        }

        /// Gets the contents of the cell, initializing it with `f` if

        /// the cell was empty. If the cell was empty and `f` failed, an

        /// error is returned.

        ///

        /// If several threads concurrently run `get_or_init`, more than one `f` can

        /// be called. However, all threads will return the same value, produced by

        /// some `f`.

        pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E>

        where

            F: FnOnce() -> Result<Box<T>, E>,

        {

            let mut ptr = self.inner.load(Ordering::Acquire);

            if ptr.is_null() {

                let val = f()?;

                ptr = Box::into_raw(val);

                let exchange = self.inner.compare_exchange(

                    ptr::null_mut(),

                    ptr,

                    Ordering::AcqRel,

                    Ordering::Acquire,

                );

                if let Err(old) = exchange {

                    drop(unsafe { Box::from_raw(ptr) });

                    ptr = old;

                }

            };

            Ok(unsafe { &*ptr })

        }

    }

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

    /// ```compile_fail

    /// struct S(*mut ());

    /// unsafe impl Sync for S {}

    ///

    /// fn share<T: Sync>(_: &T) {}

    /// share(&once_cell::race::OnceBox::<S>::new());

    /// ```

    fn _dummy() {}

}