use core::future::Future;
use core::mem;
use core::pin::Pin;
use core::task::{Context, Poll};
#[derive(Hash, Ord, PartialOrd, Eq, PartialEq, Copy, Clone, Debug)]
pub enum Polling<Fut: Future> {
Ready(Fut::Output),
Pending(Fut),
Done,
}
pub trait FuturePollingExt: Future {
fn polling(self) -> Polling<Self>
where
Self: Sized,
{
Polling::Pending(self)
}
}
impl<Fut: Future> FuturePollingExt for Fut {}
impl<Fut: Future> Polling<Fut> {
pub fn into_ready(self) -> Option<Fut::Output> {
if let Polling::Ready(out) = self {
Some(out)
} else {
None
}
}
pub fn into_pending(self) -> Option<Fut> {
if let Polling::Pending(fut) = self {
Some(fut)
} else {
None
}
}
pub fn into_poll(self) -> Poll<Fut::Output> {
if let Polling::Ready(out) = self {
Poll::Ready(out)
} else {
Poll::Pending
}
}
pub fn is_ready(&self) -> bool {
if let Polling::Ready(_) = self {
true
} else {
false
}
}
pub fn is_pending(&self) -> bool {
if let Polling::Pending(_) = self {
true
} else {
false
}
}
pub fn is_done(&self) -> bool {
if let Polling::Done = self {
true
} else {
false
}
}
pub fn as_ready(&self) -> Option<&Fut::Output> {
if let Polling::Ready(out) = self {
Some(out)
} else {
None
}
}
pub fn as_pending(&self) -> Option<&Fut> {
if let Polling::Pending(fut) = self {
Some(fut)
} else {
None
}
}
pub fn as_poll(&self) -> Poll<&Fut::Output> {
if let Polling::Ready(out) = self {
Poll::Ready(out)
} else {
Poll::Pending
}
}
pub fn as_ready_mut(&mut self) -> Option<&mut Fut::Output> {
if let Polling::Ready(out) = self {
Some(out)
} else {
None
}
}
pub fn as_pending_mut(&mut self) -> Option<&mut Fut> {
if let Polling::Pending(fut) = self {
Some(fut)
} else {
None
}
}
pub fn as_poll_mut(&mut self) -> Poll<&mut Fut::Output> {
if let Polling::Ready(out) = self {
Poll::Ready(out)
} else {
Poll::Pending
}
}
pub fn replace(&mut self, with: Self) -> Self {
mem::replace(self, with)
}
pub fn take(&mut self) -> Polling<Fut> {
mem::replace(self, Polling::Done)
}
pub fn take_ready(&mut self) -> Option<Fut::Output> {
if self.is_ready() {
self.take().into_ready()
} else {
None
}
}
pub fn take_pending(&mut self) -> Option<Fut> {
if self.is_pending() {
self.take().into_pending()
} else {
None
}
}
pub fn take_poll(&mut self) -> Poll<Fut::Output> {
if self.is_ready() {
self.take().into_poll()
} else {
Poll::Pending
}
}
pub async fn poll_once(&mut self) -> Poll<Fut::Output> {
self.polling_once().await;
self.take_poll()
}
pub async fn polling_once(&mut self) -> &mut Self {
if self.is_ready() {
return self;
} else if self.is_done() {
panic!("output already extracted");
}
struct PollingOnce<'polling, Fut: Future> {
polling: &'polling mut Polling<Fut>,
}
impl<Fut: Future> Future for PollingOnce<'_, Fut> {
type Output = ();
fn poll(self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Self::Output> {
let this = unsafe { self.get_unchecked_mut() };
if let Polling::Pending(fut) = this.polling {
if let Poll::Ready(out) = unsafe { Pin::new_unchecked(fut) }.poll(ctx) {
this.polling.replace(Polling::Ready(out));
}
Poll::Ready(())
} else {
unreachable!();
}
}
}
PollingOnce { polling: self }.await;
self
}
}
impl<T, E, Fut: Future<Output = Result<T, E>>> Polling<Fut> {
pub fn try_into_ready(self) -> Result<Option<T>, E> {
self.into_ready().transpose()
}
pub fn try_into_poll(self) -> Result<Poll<T>, E> {
match self.into_poll() {
Poll::Ready(Ok(ok)) => Ok(Poll::Ready(ok)),
Poll::Ready(Err(err)) => Err(err),
Poll::Pending => Ok(Poll::Pending),
}
}
pub fn try_as_ready(&self) -> Result<Option<&T>, &E> {
self.as_ready().map(Result::as_ref).transpose()
}
pub fn try_as_poll(&self) -> Result<Poll<&T>, &E> {
match self.as_poll() {
Poll::Ready(Ok(ok)) => Ok(Poll::Ready(ok)),
Poll::Ready(Err(err)) => Err(err),
Poll::Pending => Ok(Poll::Pending),
}
}
pub fn try_as_ready_mut(&mut self) -> Result<Option<&mut T>, &mut E> {
self.as_ready_mut().map(Result::as_mut).transpose()
}
pub fn try_as_poll_mut(&mut self) -> Result<Poll<&mut T>, &mut E> {
match self.as_poll_mut() {
Poll::Ready(Ok(ok)) => Ok(Poll::Ready(ok)),
Poll::Ready(Err(err)) => Err(err),
Poll::Pending => Ok(Poll::Pending),
}
}
pub fn try_take_ready(&mut self) -> Result<Option<T>, E> {
self.take_ready().transpose()
}
pub fn try_take_poll(&mut self) -> Result<Poll<T>, E> {
match self.take_poll() {
Poll::Ready(Ok(ok)) => Ok(Poll::Ready(ok)),
Poll::Ready(Err(err)) => Err(err),
Poll::Pending => Ok(Poll::Pending),
}
}
pub fn try_poll(self: Pin<&mut Self>, ctx: &mut Context) -> Result<Poll<T>, E> {
match self.poll(ctx) {
Poll::Ready(Ok(ok)) => Ok(Poll::Ready(ok)),
Poll::Ready(Err(err)) => Err(err),
Poll::Pending => Ok(Poll::Pending),
}
}
pub async fn try_poll_once(&mut self) -> Result<Poll<T>, E> {
self.polling_once().await;
self.try_take_poll()
}
}
impl<Fut: Future> Future for Polling<Fut> {
type Output = Fut::Output;
fn poll(self: Pin<&mut Self>, ctx: &mut Context) -> Poll<Self::Output> {
let this = unsafe { self.get_unchecked_mut() };
match this {
Polling::Ready(_) => this.take().into_poll(),
Polling::Pending(fut) => {
let poll = unsafe { Pin::new_unchecked(fut) }.poll(ctx);
if poll.is_ready() {
this.take();
}
poll
},
Polling::Done => panic!("output already extracted"),
}
}
}
impl<Fut: Future> From<Fut> for Polling<Fut> {
fn from(fut: Fut) -> Self {
Polling::Pending(fut)
}
}