use core::cell::UnsafeCell;
use core::future::Future;
use core::sync::atomic::Ordering;
use core::{fmt, mem};
use super::state::OnceLock;
use crate::state::OnceGuard;
union TOnceState<P, T> {
pending: mem::ManuallyDrop<P>,
init: mem::ManuallyDrop<T>,
}
impl<P, T> TOnceState<P, T> {
#[inline(always)]
const fn pending(value: P) -> Self {
Self {
pending: mem::ManuallyDrop::new(value),
}
}
#[inline(always)]
const fn init(value: T) -> Self {
Self {
init: mem::ManuallyDrop::new(value),
}
}
#[inline(always)]
unsafe fn drop_pending(&mut self) -> P {
mem::ManuallyDrop::take(unsafe { &mut self.pending })
}
#[inline(always)]
unsafe fn drop_init(&mut self) -> T {
mem::ManuallyDrop::take(unsafe { &mut self.init })
}
}
struct EditSanitizer<'a, P, T> {
state: &'a mut TOnceState<P, T>,
guard: OnceGuard<'a>,
taken: bool,
}
impl<'a, P, T> EditSanitizer<'a, P, T> {
#[inline(always)]
fn new(state: &'a UnsafeCell<TOnceState<P, T>>, guard: OnceGuard<'a>) -> Self {
Self {
state: unsafe { &mut *state.get() },
taken: false,
guard,
}
}
#[inline(always)]
unsafe fn take_parameter(&mut self) -> P {
debug_assert!(!self.taken, "Parameter already taken");
self.taken = true;
unsafe { self.state.drop_pending() }
}
#[inline(always)]
unsafe fn ref_parameter(&self) -> &P {
unsafe { &self.state.pending }
}
#[inline(always)]
unsafe fn commit(self, value: T) -> &'a mut T {
unsafe {
if !self.taken {
_ = self.state.drop_pending();
}
*self.state = TOnceState::init(value);
self.guard.commit_forgotten();
let ptr = self.state as *mut TOnceState<P, T>;
mem::forget(self);
&mut (*ptr).init
}
}
}
#[cfg(debug_assertions)]
impl<P, T> Drop for EditSanitizer<'_, P, T> {
#[inline(always)]
fn drop(&mut self) {
assert!(!self.taken, "Parameter taken on return, no value in cell");
}
}
pub struct TOnce<P, T> {
value: UnsafeCell<TOnceState<P, T>>,
lock: OnceLock,
}
impl<P, T> TOnce<P, T> {
#[inline]
#[must_use]
pub const fn new(param: P) -> Self {
Self {
lock: OnceLock::new(),
value: UnsafeCell::new(TOnceState::pending(param)),
}
}
#[inline]
#[must_use]
pub const fn with_value(value: T) -> Self {
Self {
lock: OnceLock::done(),
value: UnsafeCell::new(TOnceState::init(value)),
}
}
#[inline]
pub fn is_done(&self) -> bool {
self.lock.is_done(Ordering::Relaxed)
}
#[inline]
pub fn get(&self) -> Option<&T> {
if self.is_done() {
Some(unsafe { self.get_unchecked() })
} else {
None
}
}
#[inline]
pub fn get_mut(&mut self) -> Option<&mut T> {
if self.is_done() {
Some(unsafe { self.get_unchecked_mut() })
} else {
None
}
}
#[inline]
pub fn try_set(&self, value: T) -> Result<&T, T> {
let Some(guard) = self.lock.try_lock() else {
return Err(value);
};
unsafe {
let edit = EditSanitizer::new(&self.value, guard);
let v = edit.commit(value);
Ok(v)
}
}
#[inline]
pub fn replace_mut(&mut self, value: T) -> Option<T> {
if self.lock.set_done() {
unsafe {
mem::replace(self.value.get_mut(), TOnceState::init(value)).drop_pending();
}
None
} else {
Some(unsafe { mem::replace(self.get_unchecked_mut(), value) })
}
}
#[inline]
pub fn get_mut_or_set(&mut self, value: T) -> &mut T {
if self.lock.set_done() {
unsafe {
mem::replace(self.value.get_mut(), TOnceState::init(value)).drop_pending();
}
}
unsafe { self.get_unchecked_mut() }
}
#[inline]
pub fn get_mut_or_default(&mut self) -> &mut T
where
T: Default,
{
if self.lock.set_done() {
unsafe {
mem::replace(self.value.get_mut(), TOnceState::init(T::default())).drop_pending();
}
}
unsafe { self.get_unchecked_mut() }
}
#[inline]
pub unsafe fn get_unchecked(&self) -> &T {
debug_assert!(self.is_done(), "get_unchecked called on uninitialized Once");
unsafe { &(*self.value.get()).init }
}
#[inline]
pub unsafe fn get_unchecked_mut(&mut self) -> &mut T {
debug_assert!(
self.is_done(),
"get_unchecked_mut called on uninitialized Once"
);
unsafe { &mut (*self.value.get()).init }
}
#[inline]
pub fn take(&mut self, param: P) -> Result<T, P> {
if !self.lock.set_uninit() {
return Err(param);
}
let mut mref = mem::replace(self.value.get_mut(), TOnceState::pending(param));
Ok(unsafe { mref.drop_init() })
}
#[inline]
pub fn set(&self, value: T) -> Result<(), T> {
match self.try_insert(value) {
Ok(_) => Ok(()), Err((_current_value, original_value)) => Err(original_value), }
}
#[inline]
pub fn try_insert(&self, value: T) -> Result<&T, (&T, T)> {
let mut value_opt = Some(value);
let res_ref = self.get_or_init(|_p| value_opt.take().unwrap());
match value_opt {
None => Ok(res_ref), Some(original_value) => Err((res_ref, original_value)), }
}
#[inline]
pub fn get_or_init<F>(&self, f: F) -> &T
where
F: FnOnce(P) -> T,
{
if let Some(value) = self.get() {
return value;
}
self.initialize(f);
unsafe { self.get_unchecked() }
}
#[inline]
pub fn get_mut_or_init<F>(&mut self, f: F) -> &mut T
where
F: FnOnce(P) -> T,
{
if self.is_done() {
} else if let Some(guard) = self.lock.try_lock() {
unsafe {
let mut edit = EditSanitizer::new(&self.value, guard);
let value = f(edit.take_parameter());
edit.commit(value);
}
} else {
unreachable!("Could not lock for init despite having exclusive access");
}
unsafe { self.get_unchecked_mut() }
}
pub fn get_or_try_init<F, E>(&self, f: F) -> Result<&T, E>
where
F: FnOnce(&P) -> Result<T, E>,
{
if let Some(value) = self.get() {
return Ok(value);
}
self.try_initialize(f)?;
debug_assert!(self.is_done());
Ok(unsafe { self.get_unchecked() })
}
pub fn get_mut_or_try_init<F, E>(&mut self, f: F) -> Result<&mut T, E>
where
F: FnOnce(&P) -> Result<T, E>,
{
if let Some(guard) = self.lock.try_lock() {
unsafe {
let edit = EditSanitizer::new(&self.value, guard);
let value = f(edit.ref_parameter())?;
edit.commit(value);
}
}
Ok(unsafe { self.get_unchecked_mut() })
}
#[inline]
pub async fn get_or_init_async<F, Fut>(&self, f: F) -> &T
where
F: FnOnce(&P) -> Fut,
Fut: Future<Output = T>,
{
if let Some(value) = self.get() {
return value;
}
self.initialize_async(f).await;
unsafe { self.get_unchecked() }
}
#[inline]
pub async fn get_mut_or_init_async<F, Fut>(&mut self, f: F) -> &mut T
where
F: FnOnce(&P) -> Fut,
Fut: Future<Output = T>,
{
if let Some(guard) = self.lock.try_lock() {
unsafe {
let edit = EditSanitizer::new(&self.value, guard);
let value = f(edit.ref_parameter()).await;
edit.commit(value);
}
}
unsafe { self.get_unchecked_mut() }
}
pub async fn get_or_try_init_async<F, Fut, E>(&self, f: F) -> Result<&T, E>
where
F: FnOnce(&P) -> Fut,
Fut: Future<Output = Result<T, E>>,
{
if let Some(value) = self.get() {
return Ok(value);
}
self.try_initialize_async(f).await?;
debug_assert!(self.is_done());
Ok(unsafe { self.get_unchecked() })
}
pub async fn get_mut_or_try_init_async<F, Fut, E>(&mut self, f: F) -> Result<&mut T, E>
where
F: FnOnce(&P) -> Fut,
Fut: Future<Output = Result<T, E>>,
{
if let Some(guard) = self.lock.try_lock() {
unsafe {
let edit = EditSanitizer::new(&self.value, guard);
let value = f(edit.ref_parameter()).await?;
edit.commit(value);
}
}
Ok(unsafe { self.get_unchecked_mut() })
}
#[cold]
async fn initialize_async<F, Fut>(&self, f: F)
where
F: FnOnce(&P) -> Fut,
Fut: Future<Output = T>,
{
let Some(guard) = self.lock.lock_async().await else {
return; };
unsafe {
let edit = EditSanitizer::new(&self.value, guard);
let value = f(edit.ref_parameter()).await;
edit.commit(value);
}
}
#[cold]
async fn try_initialize_async<Fn, Fut, E>(&self, f: Fn) -> Result<(), E>
where
Fn: FnOnce(&P) -> Fut,
Fut: Future<Output = Result<T, E>>,
{
let Some(guard) = self.lock.lock_async().await else {
return Ok(()); };
unsafe {
let edit = EditSanitizer::new(&self.value, guard);
let value = f(edit.ref_parameter()).await?;
edit.commit(value);
}
Ok(())
}
#[cold]
fn initialize<F>(&self, f: F)
where
F: FnOnce(P) -> T,
{
let Some(guard) = self.lock.lock() else {
return; };
unsafe {
let mut edit = EditSanitizer::new(&self.value, guard);
let value = f(edit.take_parameter());
edit.commit(value);
}
}
#[cold]
fn try_initialize<F, E>(&self, f: F) -> Result<(), E>
where
F: FnOnce(&P) -> Result<T, E>,
{
let Some(guard) = self.lock.lock() else {
return Ok(()); };
unsafe {
let edit = EditSanitizer::new(&self.value, guard);
let value = f(edit.ref_parameter())?;
edit.commit(value);
}
Ok(())
}
}
impl<P, T> From<Option<T>> for TOnce<P, T>
where
P: Default,
{
fn from(value: Option<T>) -> Self {
match value {
Some(value) => Self::with_value(value),
None => Self::new(P::default()),
}
}
}
impl<P, T> From<Result<T, P>> for TOnce<P, T> {
fn from(value: Result<T, P>) -> Self {
match value {
Ok(value) => Self::with_value(value),
Err(param) => Self::new(param),
}
}
}
unsafe impl<P: Sync + Send, T: Send> Sync for TOnce<P, T> {}
unsafe impl<P: Send, T: Send> Send for TOnce<P, T> {}
impl<P: Default, T> Default for TOnce<P, T> {
#[inline]
fn default() -> Self {
Self::new(P::default())
}
}
impl<P, T: fmt::Display> fmt::Display for TOnce<P, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self.get() {
Some(v) => fmt::Display::fmt(v, f),
None => f.write_str("<uninit>"),
}
}
}
impl<P, T: fmt::Debug> fmt::Debug for TOnce<P, T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let mut d = f.debug_tuple("Once");
match self.get() {
Some(v) => d.field(v),
None => d.field(&format_args!("<uninit>")),
};
d.finish()
}
}
impl<P: Clone, T: Clone> Clone for TOnce<P, T> {
#[inline]
fn clone(&self) -> Self {
if let Some(value) = self.get() {
Self::with_value(value.clone())
} else {
let Some(_guard) = self.lock.lock() else {
return Self::with_value(
self
.get()
.expect("is_done() was false but value is init")
.clone(),
);
};
Self::new(unsafe { (*(*self.value.get()).pending).clone() })
}
}
}
impl<T> From<T> for TOnce<(), T> {
#[inline]
fn from(value: T) -> Self {
Self::with_value(value)
}
}
impl<P, T: PartialEq> PartialEq for TOnce<P, T> {
#[inline]
fn eq(&self, other: &Self) -> bool {
self.get() == other.get()
}
}
impl<P, T: Eq> Eq for TOnce<P, T> {}
impl<P, T> Drop for TOnce<P, T> {
#[inline]
fn drop(&mut self) {
unsafe {
if self.is_done() {
self.value.get_mut().drop_init();
} else {
self.value.get_mut().drop_pending();
}
}
}
}