#![feature(core_intrinsics)]

//
// Copyright (C) 2020 Nathan Sharp.
//
// This file is available under either the terms of the Apache License, Version
// 2.0 or the MIT License, at your discretion.
//

//! This crate provides the [`WaitCell`] type.
//!
//! A `WaitCell<T>` is a thread-safe cell which can only be written once while
//! shared. Attempts to read the value before it is written cause the reading
//! thread to [block]. In this way, it is like a "synchronous" [`Future`].
//!
//! Because the value can be written to only once while shared, `WaitCell`
//! provides access to its value by shared reference (`&T`). To use this
//! reference from multiple threads, `T` must be [`Sync`].
//!
//! # Example
//! ```
//! use std::time::Duration;
//! use std::thread::{ sleep, spawn };
//! use waitcell::WaitCell;
//!
//! static CELL: WaitCell<u64> = WaitCell::new();
//!
//! fn main() {
//!     let setter = spawn(|| {
//!         sleep(Duration::from_secs(1));
//!         CELL.init(42);
//!     });
//!
//!     // Prints "42".
//!     println!("{}", CELL.get());
//! #
//! #   setter.join();
//! }
//! ```
//!
//! # License
//! `waitcell` is licensed under the terms of the
//! [Apache License, Version 2.0][Apache2] or the [MIT License][MIT].
//!
//! # Development
//! `waitcell` is developed at [GitLab].
//!
//! [Apache2]: https://www.apache.org/licenses/LICENSE-2.0
//! [block]: std::thread::park
//! [`Future`]: std::future::Future
//! [GitLab]: https://gitlab.com/nwsharp/waitcell
//! [MIT]: https://opensource.org/licenses/MIT
//! [`Sync`]: std::marker::Sync

use std::cell::UnsafeCell;
use std::fmt::{self, Debug, Formatter};
use std::intrinsics::{ atomic_load_acq, atomic_store_rel };
use std::mem::{self, MaybeUninit};
use std::panic::{self, AssertUnwindSafe};
use std::ptr;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::thread::{self, Thread};

#[cfg(test)]
mod tests;

const STATE_CLEAR: usize = 0;
const STATE_INIT: usize = usize::MAX;
const BUSY_BIT: usize = 1;

// This is a singly linked list node used by parked threads to allow the
// initializing thread to wake them.
#[repr(align(2))]
struct Parked {
    // A bit indicating that this structure is no longer externally referenced.
    // Synchronizes with the availability of the value.
    pub wake: UnsafeCell<usize>,

    // A handle for the parked thread.
    pub thread: Thread,

    // The next parked thread, or NULL if there is none.
    pub next: *const Parked,
}

unsafe impl Sync for Parked {}

/// A cell type containing a value which may not yet be available.
///
/// A `WaitCell` generally begins in an uninitialized state. While in this
/// state, attempts to [reference] the value [block] until some other thread
/// [provides a value]. While shared, a value can only be set at most once.
///
/// If you require the value to be set multiple times while shared, consider
/// using [`RwLock`] instead. `WaitCell` is more performant than [`RwLock`],
/// however, as expensive atomic writes are not required to track readers and
/// writers.
///
/// In asynchronous contexts, a type implementing [`Future`] should be used
/// instead.
///
/// # Example
/// See the [crate-level documentation].
///
/// [block]: std::thread::park
/// [crate-level documentation]: crate#Example
/// [reference]: WaitCell::as_ref
/// [`Future`]: std::future::Future
/// [provides a value]: WaitCell::init
/// [`RwLock`]: std::sync::RwLock
pub struct WaitCell<T> {
    // Bit zero of `state` is the BUSY_BIT. If it is set, a thread is currently
    // initializing the WaitCell. The remaining bits are the high bits of a
    // possibly-null pointer to the first `Parked` thread. Since `Parked` is at
    // least align(2), we know that the bit taken by BUSY_BIT is always 0. However,
    // if all bits are set, the cell is instead initialized and `value` is valid.
    state: AtomicUsize,
    value: UnsafeCell<MaybeUninit<T>>,
}

impl<T> WaitCell<T> {
    /// Constructs an uninitialized `WaitCell` value.
    ///
    /// Attempts to dereference this value will [block] until an initialization
    /// function such as [`init`] is invoked.
    ///
    /// [block]: std::thread::park
    /// [`init`]: Self::init
    #[must_use]
    pub const fn new() -> Self {
        Self {
            state: AtomicUsize::new(STATE_CLEAR),
            value: UnsafeCell::new(MaybeUninit::uninit()),
        }
    }

    /// Constructs an initialized `WaitCell` value.
    ///
    /// Attempts to dereference this value will *not* block.
    #[must_use]
    pub const fn initialized(value: T) -> Self {
        Self {
            state: AtomicUsize::new(STATE_INIT),
            value: UnsafeCell::new(MaybeUninit::new(value)),
        }
    }

    #[must_use]
    unsafe fn get_ptr(&self) -> *mut T {
        (&mut *self.value.get()).as_mut_ptr()
    }

    #[must_use]
    unsafe fn get_ref(&self) -> &T {
        &*self.get_ptr()
    }

    #[allow(clippy::mut_from_ref)]
    #[must_use]
    unsafe fn get_mut(&self) -> &mut T {
        &mut *self.get_ptr()
    }

    unsafe fn wake(&self, value: T) -> &T {
        ptr::write(self.get_ptr(), value);

        let state = self.state.swap(STATE_INIT, Ordering::Release);
        debug_assert!((state & BUSY_BIT) != 0);

        let mut parked_ptr = (state & !BUSY_BIT) as *const Parked;
        while !parked_ptr.is_null() {
            let parked_wake = (*parked_ptr).wake.get();
            let parked_thread = (*parked_ptr).thread.clone();
            parked_ptr = (*parked_ptr).next;

            // As soon a nonzero value is stored in parked_wake, the prior value of
            // parked_ptr is potentially no longer valid to dereference.
            atomic_store_rel(parked_wake, 1);
            parked_thread.unpark();
        }

        self.get_ref()
    }

    unsafe fn wake_with<F: FnOnce() -> T>(&self, func: F) -> &T {
        match panic::catch_unwind(AssertUnwindSafe(func)) {
            Ok(value) => self.wake(value),
            Err(panic) => {
                self.state.fetch_and(!BUSY_BIT, Ordering::Relaxed);
                panic::resume_unwind(panic);
            }
        }
    }

    unsafe fn wait(&self, mut state: usize) -> &T {
        if state == STATE_INIT {
            return self.get_ref();
        }

        // Try to get in the Parked list. We may notice that we are suddenly init at
        // any time.
        let mut parked = Parked {
            wake: UnsafeCell::new(0),
            thread: thread::current(),
            next: ptr::null(),
        };

        let my_ptr = &parked as *const Parked as usize;
        debug_assert!((my_ptr & BUSY_BIT) == 0);

        loop {
            parked.next = (state & !BUSY_BIT) as *const Parked;
            match self.state.compare_exchange_weak(
                state,
                my_ptr | (state & BUSY_BIT),
                Ordering::Release,
                Ordering::Relaxed,
            ) {
                Ok(..) => {
                    break;
                }

                Err(STATE_INIT) => {
                    // Synchronize with the release operation that stored STATE_INIT.
                    state = self.state.load(Ordering::Acquire);
                    debug_assert_eq!(state, STATE_INIT);

                    return self.get_ref();
                }

                Err(current_state) => {
                    state = current_state;
                }
            };
        }

        // Park until the value is available and no reference exists to our Parked
        // entry.
        loop {
            thread::park();

            if atomic_load_acq(parked.wake.get()) != 0 {
                return self.get_ref();
            }
        }
    }

    /// Initializes the value, allowing dereferencing to proceed.
    ///
    /// # Panics
    /// If the value is already initialized or is currently undergoing
    /// initialization. To conditionally initialize the value, consider using
    /// [`try_init`].
    ///
    /// [`try_init`]: Self::try_init
    pub fn init(&self, value: T) -> &T {
        let state = self.state.fetch_or(BUSY_BIT, Ordering::Relaxed);
        if (state & BUSY_BIT) == 0 {
            unsafe { self.wake(value) }
        } else if state == STATE_INIT {
            panic!("WaitCell is already initialized");
        } else {
            panic!("WaitCell is already initializing");
        }
    }

    /// Conditionally initializes the value, allowing dereferencing to proceed.
    ///
    /// # Returns
    /// * `true` -- The value was initialized using `func`.
    /// * `false` -- The value is already initialized or is currently initializing.
    ///   `func` was not invoked.
    ///
    /// # Panics
    /// If `func` panics, the `WaitCell` remains uninitialized. Concurrent
    /// initialization attempts may fail.
    pub fn try_init<F: FnOnce() -> T>(&self, func: F) -> bool {
        let state = self.state.fetch_or(BUSY_BIT, Ordering::Relaxed);
        if (state & BUSY_BIT) == 0 {
            unsafe {
                self.wake_with(func);
            }
            true
        } else {
            false
        }
    }

    /// Conditionally initializes the value or waits for the value to become
    /// available if it is not already so.
    ///
    /// # Panics
    /// If `func` panics, the `WaitCell` remains uninitialized. Concurrent
    /// initialization attempts may fail.
    ///
    /// # Notes
    /// `func` is only invoked in the case where the value is not currently
    /// initialized or undergoing initialization.
    ///
    /// This function blocks if the value is currently undergoing initialization.
    #[must_use]
    pub fn get_or_init<F: FnOnce() -> T>(&self, func: F) -> &T {
        let state = self.state.fetch_or(BUSY_BIT, Ordering::Acquire);
        if (state & BUSY_BIT) == 0 {
            unsafe { self.wake_with(func) }
        } else {
            unsafe { self.wait(state) }
        }
    }

    /// Waits for the value to become initialized and returns a reference to it.
    ///
    /// # Notes
    /// This function will block until the value is initialized by another thread.
    #[must_use]
    pub fn get(&self) -> &T {
        unsafe { self.wait(self.state.load(Ordering::Acquire)) }
    }

    /// Returns a reference to the initialized value, or `None` if it is not yet
    /// initialized.
    #[must_use]
    pub fn try_get(&self) -> Option<&T> {
        if self.state.load(Ordering::Acquire) == STATE_INIT {
            Some(unsafe { self.get_ref() })
        } else {
            None
        }
    }

    /// Sets the initialized value, returning a mutable reference to it.
    ///
    /// The previous value, if any, is [dropped].
    ///
    /// [dropped]: std::ops::Drop::drop
    pub fn set(&mut self, value: T) -> &mut T {
        let state = *self.state.get_mut();
        if state == STATE_INIT {
            let inner = unsafe { self.get_mut() };
            *inner = value;
            inner
        } else {
            debug_assert_eq!(state, STATE_CLEAR);

            unsafe {
                ptr::write(self.get_ptr(), value);
                *self.state.get_mut() = STATE_INIT;

                self.get_mut()
            }
        }
    }

    /// Drops the initialized value, if any. The value may be re-initialized again
    /// later.
    ///
    /// # Returns
    /// * `true` - A value was present and was dropped.
    /// * `false` - A value was not present.
    pub fn unset(&mut self) -> bool {
        let state = *self.state.get_mut();
        if state == STATE_INIT {
            *self.state.get_mut() = STATE_CLEAR;

            unsafe {
                ptr::drop_in_place(self.get_ptr());
            }

            true
        } else {
            debug_assert_eq!(state, STATE_CLEAR);
            false
        }
    }

    /// Returns `true` if the value is currently initialized and `false` otherwise.
    ///
    /// # Notes
    /// This method is only available when the `WaitCell` is exclusively borrowed.
    /// If the `WaitCell` is shared, consider using
    /// `wait_cell.`[`try_get`]`().`[`is_some`]`()` instead.
    ///
    /// [`is_some`]: std::option::Option::is_some
    /// [`try_get`]: Self::try_get
    #[must_use]
    pub fn is_set(&mut self) -> bool {
        let state = *self.state.get_mut();
        debug_assert!(state == STATE_INIT || state == STATE_CLEAR);
        state == STATE_INIT
    }

    /// Returns a mutable reference to the initialized value, or `None` if the value
    /// is not initialized.
    #[must_use]
    pub fn as_inner(&mut self) -> Option<&mut T> {
        let state = *self.state.get_mut();
        if state == STATE_INIT {
            // Once unsafe_cell_get_mut is stabilized, this can be `self.value.get_mut()`.
            Some(unsafe { self.get_mut() })
        } else {
            debug_assert_eq!(state, STATE_CLEAR);
            None
        }
    }

    /// Returns the initialized value, or `None` if the value is not initialized.
    #[must_use]
    pub fn into_inner(mut self) -> Option<T> {
        let state = *self.state.get_mut();
        if state == STATE_INIT {
            unsafe {
                let result = ptr::read(self.get_ptr());
                mem::forget(self);
                Some(result)
            }
        } else {
            debug_assert_eq!(*self.state.get_mut(), STATE_CLEAR);
            None
        }
    }

    /// Provides direct immutable access to the value.
    ///
    /// # Safety
    /// The value may or may not be initialized and may or may not have concurrent
    /// immutable references. If a concurrent mutable reference exists, it was
    /// created unsafely.
    #[must_use]
    pub fn as_ptr(&self) -> *const T {
        unsafe { self.get_ptr() }
    }

    /// Provides direct mutable access to the value.
    ///
    /// # Safety
    /// The value may or may not be initialized, but it is guaranteed to have no
    /// concurrent references except those created unsafely.
    #[must_use]
    pub fn as_mut_ptr(&mut self) -> *mut T {
        unsafe { self.get_ptr() }
    }

    /// Directly sets the state of the value as initialized.
    ///
    /// # Safety
    /// If the value is later accessed, it must actually have been initialized.
    pub unsafe fn set_init(&mut self) {
        *self.state.get_mut() = STATE_INIT;
    }

    /// Directly clears the initialization state of the value, as if by
    /// [core::mem::forget].
    pub fn set_uninit(&mut self) {
        *self.state.get_mut() = STATE_CLEAR;
    }
}

impl<T: Default> WaitCell<T> {
    /// Constructs a default-initialized `WaitCell` value.
    ///
    /// Attempts to dereference this value will *not* block.
    #[must_use]
    pub fn with_default() -> Self {
        Self::initialized(Default::default())
    }
}

impl<T> Drop for WaitCell<T> {
    fn drop(&mut self) {
        let state = *self.state.get_mut();
        if state == STATE_INIT {
            unsafe {
                ptr::drop_in_place(self.get_ptr());
            }
        } else {
            debug_assert_eq!(state, STATE_CLEAR);
        }
    }
}

impl<T> AsRef<T> for WaitCell<T> {
    fn as_ref(&self) -> &T {
        unsafe { self.get_ref() }
    }
}

impl<T> Debug for WaitCell<T> {
    fn fmt(&self, f: &mut Formatter) -> fmt::Result {
        let state = self.state.load(Ordering::Relaxed);
        if state == STATE_INIT {
            f.debug_struct("WaitCell")
                .field("state", &"initialized")
                .field("has_waiter", &false)
                .finish()
        } else if (state & BUSY_BIT) == 0 {
            f.debug_struct("WaitCell")
                .field("state", &"uninitialized")
                .field("has_waiter", &((state & !BUSY_BIT) != 0))
                .finish()
        } else {
            f.debug_struct("WaitCell")
                .field("state", &"initializing")
                .field("has_waiter", &((state & !BUSY_BIT) != 0))
                .finish()
        }
    }
}

impl<T> Default for WaitCell<T> {
    fn default() -> Self {
        Self::new()
    }
}

impl<T> From<T> for WaitCell<T> {
    fn from(value: T) -> Self {
        Self::initialized(value)
    }
}

unsafe impl<T: Sync> Sync for WaitCell<T> {}