#![cfg_attr(feature = "critical-section", no_std)]
extern crate alloc;
#[cfg(any(not(feature = "critical-section"), feature = "std"))]
extern crate std;
use alloc::{boxed::Box, vec, vec::Vec};
use core::{
cell::UnsafeCell,
convert::Infallible,
fmt,
future::{Future, IntoFuture},
marker::{PhantomData, PhantomPinned},
mem::{self, ManuallyDrop, MaybeUninit},
panic::{RefUnwindSafe, UnwindSafe},
pin::{pin, Pin},
ptr,
sync::atomic::{AtomicPtr, AtomicUsize, Ordering},
task,
};
#[cfg(feature = "critical-section")]
struct Mutex<T> {
data: UnsafeCell<T>,
locked: core::sync::atomic::AtomicBool,
}
#[cfg(feature = "critical-section")]
impl<T> Mutex<T> {
const fn new(data: T) -> Self {
Mutex { data: UnsafeCell::new(data), locked: core::sync::atomic::AtomicBool::new(false) }
}
}
#[cfg(not(feature = "critical-section"))]
use std::sync::Mutex;
#[cfg(feature = "critical-section")]
fn with_lock<T, R>(mutex: &Mutex<T>, f: impl FnOnce(&mut T) -> R) -> R {
struct Guard<'a, T>(&'a Mutex<T>);
impl<'a, T> Drop for Guard<'a, T> {
fn drop(&mut self) {
self.0.locked.store(false, Ordering::Relaxed);
}
}
critical_section::with(|_| {
if mutex.locked.swap(true, Ordering::Relaxed) {
panic!("Attempted reentrant locking");
}
let guard = Guard(mutex);
let rv = unsafe { f(&mut *mutex.data.get()) };
drop(guard);
rv
})
}
#[cfg(not(feature = "critical-section"))]
fn with_lock<T, R>(mutex: &Mutex<T>, f: impl FnOnce(&mut T) -> R) -> R {
f(&mut *mutex.lock().unwrap())
}
pub struct OnceCell<T> {
value: UnsafeCell<MaybeUninit<T>>,
inner: Inner,
_marker: PhantomData<T>,
}
unsafe impl<T: Sync + Send> Sync for OnceCell<T> {}
unsafe impl<T: Send> Send for OnceCell<T> {}
impl<T> Unpin for OnceCell<T> {}
impl<T: RefUnwindSafe + UnwindSafe> RefUnwindSafe for OnceCell<T> {}
impl<T: UnwindSafe> UnwindSafe for OnceCell<T> {}
struct Inner {
state: AtomicUsize,
queue: AtomicPtr<Queue>,
}
impl fmt::Debug for Inner {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
let state = self.state.load(Ordering::Relaxed);
let queue = self.queue.load(Ordering::Relaxed);
fmt.debug_struct("Inner")
.field("ready", &(state & READY_BIT != 0))
.field("quick_init", &(state & QINIT_BIT != 0))
.field("refcount", &(state & (QINIT_BIT - 1)))
.field("queue", &queue)
.finish()
}
}
struct Queue {
wakers: Mutex<Option<Vec<task::Waker>>>,
}
struct QueueRef<'a> {
inner: &'a Inner,
queue: *const Queue,
}
unsafe impl<'a> Sync for QueueRef<'a> {}
unsafe impl<'a> Send for QueueRef<'a> {}
#[derive(Debug)]
struct QuickInitGuard<'a> {
inner: &'a Inner,
ready: bool,
}
struct QueueWaiter<'a> {
guard: Option<QueueRef<'a>>,
}
struct QueueHead<'a> {
guard: QueueRef<'a>,
}
const NEW: usize = 0x0;
const QINIT_BIT: usize = 1 + (usize::MAX >> 2);
const READY_BIT: usize = 1 + (usize::MAX >> 1);
const EMPTY_STATE: usize = !0;
impl Inner {
const fn new() -> Self {
Inner { state: AtomicUsize::new(NEW), queue: AtomicPtr::new(ptr::null_mut()) }
}
const fn new_ready() -> Self {
Inner { state: AtomicUsize::new(READY_BIT), queue: AtomicPtr::new(ptr::null_mut()) }
}
#[cold]
fn initialize(&self, try_quick: bool) -> Result<QueueWaiter, QuickInitGuard> {
if try_quick {
if self
.state
.compare_exchange(NEW, QINIT_BIT, Ordering::Acquire, Ordering::Relaxed)
.is_ok()
{
return Err(QuickInitGuard { inner: self, ready: false });
}
}
let prev_state = self.state.fetch_add(1, Ordering::Acquire);
let mut guard = QueueRef { inner: self, queue: self.queue.load(Ordering::Acquire) };
if guard.queue.is_null() && prev_state & READY_BIT == 0 {
let wakers = Mutex::new(None);
let new_queue = Box::into_raw(Box::new(Queue { wakers }));
match self.queue.compare_exchange(
ptr::null_mut(),
new_queue,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_null) => {
guard.queue = new_queue;
}
Err(actual) => {
guard.queue = actual;
unsafe {
drop(Box::from_raw(new_queue));
}
}
}
}
Ok(QueueWaiter { guard: Some(guard) })
}
fn set_ready(&self) {
let prev_state = self.state.fetch_or(READY_BIT, Ordering::Release);
debug_assert_eq!(prev_state & READY_BIT, 0, "Invalid state: someone else set READY_BIT");
}
}
impl<'a> Drop for QueueRef<'a> {
fn drop(&mut self) {
let prev_state = self.inner.state.fetch_sub(1, Ordering::Release);
let curr_state = prev_state - 1;
if curr_state == READY_BIT || curr_state == READY_BIT | QINIT_BIT {
let queue = self.inner.queue.swap(ptr::null_mut(), Ordering::Acquire);
if !queue.is_null() {
unsafe {
drop(Box::from_raw(queue));
}
}
}
}
}
impl<'a> Drop for QuickInitGuard<'a> {
fn drop(&mut self) {
let fast_target = if self.ready { READY_BIT } else { NEW };
if self
.inner
.state
.compare_exchange(QINIT_BIT, fast_target, Ordering::Release, Ordering::Relaxed)
.is_ok()
{
if self.ready {
let queue = self.inner.queue.swap(ptr::null_mut(), Ordering::Relaxed);
if !queue.is_null() {
core::sync::atomic::fence(Ordering::Acquire);
unsafe {
drop(Box::from_raw(queue));
}
}
}
return;
}
let waiter = self.inner.initialize(false).expect("Got a QuickInitGuard in slow init");
let guard = waiter.guard.expect("No guard available even without polling");
debug_assert!(!guard.queue.is_null(), "Queue must not be NULL when READY_BIT is not set");
let queue = unsafe { &*guard.queue };
with_lock(&queue.wakers, |lock| {
lock.get_or_insert_with(Vec::new);
let prev_state = if self.ready {
self.inner.state.fetch_add(QINIT_BIT, Ordering::Release)
} else {
self.inner.state.fetch_sub(QINIT_BIT, Ordering::Relaxed)
};
debug_assert_eq!(
prev_state & (QINIT_BIT | READY_BIT),
QINIT_BIT,
"Invalid state during QuickInitGuard drop"
);
});
drop(QueueHead { guard })
}
}
impl Drop for Inner {
fn drop(&mut self) {
let queue = *self.queue.get_mut();
if !queue.is_null() {
unsafe {
drop(Box::from_raw(queue));
}
}
}
}
impl<'a> Future for QueueWaiter<'a> {
type Output = Option<QueueHead<'a>>;
fn poll(
mut self: Pin<&mut Self>,
cx: &mut task::Context<'_>,
) -> task::Poll<Option<QueueHead<'a>>> {
let guard = self.guard.as_ref().expect("Polled future after finished");
let state = guard.inner.state.load(Ordering::Acquire);
if state & READY_BIT != 0 {
return task::Poll::Ready(None);
}
let queue = unsafe { &*guard.queue };
let rv = with_lock(&queue.wakers, |lock| {
let state = guard.inner.state.load(Ordering::Acquire);
if state & READY_BIT != 0 {
return task::Poll::Ready(None);
}
match lock.as_mut() {
None if state & QINIT_BIT == 0 => {
*lock = Some(Vec::new());
task::Poll::Ready(Some(()))
}
None => {
let waker = cx.waker().clone();
*lock = Some(vec![waker]);
task::Poll::Pending
}
Some(wakers) => {
let my_waker = cx.waker();
for waker in wakers.iter() {
if waker.will_wake(my_waker) {
return task::Poll::Pending;
}
}
wakers.push(my_waker.clone());
task::Poll::Pending
}
}
});
rv.map(|o| o.map(|()| QueueHead { guard: self.guard.take().unwrap() }))
}
}
impl<'a> Drop for QueueHead<'a> {
fn drop(&mut self) {
let queue = unsafe { &*self.guard.queue };
let wakers =
with_lock(&queue.wakers, Option::take).expect("QueueHead dropped without a waker list");
for waker in wakers {
waker.wake();
}
}
}
enum Step<'a> {
Start { inner: &'a Inner },
Quick { guard: QuickInitGuard<'a> },
Wait { guard: QueueWaiter<'a> },
Run { head: QueueHead<'a> },
Done,
}
enum EitherHead<'a, 'b> {
Quick(&'b mut QuickInitGuard<'a>),
Normal(&'b QueueHead<'a>),
}
impl EitherHead<'_, '_> {
fn set_ready(&mut self) {
match self {
Self::Quick(guard) => guard.ready = true,
Self::Normal(head) => head.guard.inner.set_ready(),
}
}
}
impl<'a> Step<'a> {
fn poll_init<F, R>(&mut self, cx: &mut task::Context<'_>, done: R, mut init: F) -> task::Poll<R>
where
F: FnMut(&mut task::Context<'_>, EitherHead<'a, '_>) -> task::Poll<R>,
{
loop {
match mem::replace(self, Step::Done) {
Step::Start { inner } => {
let state = inner.state.load(Ordering::Acquire);
if state & READY_BIT == 0 {
*self = match inner.initialize(state == NEW) {
Err(guard) => Step::Quick { guard },
Ok(guard) => Step::Wait { guard },
};
continue;
}
return task::Poll::Ready(done);
}
Step::Quick { mut guard } => {
let rv = init(cx, EitherHead::Quick(&mut guard));
if rv.is_pending() {
*self = Step::Quick { guard };
}
return rv;
}
Step::Wait { mut guard } => match Pin::new(&mut guard).poll(cx) {
task::Poll::Pending => {
*self = Step::Wait { guard };
return task::Poll::Pending;
}
task::Poll::Ready(None) => {
return task::Poll::Ready(done);
}
task::Poll::Ready(Some(head)) => {
*self = Step::Run { head };
continue;
}
},
Step::Run { head } => {
let rv = init(cx, EitherHead::Normal(&head));
if rv.is_pending() {
*self = Step::Run { head };
}
return rv;
}
Step::Done => {
panic!("Polled future after completion");
}
}
}
}
}
struct InitFuture<'a, T, F> {
cell: &'a OnceCell<T>,
init: F,
step: Step<'a>,
}
impl<'a, T, F> InitFuture<'a, T, F> {
fn new<R>(cell: &'a OnceCell<T>, init: F) -> Self
where
F: for<'c> FnMut(&mut task::Context<'c>) -> task::Poll<R> + Unpin,
{
Self { cell, init, step: Step::Start { inner: &cell.inner } }
}
}
impl<'a, T, F, E> Future for InitFuture<'a, T, F>
where
F: for<'c> FnMut(&mut task::Context<'c>) -> task::Poll<Result<T, E>> + Unpin,
{
type Output = Result<&'a T, E>;
fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Self::Output> {
struct Filled;
let this = self.get_mut();
let cell = this.cell;
let init = &mut this.init;
this.step
.poll_init(cx, Ok(Filled), |cx, mut head| {
let value = task::ready!(init(cx))?;
unsafe {
(*cell.value.get()).write(value);
}
head.set_ready();
task::Poll::Ready(Ok(Filled))
})
.map(|r| {
r.map(|Filled| unsafe { (*this.cell.value.get()).assume_init_ref() })
})
}
}
impl<T> OnceCell<T> {
pub const fn new() -> Self {
Self {
value: UnsafeCell::new(MaybeUninit::uninit()),
inner: Inner::new(),
_marker: PhantomData,
}
}
pub const fn new_with(value: T) -> Self {
Self {
value: UnsafeCell::new(MaybeUninit::new(value)),
inner: Inner::new_ready(),
_marker: PhantomData,
}
}
pub async fn get_or_init(&self, init: impl Future<Output = T>) -> &T {
let mut init = pin!(init);
match InitFuture::new(self, |cx| init.as_mut().poll(cx).map(Ok::<T, Infallible>)).await {
Ok(t) => t,
Err(e) => match e {},
}
}
pub async fn get_or_try_init<E>(
&self,
init: impl Future<Output = Result<T, E>>,
) -> Result<&T, E> {
let mut init = pin!(init);
InitFuture::new(self, |cx| init.as_mut().poll(cx)).await
}
pub fn get(&self) -> Option<&T> {
let state = self.inner.state.load(Ordering::Acquire);
if state & READY_BIT == 0 {
None
} else {
Some(unsafe { (*self.value.get()).assume_init_ref() })
}
}
pub fn get_mut(&mut self) -> Option<&mut T> {
let state = *self.inner.state.get_mut();
if state & READY_BIT == 0 {
None
} else {
Some(unsafe { self.value.get_mut().assume_init_mut() })
}
}
pub fn take(&mut self) -> Option<T> {
let state = *self.inner.state.get_mut();
self.inner = Inner::new();
if state & READY_BIT == 0 {
None
} else {
Some(unsafe { self.value.get_mut().assume_init_read() })
}
}
pub fn into_inner(mut self) -> Option<T> {
self.take()
}
}
impl<T: fmt::Debug> fmt::Debug for OnceCell<T> {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
let value = self.get();
fmt.debug_struct("OnceCell").field("value", &value).field("inner", &self.inner).finish()
}
}
impl<T> Drop for OnceCell<T> {
fn drop(&mut self) {
let state = *self.inner.state.get_mut();
if state & READY_BIT != 0 {
unsafe {
self.value.get_mut().assume_init_drop();
}
}
}
}
impl<T> Default for OnceCell<T> {
fn default() -> Self {
Self::new()
}
}
impl<T> From<T> for OnceCell<T> {
fn from(value: T) -> Self {
Self::new_with(value)
}
}
#[cfg(test)]
mod test {
use super::*;
use alloc::sync::Arc;
use core::pin::pin;
#[derive(Default)]
struct CountWaker(AtomicUsize);
impl alloc::task::Wake for CountWaker {
fn wake(self: Arc<Self>) {
self.0.fetch_add(1, Ordering::Relaxed);
}
}
struct CmdWait<'a>(&'a AtomicUsize);
impl Future for CmdWait<'_> {
type Output = usize;
fn poll(self: Pin<&mut Self>, _: &mut task::Context<'_>) -> task::Poll<usize> {
match self.0.load(Ordering::Relaxed) {
0 => task::Poll::Pending,
n => task::Poll::Ready(n),
}
}
}
impl Drop for CmdWait<'_> {
fn drop(&mut self) {
if self.0.load(Ordering::Relaxed) == 6 {
panic!("Panic on drop");
}
}
}
async fn maybe(cmd: &AtomicUsize, cell: &OnceCell<usize>) -> Result<usize, usize> {
cell.get_or_try_init(async {
match dbg!(CmdWait(cmd).await) {
1 => Err(1),
2 => Ok(2),
_ => unreachable!(),
}
})
.await
.map(|v| *v)
}
async fn never_init(cell: &OnceCell<usize>) {
let v = cell.get_or_init(async { unreachable!() }).await;
assert_eq!(v, &2);
}
#[test]
fn slow_path() {
let w = Arc::new(CountWaker::default()).into();
let mut cx = std::task::Context::from_waker(&w);
let cmd = AtomicUsize::new(0);
let cell = OnceCell::new();
let mut f1 = pin!(maybe(&cmd, &cell));
let mut f2 = pin!(never_init(&cell));
println!("{:?}", cell);
assert!(f1.as_mut().poll(&mut cx).is_pending());
println!("{:?}", cell);
assert!(f2.as_mut().poll(&mut cx).is_pending());
println!("{:?}", cell);
cmd.store(2, Ordering::Relaxed);
assert!(f2.as_mut().poll(&mut cx).is_pending());
assert!(f1.as_mut().poll(&mut cx).is_ready());
println!("{:?}", cell);
assert!(f2.as_mut().poll(&mut cx).is_ready());
}
#[test]
fn fast_path_tricked() {
let w = Arc::new(CountWaker::default()).into();
let mut cx = std::task::Context::from_waker(&w);
let cmd = AtomicUsize::new(0);
let cell = OnceCell::new();
let mut f1 = pin!(maybe(&cmd, &cell));
let mut f2 = pin!(never_init(&cell));
println!("{:?}", cell);
assert!(f1.as_mut().poll(&mut cx).is_pending());
println!("{:?}", cell);
assert!(f2.as_mut().poll(&mut cx).is_pending());
println!("{:?}", cell);
cmd.store(2, Ordering::Relaxed);
f2.set(never_init(&cell));
println!("{:?}", cell);
assert!(f1.as_mut().poll(&mut cx).is_ready());
println!("{:?}", cell);
assert!(f2.as_mut().poll(&mut cx).is_ready());
}
#[test]
fn second_try() {
let waker = Arc::new(CountWaker::default());
let w = waker.clone().into();
let mut cx = std::task::Context::from_waker(&w);
let cmd = AtomicUsize::new(0);
let cell = OnceCell::new();
let mut f1 = pin!(maybe(&cmd, &cell));
let mut f2 = pin!(maybe(&cmd, &cell));
let mut f3 = pin!(maybe(&cmd, &cell));
let mut f4 = pin!(maybe(&cmd, &cell));
assert!(f1.as_mut().poll(&mut cx).is_pending());
assert_eq!(cell.inner.state.load(Ordering::Relaxed), QINIT_BIT);
assert!(f2.as_mut().poll(&mut cx).is_pending());
assert!(f3.as_mut().poll(&mut cx).is_pending());
assert!(f4.as_mut().poll(&mut cx).is_pending());
assert_eq!(cell.inner.state.load(Ordering::Relaxed), QINIT_BIT | 3);
cmd.store(1, Ordering::Relaxed);
assert!(f2.as_mut().poll(&mut cx).is_pending());
assert_eq!(waker.0.load(Ordering::Relaxed), 0);
assert_eq!(f1.as_mut().poll(&mut cx), task::Poll::Ready(Err(1)));
assert_eq!(cell.inner.state.load(Ordering::Relaxed), 3);
assert_eq!(waker.0.load(Ordering::Relaxed), 1);
f4.set(maybe(&cmd, &cell));
assert_eq!(cell.inner.state.load(Ordering::Relaxed), 2);
cmd.store(0, Ordering::Relaxed);
assert!(f2.as_mut().poll(&mut cx).is_pending());
cmd.store(2, Ordering::Relaxed);
assert!(f3.as_mut().poll(&mut cx).is_pending());
assert_eq!(waker.0.load(Ordering::Relaxed), 1);
assert_eq!(f2.as_mut().poll(&mut cx), task::Poll::Ready(Ok(2)));
cmd.store(3, Ordering::Relaxed);
assert_eq!(f4.as_mut().poll(&mut cx), task::Poll::Ready(Ok(2)));
assert_eq!(waker.0.load(Ordering::Relaxed), 2);
assert_eq!(cell.inner.state.load(Ordering::Relaxed), READY_BIT | 1);
assert!(!cell.inner.queue.load(Ordering::Relaxed).is_null());
assert_eq!(f3.as_mut().poll(&mut cx), task::Poll::Ready(Ok(2)));
assert_eq!(cell.inner.state.load(Ordering::Relaxed), READY_BIT);
assert!(cell.inner.queue.load(Ordering::Relaxed).is_null());
assert_eq!(waker.0.load(Ordering::Relaxed), 2);
}
#[test]
fn lazy_panic() {
let w = Arc::new(CountWaker::default()).into();
let cmd = AtomicUsize::new(6);
let lz = Lazy::new(CmdWait(&cmd));
assert_eq!(std::mem::size_of_val(&lz), 3 * std::mem::size_of::<usize>(), "Extra overhead?");
assert!(std::panic::catch_unwind(|| {
let mut cx = std::task::Context::from_waker(&w);
pin!(lz.get_unpin()).poll(&mut cx)
})
.is_err());
assert_eq!(lz.try_get(), Some(&6));
}
}
union LazyState<T, F> {
running: ManuallyDrop<F>,
ready: ManuallyDrop<T>,
_empty: (),
}
pub struct Lazy<T, F> {
value: UnsafeCell<LazyState<T, F>>,
inner: Inner,
}
unsafe impl<T: Send + Sync, F: Send> Sync for Lazy<T, F> {}
unsafe impl<T: Send, F: Send> Send for Lazy<T, F> {}
impl<T: Unpin, F: Unpin> Unpin for Lazy<T, F> {}
impl<T: RefUnwindSafe + UnwindSafe, F: UnwindSafe> RefUnwindSafe for Lazy<T, F> {}
impl<T: UnwindSafe, F: UnwindSafe> UnwindSafe for Lazy<T, F> {}
impl<T, F> Lazy<T, F>
where
F: Future<Output = T>,
{
pub const fn new(future: F) -> Self {
Self::from_future(future)
}
pub async fn get(self: Pin<&Self>) -> Pin<&T> {
self.await
}
}
struct LazyFuture<'a, T, F> {
lazy: &'a Lazy<T, F>,
step: Step<'a>,
_pin: PhantomPinned,
}
impl<'a, T, F> LazyFuture<'a, T, F>
where
F: Future<Output = T>,
{
fn poll(&mut self, cx: &mut task::Context<'_>) -> task::Poll<&'a T> {
struct ReplaceGuard<'a, 'b, T, F> {
this: &'a Lazy<T, F>,
value: ManuallyDrop<T>,
head: EitherHead<'a, 'b>,
}
impl<T, F> Drop for ReplaceGuard<'_, '_, T, F> {
fn drop(&mut self) {
unsafe {
let value = ManuallyDrop::take(&mut self.value);
(*self.this.value.get()).ready = ManuallyDrop::new(value);
}
self.head.set_ready();
}
}
let this = &self.lazy;
self.step
.poll_init(cx, (), |cx, head| {
let init = unsafe { Pin::new_unchecked(&mut *(*this.value.get()).running) };
let value = ManuallyDrop::new(task::ready!(init.poll(cx)));
unsafe {
let guard = ReplaceGuard { this, value, head };
ManuallyDrop::drop(&mut (*this.value.get()).running);
drop(guard);
}
task::Poll::Ready(())
})
.map(|()| {
unsafe { &*(*this.value.get()).ready }
})
}
}
pub struct LazyFuturePin<'a, T, F>(LazyFuture<'a, T, F>);
impl<'a, T, F> IntoFuture for Pin<&'a Lazy<T, F>>
where
F: Future<Output = T>,
{
type Output = Pin<&'a T>;
type IntoFuture = LazyFuturePin<'a, T, F>;
fn into_future(self) -> Self::IntoFuture {
let lazy = unsafe { Pin::into_inner_unchecked(self) };
LazyFuturePin(LazyFuture {
lazy,
step: Step::Start { inner: &lazy.inner },
_pin: PhantomPinned,
})
}
}
impl<'a, T, F> Future for LazyFuturePin<'a, T, F>
where
F: Future<Output = T>,
{
type Output = Pin<&'a T>;
fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<Pin<&'a T>> {
let inner = unsafe { &mut Pin::into_inner_unchecked(self).0 };
inner.poll(cx).map(|p| unsafe { Pin::new_unchecked(p) })
}
}
impl<T, F> Lazy<T, F>
where
F: Future<Output = T> + Unpin,
{
pub async fn get_unpin(&self) -> &T {
self.await
}
}
pub struct LazyFutureUnpin<'a, T, F>(LazyFuture<'a, T, F>);
impl<'a, T, F> IntoFuture for &'a Lazy<T, F>
where
F: Future<Output = T> + Unpin,
{
type Output = &'a T;
type IntoFuture = LazyFutureUnpin<'a, T, F>;
fn into_future(self) -> Self::IntoFuture {
LazyFutureUnpin(LazyFuture {
lazy: self,
step: Step::Start { inner: &self.inner },
_pin: PhantomPinned,
})
}
}
impl<'a, T, F> Future for LazyFutureUnpin<'a, T, F>
where
F: Future<Output = T> + Unpin,
{
type Output = &'a T;
fn poll(self: Pin<&mut Self>, cx: &mut task::Context<'_>) -> task::Poll<&'a T> {
unsafe { Pin::into_inner_unchecked(self) }.0.poll(cx)
}
}
impl<T, F> Lazy<T, F> {
pub const fn from_future(future: F) -> Self {
Self {
value: UnsafeCell::new(LazyState { running: ManuallyDrop::new(future) }),
inner: Inner::new(),
}
}
pub const fn with_value(value: T) -> Self {
Self {
value: UnsafeCell::new(LazyState { ready: ManuallyDrop::new(value) }),
inner: Inner::new_ready(),
}
}
pub fn try_get(&self) -> Option<&T> {
let state = self.inner.state.load(Ordering::Acquire);
if state & READY_BIT == 0 {
None
} else {
unsafe { Some(&(*self.value.get()).ready) }
}
}
pub fn try_get_mut(self: Pin<&mut Self>) -> Option<Pin<&mut T>> {
let this = unsafe { self.get_unchecked_mut() };
let state = *this.inner.state.get_mut();
if state & READY_BIT == 0 {
None
} else {
unsafe { Some(Pin::new_unchecked(&mut this.value.get_mut().ready)) }
}
}
pub fn try_get_mut_unpin(&mut self) -> Option<&mut T> {
let state = *self.inner.state.get_mut();
if state & READY_BIT == 0 {
None
} else {
unsafe { Some(&mut self.value.get_mut().ready) }
}
}
pub fn into_inner(self) -> Option<T> {
self.into_parts().ok()
}
pub fn into_parts(mut self) -> Result<T, F> {
let state = *self.inner.state.get_mut();
unsafe {
*self.inner.state.get_mut() = EMPTY_STATE;
if state & READY_BIT == 0 {
Err(ptr::read(&*self.value.get_mut().running))
} else {
Ok(ptr::read(&*self.value.get_mut().ready))
}
}
}
pub fn replace_and_take(self: Pin<&mut Self>, replacement: Self) -> Option<T>
where
T: Unpin,
{
let this = unsafe { self.get_unchecked_mut() };
let state = *this.inner.state.get_mut();
let value = if state & READY_BIT == 0 {
None
} else {
*this.inner.state.get_mut() = EMPTY_STATE;
Some(unsafe { ptr::read(&*this.value.get_mut().ready) })
};
*this = replacement;
value
}
}
impl<T, F> Drop for Lazy<T, F> {
fn drop(&mut self) {
let state = *self.inner.state.get_mut();
unsafe {
if state == EMPTY_STATE {
} else if state & READY_BIT == 0 {
ManuallyDrop::drop(&mut self.value.get_mut().running);
} else {
ManuallyDrop::drop(&mut self.value.get_mut().ready);
}
}
}
}
impl<T: fmt::Debug, F> fmt::Debug for Lazy<T, F> {
fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
let value = self.try_get();
fmt.debug_struct("Lazy").field("value", &value).field("inner", &self.inner).finish()
}
}