diff --git a/src/future/future.rs b/src/future/future.rs index fe68517..1b88430 100644 --- a/src/future/future.rs +++ b/src/future/future.rs @@ -1,9 +1,12 @@ cfg_unstable! { mod delay; + mod flatten; use std::time::Duration; use delay::DelayFuture; + use flatten::FlattenFuture; + use crate::future::IntoFuture; } extension_trait! { @@ -129,6 +132,30 @@ extension_trait! { { DelayFuture::new(self, dur) } + + /// Flatten out the execution of this future when the result itself + /// can be converted into another future. + /// + /// # Examples + /// + /// ``` + /// # async_std::task::block_on(async { + /// use async_std::prelude::*; + /// + /// let nested_future = async { async { 1 } }; + /// let future = nested_future.flatten(); + /// assert_eq!(future.await, 1); + /// # }) + /// ``` + #[cfg_attr(feature = "docs", doc(cfg(unstable)))] + #[cfg(any(feature = "unstable", feature = "docs"))] + fn flatten(self) -> FlattenFuture::Future> + where + Self: Future + Sized, + Self::Output: IntoFuture + { + FlattenFuture::new(self) + } } impl Future for Box { diff --git a/src/future/future/flatten.rs b/src/future/future/flatten.rs new file mode 100644 index 0000000..27ee228 --- /dev/null +++ b/src/future/future/flatten.rs @@ -0,0 +1,46 @@ +use futures_core::ready; +use std::pin::Pin; + +use crate::future::Future; +use crate::future::IntoFuture; +use crate::task::{Context, Poll}; + +#[doc(hidden)] +#[derive(Debug)] +pub enum FlattenFuture { + First(Fut1), + Second(Fut2), + Empty, +} + +impl FlattenFuture { + pub fn new(future: Fut1) -> FlattenFuture { + FlattenFuture::First(future) + } +} + +impl Future for FlattenFuture::Future> +where + Fut1: Future, + Fut1::Output: IntoFuture, +{ + type Output = ::Output; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let this = unsafe { self.get_unchecked_mut() }; + loop { + match this { + FlattenFuture::First(fut1) => { + let fut2 = ready!(unsafe { Pin::new_unchecked(fut1) }.poll(cx)).into_future(); + *this = FlattenFuture::Second(fut2); + } + FlattenFuture::Second(fut2) => { + let v = ready!(unsafe { Pin::new_unchecked(fut2) }.poll(cx)); + *this = FlattenFuture::Empty; + return Poll::Ready(v); + } + FlattenFuture::Empty => unreachable!(), + } + } + } +}