#[cfg(all(not(feature = "std"), feature = "alloc"))]
extern crate alloc;
#[doc(no_inline)]
pub use core::future::{pending, ready, Future, Pending, Ready};
use core::fmt;
use core::pin::Pin;
use pin_project_lite::pin_project;
#[cfg(feature = "std")]
use std::{
any::Any,
panic::{catch_unwind, AssertUnwindSafe, UnwindSafe},
};
#[cfg(feature = "race")]
use fastrand::Rng;
#[cfg(all(not(feature = "std"), feature = "alloc"))]
use alloc::boxed::Box;
use core::task::{Context, Poll};
#[cfg(feature = "std")]
pub fn block_on<T>(future: impl Future<Output = T>) -> T {
use core::cell::RefCell;
use core::task::Waker;
use parking::Parker;
crate::pin!(future);
fn parker_and_waker() -> (Parker, Waker) {
let parker = Parker::new();
let unparker = parker.unparker();
let waker = Waker::from(unparker);
(parker, waker)
}
thread_local! {
static CACHE: RefCell<(Parker, Waker)> = RefCell::new(parker_and_waker());
}
CACHE.with(|cache| {
let tmp_cached;
let tmp_fresh;
let (parker, waker) = match cache.try_borrow_mut() {
Ok(cache) => {
tmp_cached = cache;
&*tmp_cached
}
Err(_) => {
tmp_fresh = parker_and_waker();
&tmp_fresh
}
};
let cx = &mut Context::from_waker(waker);
loop {
match future.as_mut().poll(cx) {
Poll::Ready(output) => return output,
Poll::Pending => parker.park(),
}
}
})
}
pub fn poll_once<T, F>(f: F) -> PollOnce<F>
where
F: Future<Output = T>,
{
PollOnce { f }
}
pin_project! {
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct PollOnce<F> {
#[pin]
f: F,
}
}
impl<F> fmt::Debug for PollOnce<F> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("PollOnce").finish()
}
}
impl<T, F> Future for PollOnce<F>
where
F: Future<Output = T>,
{
type Output = Option<T>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
match self.project().f.poll(cx) {
Poll::Ready(t) => Poll::Ready(Some(t)),
Poll::Pending => Poll::Ready(None),
}
}
}
pub fn poll_fn<T, F>(f: F) -> PollFn<F>
where
F: FnMut(&mut Context<'_>) -> Poll<T>,
{
PollFn { f }
}
pin_project! {
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct PollFn<F> {
f: F,
}
}
impl<F> fmt::Debug for PollFn<F> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("PollFn").finish()
}
}
impl<T, F> Future for PollFn<F>
where
F: FnMut(&mut Context<'_>) -> Poll<T>,
{
type Output = T;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<T> {
let this = self.project();
(this.f)(cx)
}
}
pub fn yield_now() -> YieldNow {
YieldNow(false)
}
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct YieldNow(bool);
impl Future for YieldNow {
type Output = ();
fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
if !self.0 {
self.0 = true;
cx.waker().wake_by_ref();
Poll::Pending
} else {
Poll::Ready(())
}
}
}
pub fn zip<F1, F2>(future1: F1, future2: F2) -> Zip<F1, F2>
where
F1: Future,
F2: Future,
{
Zip {
future1,
future2,
output1: None,
output2: None,
}
}
pin_project! {
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct Zip<F1, F2>
where
F1: Future,
F2: Future,
{
#[pin]
future1: F1,
output1: Option<F1::Output>,
#[pin]
future2: F2,
output2: Option<F2::Output>,
}
}
fn take_zip_from_parts<T1, T2>(o1: &mut Option<T1>, o2: &mut Option<T2>) -> Poll<(T1, T2)> {
match (o1.take(), o2.take()) {
(Some(t1), Some(t2)) => Poll::Ready((t1, t2)),
(o1x, o2x) => {
*o1 = o1x;
*o2 = o2x;
Poll::Pending
}
}
}
impl<F1, F2> Future for Zip<F1, F2>
where
F1: Future,
F2: Future,
{
type Output = (F1::Output, F2::Output);
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
if this.output1.is_none() {
if let Poll::Ready(out) = this.future1.poll(cx) {
*this.output1 = Some(out);
}
}
if this.output2.is_none() {
if let Poll::Ready(out) = this.future2.poll(cx) {
*this.output2 = Some(out);
}
}
take_zip_from_parts(this.output1, this.output2)
}
}
pub fn try_zip<T1, T2, E, F1, F2>(future1: F1, future2: F2) -> TryZip<F1, T1, F2, T2>
where
F1: Future<Output = Result<T1, E>>,
F2: Future<Output = Result<T2, E>>,
{
TryZip {
future1,
future2,
output1: None,
output2: None,
}
}
pin_project! {
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct TryZip<F1, T1, F2, T2> {
#[pin]
future1: F1,
output1: Option<T1>,
#[pin]
future2: F2,
output2: Option<T2>,
}
}
impl<T1, T2, E, F1, F2> Future for TryZip<F1, T1, F2, T2>
where
F1: Future<Output = Result<T1, E>>,
F2: Future<Output = Result<T2, E>>,
{
type Output = Result<(T1, T2), E>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
if this.output1.is_none() {
if let Poll::Ready(out) = this.future1.poll(cx) {
match out {
Ok(t) => *this.output1 = Some(t),
Err(err) => return Poll::Ready(Err(err)),
}
}
}
if this.output2.is_none() {
if let Poll::Ready(out) = this.future2.poll(cx) {
match out {
Ok(t) => *this.output2 = Some(t),
Err(err) => return Poll::Ready(Err(err)),
}
}
}
take_zip_from_parts(this.output1, this.output2).map(Ok)
}
}
pub fn or<T, F1, F2>(future1: F1, future2: F2) -> Or<F1, F2>
where
F1: Future<Output = T>,
F2: Future<Output = T>,
{
Or { future1, future2 }
}
pin_project! {
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct Or<F1, F2> {
#[pin]
future1: F1,
#[pin]
future2: F2,
}
}
impl<T, F1, F2> Future for Or<F1, F2>
where
F1: Future<Output = T>,
F2: Future<Output = T>,
{
type Output = T;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
if let Poll::Ready(t) = this.future1.poll(cx) {
return Poll::Ready(t);
}
if let Poll::Ready(t) = this.future2.poll(cx) {
return Poll::Ready(t);
}
Poll::Pending
}
}
#[cfg(all(feature = "race", feature = "std"))]
pub fn race<T, F1, F2>(future1: F1, future2: F2) -> Race<F1, F2>
where
F1: Future<Output = T>,
F2: Future<Output = T>,
{
Race {
future1,
future2,
rng: Rng::new(),
}
}
#[cfg(feature = "race")]
pub fn race_with_seed<T, F1, F2>(future1: F1, future2: F2, seed: u64) -> Race<F1, F2>
where
F1: Future<Output = T>,
F2: Future<Output = T>,
{
Race {
future1,
future2,
rng: Rng::with_seed(seed),
}
}
#[cfg(feature = "race")]
pin_project! {
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct Race<F1, F2> {
#[pin]
future1: F1,
#[pin]
future2: F2,
rng: Rng,
}
}
#[cfg(feature = "race")]
impl<T, F1, F2> Future for Race<F1, F2>
where
F1: Future<Output = T>,
F2: Future<Output = T>,
{
type Output = T;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
if this.rng.bool() {
if let Poll::Ready(t) = this.future1.poll(cx) {
return Poll::Ready(t);
}
if let Poll::Ready(t) = this.future2.poll(cx) {
return Poll::Ready(t);
}
} else {
if let Poll::Ready(t) = this.future2.poll(cx) {
return Poll::Ready(t);
}
if let Poll::Ready(t) = this.future1.poll(cx) {
return Poll::Ready(t);
}
}
Poll::Pending
}
}
#[cfg(feature = "std")]
pin_project! {
#[derive(Debug)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct CatchUnwind<F> {
#[pin]
inner: F,
}
}
#[cfg(feature = "std")]
impl<F: Future + UnwindSafe> Future for CatchUnwind<F> {
type Output = Result<F::Output, Box<dyn Any + Send>>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.project();
catch_unwind(AssertUnwindSafe(|| this.inner.poll(cx)))?.map(Ok)
}
}
#[cfg(feature = "alloc")]
pub type Boxed<T> = Pin<Box<dyn Future<Output = T> + Send + 'static>>;
#[cfg(feature = "alloc")]
pub type BoxedLocal<T> = Pin<Box<dyn Future<Output = T> + 'static>>;
pub trait FutureExt: Future {
fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Self::Output>
where
Self: Unpin,
{
Future::poll(Pin::new(self), cx)
}
fn or<F>(self, other: F) -> Or<Self, F>
where
Self: Sized,
F: Future<Output = Self::Output>,
{
Or {
future1: self,
future2: other,
}
}
#[cfg(all(feature = "std", feature = "race"))]
fn race<F>(self, other: F) -> Race<Self, F>
where
Self: Sized,
F: Future<Output = Self::Output>,
{
Race {
future1: self,
future2: other,
rng: Rng::new(),
}
}
#[cfg(feature = "std")]
fn catch_unwind(self) -> CatchUnwind<Self>
where
Self: Sized + UnwindSafe,
{
CatchUnwind { inner: self }
}
#[cfg(feature = "alloc")]
fn boxed<'a>(self) -> Pin<Box<dyn Future<Output = Self::Output> + Send + 'a>>
where
Self: Sized + Send + 'a,
{
Box::pin(self)
}
#[cfg(feature = "alloc")]
fn boxed_local<'a>(self) -> Pin<Box<dyn Future<Output = Self::Output> + 'a>>
where
Self: Sized + 'a,
{
Box::pin(self)
}
}
impl<F: Future + ?Sized> FutureExt for F {}