waitcell 0.4.1

//
// Copyright (C) 2022 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::mem::{self, MaybeUninit};
use std::ptr;
use std::sync::atomic::{self, AtomicBool, AtomicUsize, Ordering};
use std::thread::{self, Thread};

use scopeguard::defer_on_unwind;

#[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.
    wake: AtomicBool,

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

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

impl Parked {
    // Safety: 'state' must be a valid state.
    pub unsafe fn park(state: &AtomicUsize, mut guess: usize) {
        let mut parked = Self {
            wake: AtomicBool::new(false),
            thread: thread::current(),
            next: ptr::null(),
        };

        let parked_ptr = &parked as *const Self as usize;
        debug_assert!((parked_ptr & BUSY_BIT) == 0);

        // Try to get in the Parked list. We may notice that we are suddenly STATE_INIT
        // at any time.
        loop {
            parked.next = (guess & !BUSY_BIT) as *const Parked;
            match state.compare_exchange_weak(
                guess,
                parked_ptr | (guess & BUSY_BIT),
                Ordering::Release,
                Ordering::Relaxed,
            ) {
                Ok(..) => {
                    break;
                }

                Err(STATE_INIT) => {
                    // Synchronize with the release operation that stored STATE_INIT.
                    atomic::fence(Ordering::Acquire);
                    return;
                }

                Err(current) => {
                    guess = current;
                }
            };
        }

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

            if parked.wake.load(Ordering::Acquire) {
                return;
            }
        }
    }

    pub unsafe fn unpark_all(mut head: *const Self) {
        while !head.is_null() {
            head = Self::unpark(head);
        }
    }

    unsafe fn unpark(parked: *const Self) -> *const Self {
        let thread = (*parked).thread.clone();
        let next = (*parked).next;

        assert!(!(*parked).wake.swap(true, Ordering::Release));
        thread.unpark();

        next
    }
}

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::get
/// [`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);

        Parked::unpark_all((state & !BUSY_BIT) as *const Parked);
        self.get_ref()
    }

    unsafe fn wake_with<F: FnOnce() -> T>(&self, func: F) -> &T {
        self.wake({
            defer_on_unwind! {
                self.state.fetch_and(!BUSY_BIT, Ordering::Relaxed);
            }

            func()
        })
    }

    unsafe fn wait(&self, state: usize) -> &T {
        if state != STATE_INIT {
            Parked::park(&self.state, state);
        }

        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);
        assert!((state & BUSY_BIT) == 0, "WaitCell is not uninitialized");
        unsafe { self.wake(value) }
    }

    /// 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 {
            return false;
        }

        unsafe {
            self.wake_with(func);
        }

        true
    }

    /// 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> {
        (self.state.load(Ordering::Acquire) == STATE_INIT).then(|| unsafe { self.get_ref() })
    }

    /// 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 {
        unsafe {
            if self.is_set() {
                *self.get_mut() = value;
            } else {
                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 {
        self.is_set()
            .then(|| unsafe {
                *self.state.get_mut() = STATE_CLEAR;
                ptr::drop_in_place(self.get_ptr());
            })
            .is_some()
    }

    /// 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_CLEAR || state == STATE_INIT);
        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> {
        self.is_set().then(move || unsafe { self.get_mut() })
    }

    /// Returns the initialized value, or `None` if the value is not initialized.
    #[must_use]
    pub fn into_inner(mut self) -> Option<T> {
        self.is_set().then(|| unsafe {
            let result = ptr::read(self.get_ptr());
            mem::forget(self);
            result
        })
    }

    /// 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 *self.state.get_mut() == STATE_INIT {
            unsafe { ptr::drop_in_place(self.get_ptr()); }
        } else {
            debug_assert_eq!(state & !BUSY_BIT, STATE_CLEAR);
        }
    }
}

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> {}