mod spin_mutex { use core::{ cell::UnsafeCell, sync::atomic::{AtomicBool, Ordering}, }; pub struct SpinMutex { lock: AtomicBool, data: UnsafeCell, } pub struct SpinMutexGuard<'a, T: ?Sized + 'a> { mutex: &'a SpinMutex, } impl core::ops::Deref for SpinMutexGuard<'_, T> { type Target = T; fn deref(&self) -> &Self::Target { unsafe { &*self.mutex.data.get() } } } impl core::ops::DerefMut for SpinMutexGuard<'_, T> { fn deref_mut(&mut self) -> &mut Self::Target { unsafe { &mut *self.mutex.data.get() } } } impl Drop for SpinMutexGuard<'_, T> { fn drop(&mut self) { unsafe { self.mutex.unlock() } } } unsafe impl Send for SpinMutex {} unsafe impl Sync for SpinMutex {} unsafe impl Send for SpinMutexGuard<'_, T> where for<'a> &'a mut T: Send {} unsafe impl Sync for SpinMutexGuard<'_, T> where for<'a> &'a mut T: Sync {} impl SpinMutex { pub const fn new(data: T) -> Self { Self { lock: AtomicBool::new(false), data: UnsafeCell::new(data), } } pub fn into_inner(self) -> T { self.data.into_inner() } pub fn as_mut_ptr(&self) -> *mut T { self.data.get() } } impl SpinMutex { pub fn try_lock(&self) -> bool { self.lock .compare_exchange(false, true, Ordering::Acquire, Ordering::Relaxed) .is_ok() } pub fn try_lock_weak(&self) -> bool { self.lock .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed) .is_ok() } pub fn lock(&self) -> SpinMutexGuard<'_, T> { while !self.try_lock() { core::hint::spin_loop(); } SpinMutexGuard { mutex: self } } /// # Safety /// The caller must be the logical owner of the locked mutex. pub unsafe fn unlock(&self) { self.lock.store(false, Ordering::Release); } } } pub use spin_mutex::SpinMutex; mod once { use core::{ cell::{Cell, UnsafeCell}, mem::{ManuallyDrop, MaybeUninit}, sync::atomic::{AtomicU8, Ordering}, }; const UNINITIALIZED: u8 = 0; const INITIALIZING: u8 = 1; const INITIALIZED: u8 = 2; const PANICKED: u8 = 3; pub struct OnceState { poisoned: bool, state_to_set: Cell, } impl OnceState { pub fn is_poisoned(&self) -> bool { self.poisoned } pub fn poison(&self) { self.state_to_set.set(PANICKED); } } pub struct Once { state: AtomicU8, } const impl Default for Once { fn default() -> Self { Self { state: AtomicU8::new(UNINITIALIZED), } } } impl Once { pub const fn new() -> Self { Self { state: AtomicU8::new(UNINITIALIZED), } } pub fn state(&self) -> u8 { self.state.load(Ordering::Acquire) } pub fn set_state(&self, state: u8) { self.state.store(state, Ordering::Release); } pub fn is_completed(&self) -> bool { self.state.load(Ordering::Acquire) == INITIALIZED } pub fn call_once(&self, f: F) where F: FnOnce(&OnceState), { if self.is_completed() { return; } self.call_once_slow(false, f); } pub fn call_once_force(&self, f: F) where F: FnOnce(&OnceState), { if self.is_completed() { return; } self.call_once_slow(true, f); } #[cold] pub fn call_once_slow(&self, ignore_poison: bool, f: F) where F: FnOnce(&OnceState), { let mut state = self.state.load(Ordering::Acquire); loop { match state { INITIALIZED => return, PANICKED if !ignore_poison => { panic!("Once instance has previously been poisoned") } PANICKED | UNINITIALIZED => { match self.state.compare_exchange( state, INITIALIZING, Ordering::Acquire, // if we get `Err(INITIALIZED)`, we want to have // acquired what the lock is protecting. Ordering::Acquire, ) { Err(new) => { state = new; continue; } Ok(_) => { // even though we don't have unwinding, for // completeness sake we'll poison the lock if // the closure panics. struct Guard<'a>(&'a AtomicU8); impl Drop for Guard<'_> { fn drop(&mut self) { self.0.store(PANICKED, Ordering::Release); } } let guard = Guard(&self.state); let state = OnceState { poisoned: state == PANICKED, state_to_set: Cell::new(INITIALIZED), }; f(&state); core::mem::forget(guard); self.state .store(state.state_to_set.take(), Ordering::Release); return; } } } _ => { assert_eq!(state, INITIALIZING); loop { state = self.state.load(Ordering::Acquire); if state != INITIALIZING { break; } core::hint::spin_loop(); } } } } } } pub struct OnceLock { once: Once, data: UnsafeCell>, } unsafe impl Send for OnceLock {} unsafe impl Sync for OnceLock {} const impl Default for OnceLock { fn default() -> Self { Self { once: Once::new(), data: UnsafeCell::new(MaybeUninit::uninit()), } } } impl OnceLock { pub const fn new() -> Self { Self { once: Once::new(), data: UnsafeCell::new(MaybeUninit::uninit()), } } pub const fn from(t: T) -> Self { Self { once: Once::new(), data: UnsafeCell::new(MaybeUninit::new(t)), } } /// # Safety /// The caller must ensure that the `OnceLock` has been initialized before calling this method. pub unsafe fn get_unchecked(&self) -> &T { unsafe { self.data.get().as_ref_unchecked().assume_init_ref() } } /// # Safety /// The caller must ensure that the `OnceLock` has been initialized before calling this method. pub unsafe fn get_mut_unchecked(&mut self) -> &mut T { unsafe { self.data.get().as_mut_unchecked().assume_init_mut() } } pub fn get(&self) -> Option<&T> { if self.is_completed() { Some(unsafe { self.get_unchecked() }) } else { None } } pub fn get_mut(&mut self) -> Option<&mut T> { if self.is_completed_mut() { Some(unsafe { self.get_mut_unchecked() }) } else { None } } pub fn is_completed(&self) -> bool { self.once.is_completed() } pub fn is_completed_mut(&mut self) -> bool { // avoid atomic load when we have mutable access to self *self.once.state.get_mut() == INITIALIZED } pub fn get_spinning(&self) -> Option<&T> { loop { match self.once.state() { INITIALIZED => return Some(unsafe { self.get_unchecked() }), INITIALIZING => core::hint::spin_loop(), UNINITIALIZED | PANICKED => return None, _ => unreachable!(), } } } pub fn initialize(&self, f: F) -> Result<(), E> where F: FnOnce() -> Result, { let mut res: Result<(), E> = Ok(()); let cell = unsafe { &mut *self.data.get() }; self.once.call_once_force(|state| match f() { Ok(value) => { unsafe { cell.as_mut_ptr().write(value) }; } Err(e) => { state.poison(); res = Err(e); } }); res } pub fn get_or_init(&self, f: F) -> &T where F: FnOnce() -> T, { match self.get_or_try_init(|| Ok::(f())) { Ok(value) => value, Err(_) => unreachable!(), } } pub fn get_mut_or_init(&mut self, f: F) -> &mut T where F: FnOnce() -> T, { match self.get_mut_or_try_init(|| Ok::(f())) { Ok(value) => value, Err(_) => unreachable!(), } } pub fn get_or_try_init(&self, f: F) -> Result<&T, E> where F: FnOnce() -> Result, { if !self.is_completed() { self.initialize(f)?; } unsafe { Ok(self.get_unchecked()) } } pub fn get_mut_or_try_init(&mut self, f: F) -> Result<&mut T, E> where F: FnOnce() -> Result, { if self.get_mut().is_none() { self.initialize(f)?; } Ok(unsafe { self.get_mut_unchecked() }) } pub fn try_insert(&self, value: T) -> Result<&T, (&T, T)> { let mut val = Some(value); let res = self.get_or_init(|| unsafe { val.take().unwrap_unchecked() }); match val { Some(value) => Err((res, value)), None => Ok(res), } } } impl Drop for OnceLock { fn drop(&mut self) { if self.is_completed() { unsafe { self.data.get().as_mut_unchecked().assume_init_drop() } } } } pub struct LazyLock T> { once: Once, f: UnsafeCell>, t: UnsafeCell>, } unsafe impl Sync for LazyLock {} impl LazyLock where F: FnOnce() -> T, { pub const fn new(f: F) -> Self { Self { once: Once::new(), f: UnsafeCell::new(ManuallyDrop::new(f)), t: UnsafeCell::new(MaybeUninit::uninit()), } } pub fn get(&self) -> Option<&T> { if self.once.is_completed() { Some(unsafe { (&*self.t.get()).assume_init_ref() }) } else { None } } fn force(this: &Self) -> &T { this.once.call_once(|_| { // SAFETY: because `call_once` will panic if poisoned, this // closure will only be called once. let f = unsafe { ManuallyDrop::take(&mut *this.f.get()) }; let val = f(); unsafe { (&mut *this.t.get()).write(val) }; }); unsafe { (&*this.t.get()).assume_init_ref() } } fn force_mut(this: &mut Self) -> &mut T { this.once.call_once(|_| { // SAFETY: because `call_once` will panic if poisoned, this // closure will only be called once. let f = unsafe { ManuallyDrop::take(&mut *this.f.get()) }; let val = f(); unsafe { (&mut *this.t.get()).write(val) }; }); unsafe { (&mut *this.t.get()).assume_init_mut() } } } impl T> core::ops::Deref for LazyLock { type Target = T; fn deref(&self) -> &Self::Target { LazyLock::force(self) } } impl T> core::ops::DerefMut for LazyLock { fn deref_mut(&mut self) -> &mut Self::Target { LazyLock::force_mut(self) } } impl Drop for LazyLock { fn drop(&mut self) { match self.once.state() { UNINITIALIZED => unsafe { ManuallyDrop::drop(self.f.get_mut()); }, INITIALIZED => unsafe { (&mut *self.t.get()).assume_init_drop(); }, _ => {} } } } } pub use once::{LazyLock, Once, OnceLock};