#![allow(unsafe_code)]
use core::fmt;
use core::future::Future;
use core::mem::MaybeUninit;
use core::pin::Pin;
use core::task::{Context, Poll};
use super::CancellationToken;
use crate::sync::event::{Event, EventWait};
pub(super) struct RunUntilCancelled<'a, F: Future> {
pub(super) cancellation: WaitForCancellationFuture<'a>,
pub(super) future: F,
}
impl<F: Future> Future for RunUntilCancelled<'_, F> {
type Output = Option<F::Output>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = unsafe { self.get_unchecked_mut() };
let future = unsafe { Pin::new_unchecked(&mut this.future) };
if let Poll::Ready(output) = future.poll(cx) {
Poll::Ready(Some(output))
} else if Pin::new(&mut this.cancellation).poll(cx).is_ready() {
Poll::Ready(None)
} else {
Poll::Pending
}
}
}
#[must_use = "futures do nothing unless polled"]
pub struct WaitForCancellationFuture<'a> {
pub(super) wait: EventWait<'a>,
}
impl fmt::Debug for WaitForCancellationFuture<'_> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("WaitForCancellationFuture").finish()
}
}
impl Future for WaitForCancellationFuture<'_> {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
Pin::new(&mut self.get_mut().wait).poll(cx)
}
}
#[repr(transparent)]
struct MaybeDangling<T>(MaybeUninit<T>);
impl<T> MaybeDangling<T> {
fn new(inner: T) -> Self {
Self(MaybeUninit::new(inner))
}
fn get_mut(&mut self) -> &mut T {
unsafe { self.0.assume_init_mut() }
}
}
impl<T> Drop for MaybeDangling<T> {
fn drop(&mut self) {
unsafe { self.0.assume_init_drop() };
}
}
#[must_use = "futures do nothing unless polled"]
pub struct WaitForCancellationFutureOwned {
wait: MaybeDangling<EventWait<'static>>,
_token: CancellationToken,
}
impl WaitForCancellationFutureOwned {
pub(super) fn new(token: CancellationToken) -> Self {
let event: *const Event = &token.node.cancelled;
let wait = unsafe { (*event).wait() };
Self {
wait: MaybeDangling::new(wait),
_token: token,
}
}
}
impl fmt::Debug for WaitForCancellationFutureOwned {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("WaitForCancellationFutureOwned").finish()
}
}
impl Future for WaitForCancellationFutureOwned {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
Pin::new(self.get_mut().wait.get_mut()).poll(cx)
}
}