src/spin/once.rs

use core::cell::UnsafeCell;
use core::fmt;
use core::hint::spin_loop;
use core::sync::atomic::{AtomicUsize, Ordering};

/// A synchronization primitive which can be used to run a one-time global
/// initialization. Unlike its std equivalent, this is generalized so that the
/// closure returns a value and it is stored. Once therefore acts something like
/// a future, too.
pub struct Once<T> {
    state: AtomicUsize,
    data: UnsafeCell<Option<T>>, // TODO remove option and use mem::uninitialized
}

impl<T: fmt::Debug> fmt::Debug for Once<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self.r#try() {
            Some(s) => write!(f, "Once {{ data: ")
                .and_then(|()| s.fmt(f))
                .and_then(|()| write!(f, "}}")),
            None => write!(f, "Once {{ <uninitialized> }}"),
        }
    }
}

// Same unsafe impls as `std::sync::RwLock`, because this also allows for
// concurrent reads.
unsafe impl<T: Send + Sync> Sync for Once<T> {}
unsafe impl<T: Send> Send for Once<T> {}

// Four states that a Once can be in, encoded into the lower bits of `state` in
// the Once structure.
const INCOMPLETE: usize = 0x0;
const RUNNING: usize = 0x1;
const COMPLETE: usize = 0x2;
const PANICKED: usize = 0x3;

use core::hint::unreachable_unchecked as unreachable;

impl<T> Once<T> {
    /// Initialization constant of `Once`.
    pub const INIT: Self = Once {
        state: AtomicUsize::new(INCOMPLETE),
        data: UnsafeCell::new(None),
    };

    /// Creates a new `Once` value.
    pub const fn new() -> Once<T> {
        Self::INIT
    }

    fn force_get<'a>(&'a self) -> &'a T {
        match unsafe { &*self.data.get() }.as_ref() {
            None => unsafe { unreachable() },
            Some(p) => p,
        }
    }

    /// Performs an initialization routine once and only once. The given closure
    /// will be executed if this is the first time `call_once` has been called,
    /// and otherwise the routine will *not* be invoked.
    ///
    /// This method will block the calling thread if another initialization
    /// routine is currently running.
    ///
    /// When this function returns, it is guaranteed that some initialization
    /// has run and completed (it may not be the closure specified). The
    /// returned pointer will point to the result from the closure that was
    /// run.
    pub fn call_once<'a, F>(&'a self, builder: F) -> &'a T
    where
        F: FnOnce() -> T,
    {
        let mut status = self.state.load(Ordering::SeqCst);

        if status == INCOMPLETE {
            status = match self.state.compare_exchange(
                INCOMPLETE,
                RUNNING,
                Ordering::SeqCst,
                Ordering::SeqCst,
            ) {
                Ok(status) => {
                    debug_assert_eq!(
                        status, INCOMPLETE,
                        "if compare_exchange succeeded, previous status must be incomplete",
                    );
                    // We init
                    // We use a guard (Finish) to catch panics caused by builder
                    let mut finish = Finish {
                        state: &self.state,
                        panicked: true,
                    };
                    unsafe { *self.data.get() = Some(builder()) };
                    finish.panicked = false;

                    self.state.store(COMPLETE, Ordering::SeqCst);

                    // This next line is strictly an optimization
                    return self.force_get();
                }
                Err(status) => status,
            }
        }

        loop {
            match status {
                INCOMPLETE => unreachable!(),
                RUNNING => {
                    // We spin
                    spin_loop();
                    status = self.state.load(Ordering::SeqCst)
                }
                PANICKED => panic!("Once has panicked"),
                COMPLETE => return self.force_get(),
                _ => unsafe { unreachable() },
            }
        }
    }

    /// Returns a pointer iff the `Once` was previously initialized
    pub fn r#try<'a>(&'a self) -> Option<&'a T> {
        match self.state.load(Ordering::SeqCst) {
            COMPLETE => Some(self.force_get()),
            _ => None,
        }
    }

    /// Like try, but will spin if the `Once` is in the process of being
    /// initialized
    pub fn wait<'a>(&'a self) -> Option<&'a T> {
        loop {
            match self.state.load(Ordering::SeqCst) {
                INCOMPLETE => return None,

                RUNNING => {
                    spin_loop() // We spin
                }
                COMPLETE => return Some(self.force_get()),
                PANICKED => panic!("Once has panicked"),
                _ => unsafe { unreachable() },
            }
        }
    }
}

struct Finish<'a> {
    state: &'a AtomicUsize,
    panicked: bool,
}

impl<'a> Drop for Finish<'a> {
    fn drop(&mut self) {
        if self.panicked {
            self.state.store(PANICKED, Ordering::SeqCst);
        }
    }
}