use super::assert_future;
use core::pin::Pin;
use futures_core::task::{Context, Poll};
use futures_core::{FusedFuture, Future, Stream};
use pin_project_lite::pin_project;
pin_project! {
#[derive(Debug, Clone)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct PollImmediate<T> {
#[pin]
future: Option<T>
}
}
impl<T, F> Future for PollImmediate<F>
where
F: Future<Output = T>,
{
type Output = Option<T>;
#[inline]
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<T>> {
let mut this = self.project();
let inner =
this.future.as_mut().as_pin_mut().expect("PollImmediate polled after completion");
match inner.poll(cx) {
Poll::Ready(t) => {
this.future.set(None);
Poll::Ready(Some(t))
}
Poll::Pending => Poll::Ready(None),
}
}
}
impl<T: Future> FusedFuture for PollImmediate<T> {
fn is_terminated(&self) -> bool {
self.future.is_none()
}
}
impl<T, F> Stream for PollImmediate<F>
where
F: Future<Output = T>,
{
type Item = Poll<T>;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let mut this = self.project();
match this.future.as_mut().as_pin_mut() {
None => Poll::Ready(None),
Some(fut) => Poll::Ready(Some(fut.poll(cx).map(|t| {
this.future.set(None);
t
}))),
}
}
}
pub fn poll_immediate<F: Future>(f: F) -> PollImmediate<F> {
assert_future::<Option<F::Output>, PollImmediate<F>>(PollImmediate { future: Some(f) })
}