use std::{ cell::UnsafeCell, hint, panic::{RefUnwindSafe, UnwindSafe}, sync::atomic::{AtomicBool, Ordering}, }; use parking_lot::Mutex; pub(crate) struct OnceCell { mutex: Mutex<()>, is_initialized: AtomicBool, value: UnsafeCell>, } // Why do we need `T: Send`? // Thread A creates a `OnceCell` and shares it with // scoped thread B, which fills the cell, which is // then destroyed by A. That is, destructor observes // a sent value. unsafe impl Sync for OnceCell {} unsafe impl Send for OnceCell {} impl RefUnwindSafe for OnceCell {} impl UnwindSafe for OnceCell {} impl OnceCell { pub(crate) const fn new() -> OnceCell { OnceCell { mutex: parking_lot::const_mutex(()), is_initialized: AtomicBool::new(false), value: UnsafeCell::new(None), } } /// Safety: synchronizes with store to value via Release/Acquire. #[inline] pub(crate) fn is_initialized(&self) -> bool { self.is_initialized.load(Ordering::Acquire) } /// Safety: synchronizes with store to value via `is_initialized` or mutex /// lock/unlock, writes value only once because of the mutex. #[cold] pub(crate) fn initialize(&self, f: F) -> Result<(), E> where F: FnOnce() -> Result, { let _guard = self.mutex.lock(); if !self.is_initialized() { // We are calling user-supplied function and need to be careful. // - if it returns Err, we unlock mutex and return without touching anything // - if it panics, we unlock mutex and propagate panic without touching anything // - if it calls `set` or `get_or_try_init` re-entrantly, we get a deadlock on // mutex, which is important for safety. We *could* detect this and panic, // but that is more complicated // - finally, if it returns Ok, we store the value and store the flag with // `Release`, which synchronizes with `Acquire`s. let value = f()?; // Safe b/c we have a unique access and no panic may happen // until the cell is marked as initialized. let slot: &mut Option = unsafe { &mut *self.value.get() }; debug_assert!(slot.is_none()); *slot = Some(value); self.is_initialized.store(true, Ordering::Release); } Ok(()) } /// Get the reference to the underlying value, without checking if the cell /// is initialized. /// /// # Safety /// /// Caller must ensure that the cell is in initialized state, and that /// the contents are acquired by (synchronized to) this thread. pub(crate) unsafe fn get_unchecked(&self) -> &T { debug_assert!(self.is_initialized()); let slot: &Option = &*self.value.get(); match slot { Some(value) => value, // This unsafe does improve performance, see `examples/bench`. None => { debug_assert!(false); hint::unreachable_unchecked() } } } /// Gets the mutable reference to the underlying value. /// Returns `None` if the cell is empty. pub(crate) fn get_mut(&mut self) -> Option<&mut T> { // Safe b/c we have an exclusive access let slot: &mut Option = unsafe { &mut *self.value.get() }; slot.as_mut() } /// Consumes this `OnceCell`, returning the wrapped value. /// Returns `None` if the cell was empty. pub(crate) fn into_inner(self) -> Option { self.value.into_inner() } } #[test] fn test_size() { use std::mem::size_of; assert_eq!(size_of::>(), 2 * size_of::() + size_of::()); }