use super::mutual::Mutual;
use super::{InitState, Waiter, Waiters, WaitersExt};
use alloc::sync::Arc;
use core::cell::UnsafeCell;
use core::future::Future;
use core::mem::MaybeUninit;
use core::sync::atomic::Ordering;
pub struct Lazy<T, F> {
state: Mutual<InitState>,
value: UnsafeCell<MaybeUninit<T>>,
init_fn: UnsafeCell<Option<F>>,
waiters: Waiters,
}
impl<T, F> Lazy<T, F> {
pub const fn new(init: F) -> Self {
Self {
state: Mutual::new(),
value: UnsafeCell::new(MaybeUninit::uninit()),
init_fn: UnsafeCell::new(Some(init)),
waiters: Waiters::new(),
}
}
}
impl<T: Send + Sync, F: Future<Output = T>> Lazy<T, F> {
pub async fn get(&self) -> &T {
loop {
if self.state.is(&InitState::Initialized) {
return unsafe { &*(*self.value.get()).as_ptr() };
}
if self.state.is(&InitState::Initializing) {
let waiter = Waiter::default();
unsafe { self.waiters.enqueue(waiter) };
if self.state.is(&InitState::Initialized) {
return unsafe { &*(*self.value.get()).as_ptr() };
}
continue;
}
match self.state.compare_exchange(
InitState::Uninitialized,
InitState::Initializing,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
let init_fn = unsafe {
(*self.init_fn.get())
.take()
.expect("init function already taken")
};
let value = init_fn.await;
unsafe {
(*self.value.get()).write(value);
}
self.state.set(InitState::Initialized);
self.waiters.notify_all();
return unsafe { &*(*self.value.get()).as_ptr() };
}
Err(_) => {
continue;
}
}
}
}
pub fn try_get(&self) -> Option<&T> {
if self.state.is(&InitState::Initialized) {
unsafe { Some(&*(*self.value.get()).as_ptr()) }
} else {
None
}
}
}
unsafe impl<T: Send + Sync, F: Send> Send for Lazy<T, F> {}
unsafe impl<T: Send + Sync, F: Send> Sync for Lazy<T, F> {}
pub struct Once<T> {
state: Mutual<InitState>,
value: UnsafeCell<MaybeUninit<T>>,
waiters: Waiters,
}
impl<T> Once<T> {
pub const fn new() -> Self {
Self {
state: Mutual::new(),
value: UnsafeCell::new(MaybeUninit::uninit()),
waiters: Waiters::new(),
}
}
pub fn init(&self, value: T) -> Result<(), T> {
match self.state.compare_exchange(
InitState::Uninitialized,
InitState::Initialized,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
unsafe {
(*self.value.get()).write(value);
}
self.waiters.notify_all();
Ok(())
}
Err(_) => Err(value),
}
}
pub async fn get(&self) -> &T {
loop {
if self.state.is(&InitState::Initialized) {
return unsafe { &*(*self.value.get()).as_ptr() };
}
let waiter = Waiter::default();
unsafe { self.waiters.enqueue(waiter) };
}
}
pub fn try_get(&self) -> Option<&T> {
if self.state.is(&InitState::Initialized) {
unsafe { Some(&*(*self.value.get()).as_ptr()) }
} else {
None
}
}
}
unsafe impl<T: Send + Sync> Send for Once<T> {}
unsafe impl<T: Send + Sync> Sync for Once<T> {}
pub type SharedLazy<T, F> = Arc<Lazy<T, F>>;