use core::fmt;
pub struct Lazy<T, F = fn() -> T>(lazy::Lazy<T, F>);
impl<T, F> Lazy<T, F> {
pub const fn new(create: F) -> Lazy<T, F> {
Lazy(lazy::Lazy::new(create))
}
}
impl<T, F: Fn() -> T> Lazy<T, F> {
pub fn get(this: &Lazy<T, F>) -> &T {
this.0.get()
}
}
impl<T, F: Fn() -> T> core::ops::Deref for Lazy<T, F> {
type Target = T;
fn deref(&self) -> &T {
Lazy::get(self)
}
}
impl<T: fmt::Debug, F: Fn() -> T> fmt::Debug for Lazy<T, F> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
self.0.fmt(f)
}
}
#[cfg(feature = "alloc")]
mod lazy {
use core::{
fmt,
marker::PhantomData,
sync::atomic::{AtomicPtr, Ordering},
};
use alloc::boxed::Box;
pub(super) struct Lazy<T, F> {
data: AtomicPtr<T>,
create: F,
owned: PhantomData<Box<T>>,
}
unsafe impl<T: Send + Sync, F: Send + Sync> Sync for Lazy<T, F> {}
impl<T, F> Lazy<T, F> {
pub(super) const fn new(create: F) -> Lazy<T, F> {
Lazy {
data: AtomicPtr::new(core::ptr::null_mut()),
create,
owned: PhantomData,
}
}
}
impl<T, F: Fn() -> T> Lazy<T, F> {
pub(super) fn get(&self) -> &T {
if let Some(data) = self.poll() {
return data;
}
let data = (self.create)();
let mut ptr = Box::into_raw(Box::new(data));
let result = self.data.compare_exchange(
core::ptr::null_mut(),
ptr,
Ordering::AcqRel,
Ordering::Acquire,
);
if let Err(old) = result {
drop(unsafe { Box::from_raw(ptr) });
ptr = old;
}
unsafe { &*ptr }
}
fn poll(&self) -> Option<&T> {
let ptr = self.data.load(Ordering::Acquire);
if ptr.is_null() {
return None;
}
Some(unsafe { &*ptr })
}
}
impl<T: fmt::Debug, F: Fn() -> T> fmt::Debug for Lazy<T, F> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("Lazy").field("data", &self.poll()).finish()
}
}
impl<T, F> Drop for Lazy<T, F> {
fn drop(&mut self) {
let ptr = *self.data.get_mut();
if !ptr.is_null() {
drop(unsafe { Box::from_raw(ptr) });
}
}
}
}
#[cfg(not(feature = "alloc"))]
mod lazy {
use core::{
cell::Cell,
fmt,
mem::MaybeUninit,
panic::{RefUnwindSafe, UnwindSafe},
sync::atomic::{AtomicU8, Ordering},
};
const LAZY_STATE_INIT: u8 = 0;
const LAZY_STATE_BUSY: u8 = 1;
const LAZY_STATE_DONE: u8 = 2;
pub(super) struct Lazy<T, F> {
state: AtomicU8,
create: Cell<Option<F>>,
data: Cell<MaybeUninit<T>>,
}
unsafe impl<T: Send + Sync, F: Send + Sync> Sync for Lazy<T, F> {}
impl<T: UnwindSafe, F: UnwindSafe + RefUnwindSafe> RefUnwindSafe
for Lazy<T, F>
{
}
impl<T, F> Lazy<T, F> {
pub(super) const fn new(create: F) -> Lazy<T, F> {
Lazy {
state: AtomicU8::new(LAZY_STATE_INIT),
create: Cell::new(Some(create)),
data: Cell::new(MaybeUninit::uninit()),
}
}
}
impl<T, F: FnOnce() -> T> Lazy<T, F> {
pub(super) fn get(&self) -> &T {
while self.state.load(Ordering::Acquire) != LAZY_STATE_DONE {
let result = self.state.compare_exchange(
LAZY_STATE_INIT,
LAZY_STATE_BUSY,
Ordering::AcqRel,
Ordering::Acquire,
);
if let Ok(_) = result {
let create = unsafe {
(*self.create.as_ptr()).take().expect(
"Lazy's create function panicked, \
preventing initialization,
poisoning current thread",
)
};
let guard = Guard { state: &self.state };
unsafe {
(*self.data.as_ptr()).as_mut_ptr().write(create());
}
core::mem::forget(guard);
self.state.store(LAZY_STATE_DONE, Ordering::Release);
break;
}
core::hint::spin_loop();
}
self.poll().unwrap()
}
fn poll(&self) -> Option<&T> {
if self.state.load(Ordering::Acquire) == LAZY_STATE_DONE {
Some(unsafe { &*(*self.data.as_ptr()).as_ptr() })
} else {
None
}
}
}
impl<T: fmt::Debug, F: FnMut() -> T> fmt::Debug for Lazy<T, F> {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
f.debug_struct("Lazy")
.field("state", &self.state.load(Ordering::Acquire))
.field("create", &"<closure>")
.field("data", &self.poll())
.finish()
}
}
impl<T, F> Drop for Lazy<T, F> {
fn drop(&mut self) {
if *self.state.get_mut() == LAZY_STATE_DONE {
unsafe {
self.data.get_mut().assume_init_drop();
}
}
}
}
struct Guard<'a> {
state: &'a AtomicU8,
}
impl<'a> Drop for Guard<'a> {
fn drop(&mut self) {
self.state.store(LAZY_STATE_INIT, Ordering::Release);
}
}
}
#[cfg(test)]
mod tests {
use super::*;
fn assert_send<T: Send>() {}
fn assert_sync<T: Sync>() {}
fn assert_unwind<T: core::panic::UnwindSafe>() {}
fn assert_refunwind<T: core::panic::RefUnwindSafe>() {}
#[test]
fn oibits() {
assert_send::<Lazy<u64>>();
assert_sync::<Lazy<u64>>();
assert_unwind::<Lazy<u64>>();
assert_refunwind::<Lazy<u64>>();
}
}