use core::cell::UnsafeCell;
use core::mem::MaybeUninit;
use core::sync::atomic::{AtomicBool, Ordering};
use std::sync::Once;
pub(crate) struct OnceLock<T> {
once: Once,
is_initialized: AtomicBool,
value: UnsafeCell<MaybeUninit<T>>,
}
unsafe impl<T: Sync + Send> Sync for OnceLock<T> {}
unsafe impl<T: Send> Send for OnceLock<T> {}
impl<T> OnceLock<T> {
#[must_use]
pub(crate) const fn new() -> Self {
Self {
once: Once::new(),
is_initialized: AtomicBool::new(false),
value: UnsafeCell::new(MaybeUninit::uninit()),
}
}
pub(crate) fn get_or_init<F>(&self, f: F) -> &T
where
F: FnOnce() -> T,
{
if self.is_initialized() {
return unsafe { self.get_unchecked() };
}
self.initialize(f);
debug_assert!(self.is_initialized());
unsafe { self.get_unchecked() }
}
#[inline]
fn is_initialized(&self) -> bool {
self.is_initialized.load(Ordering::Acquire)
}
#[cold]
fn initialize<F>(&self, f: F)
where
F: FnOnce() -> T,
{
let slot = self.value.get().cast::<T>();
let is_initialized = &self.is_initialized;
self.once.call_once(|| {
let value = f();
unsafe {
slot.write(value);
}
is_initialized.store(true, Ordering::Release);
});
}
unsafe fn get_unchecked(&self) -> &T {
debug_assert!(self.is_initialized());
&*self.value.get().cast::<T>()
}
}
impl<T> Drop for OnceLock<T> {
fn drop(&mut self) {
if self.is_initialized() {
unsafe { self.value.get().cast::<T>().drop_in_place() };
}
}
}