diff --git a/library/core/src/future/join.rs b/library/core/src/future/join.rs new file mode 100644 index 0000000000000..aadff103ebab4 --- /dev/null +++ b/library/core/src/future/join.rs @@ -0,0 +1,147 @@ +#![allow(unused_imports)] // items are used by the macro + +use crate::cell::UnsafeCell; +use crate::future::{poll_fn, Future}; +use crate::mem; +use crate::pin::Pin; +use crate::task::{Context, Poll}; + +/// Polls multiple futures simultaneously, returning a tuple +/// of all results once complete. +/// +/// While `join!(a, b)` is similar to `(a.await, b.await)`, +/// `join!` polls both futures concurrently and is therefore more efficient. +/// +/// # Examples +/// +/// ``` +/// #![feature(future_join, future_poll_fn)] +/// +/// use std::future::join; +/// +/// async fn one() -> usize { 1 } +/// async fn two() -> usize { 2 } +/// +/// # let _ = async { +/// let x = join!(one(), two()).await; +/// assert_eq!(x, (1, 2)); +/// # }; +/// ``` +/// +/// `join!` is variadic, so you can pass any number of futures: +/// +/// ``` +/// #![feature(future_join, future_poll_fn)] +/// +/// use std::future::join; +/// +/// async fn one() -> usize { 1 } +/// async fn two() -> usize { 2 } +/// async fn three() -> usize { 3 } +/// +/// # let _ = async { +/// let x = join!(one(), two(), three()).await; +/// assert_eq!(x, (1, 2, 3)); +/// # }; +/// ``` +#[unstable(feature = "future_join", issue = "91642")] +pub macro join { + ( $($fut:expr),* $(,)?) => { + join! { @count: (), @futures: {}, @rest: ($($fut,)*) } + }, + // Recurse until we have the position of each future in the tuple + ( + // A token for each future that has been expanded: "_ _ _" + @count: ($($count:tt)*), + // Futures and their positions in the tuple: "{ a => (_), b => (_ _)) }" + @futures: { $($fut:tt)* }, + // Take a future from @rest to expand + @rest: ($current:expr, $($rest:tt)*) + ) => { + join! { + @count: ($($count)* _), + @futures: { $($fut)* $current => ($($count)*), }, + @rest: ($($rest)*) + } + }, + // Now generate the output future + ( + @count: ($($count:tt)*), + @futures: { + $( $(@$f:tt)? $fut:expr => ( $($pos:tt)* ), )* + }, + @rest: () + ) => { + async move { + let mut futures = ( $( MaybeDone::Future($fut), )* ); + + poll_fn(move |cx| { + let mut done = true; + + $( + let ( $($pos,)* fut, .. ) = &mut futures; + + // SAFETY: The futures are never moved + done &= unsafe { Pin::new_unchecked(fut).poll(cx).is_ready() }; + )* + + if done { + // Extract all the outputs + Poll::Ready(($({ + let ( $($pos,)* fut, .. ) = &mut futures; + + fut.take_output().unwrap() + }),*)) + } else { + Poll::Pending + } + }).await + } + } +} + +/// Future used by `join!` that stores it's output to +/// be later taken and doesn't panic when polled after ready. +/// +/// This type is public in a private module for use by the macro. +#[allow(missing_debug_implementations)] +#[unstable(feature = "future_join", issue = "91642")] +pub enum MaybeDone { + Future(F), + Done(F::Output), + Took, +} + +#[unstable(feature = "future_join", issue = "91642")] +impl MaybeDone { + pub fn take_output(&mut self) -> Option { + match &*self { + MaybeDone::Done(_) => match mem::replace(self, Self::Took) { + MaybeDone::Done(val) => Some(val), + _ => unreachable!(), + }, + _ => None, + } + } +} + +#[unstable(feature = "future_join", issue = "91642")] +impl Future for MaybeDone { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + // SAFETY: pinning in structural for `f` + unsafe { + match self.as_mut().get_unchecked_mut() { + MaybeDone::Future(f) => match Pin::new_unchecked(f).poll(cx) { + Poll::Ready(val) => self.set(Self::Done(val)), + Poll::Pending => return Poll::Pending, + }, + MaybeDone::Done(_) => {} + MaybeDone::Took => unreachable!(), + } + } + + Poll::Ready(()) + } +} diff --git a/library/core/src/future/mod.rs b/library/core/src/future/mod.rs index 7a3af70d6d97c..88db584aefd08 100644 --- a/library/core/src/future/mod.rs +++ b/library/core/src/future/mod.rs @@ -11,6 +11,7 @@ use crate::{ mod future; mod into_future; +mod join; mod pending; mod poll_fn; mod ready; @@ -18,6 +19,9 @@ mod ready; #[stable(feature = "futures_api", since = "1.36.0")] pub use self::future::Future; +#[unstable(feature = "future_join", issue = "91642")] +pub use self::join::join; + #[unstable(feature = "into_future", issue = "67644")] pub use into_future::IntoFuture; diff --git a/library/core/tests/future.rs b/library/core/tests/future.rs new file mode 100644 index 0000000000000..73249b1b8a435 --- /dev/null +++ b/library/core/tests/future.rs @@ -0,0 +1,85 @@ +use std::future::{join, Future}; +use std::pin::Pin; +use std::sync::Arc; +use std::task::{Context, Poll, Wake}; +use std::thread; + +struct PollN { + val: usize, + polled: usize, + num: usize, +} + +impl Future for PollN { + type Output = usize; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + self.polled += 1; + + if self.polled == self.num { + return Poll::Ready(self.val); + } + + cx.waker().wake_by_ref(); + Poll::Pending + } +} + +fn poll_n(val: usize, num: usize) -> PollN { + PollN { val, num, polled: 0 } +} + +#[test] +fn test_join() { + block_on(async move { + let x = join!(async { 0 }).await; + assert_eq!(x, 0); + + let x = join!(async { 0 }, async { 1 }).await; + assert_eq!(x, (0, 1)); + + let x = join!(async { 0 }, async { 1 }, async { 2 }).await; + assert_eq!(x, (0, 1, 2)); + + let x = join!( + poll_n(0, 1), + poll_n(1, 5), + poll_n(2, 2), + poll_n(3, 1), + poll_n(4, 2), + poll_n(5, 3), + poll_n(6, 4), + poll_n(7, 1) + ) + .await; + assert_eq!(x, (0, 1, 2, 3, 4, 5, 6, 7)); + + let y = String::new(); + let x = join!(async { + println!("{}", &y); + 1 + }) + .await; + assert_eq!(x, 1); + }); +} + +fn block_on(fut: impl Future) { + struct Waker; + impl Wake for Waker { + fn wake(self: Arc) { + thread::current().unpark() + } + } + + let waker = Arc::new(Waker).into(); + let mut cx = Context::from_waker(&waker); + let mut fut = Box::pin(fut); + + loop { + match fut.as_mut().poll(&mut cx) { + Poll::Ready(_) => break, + Poll::Pending => thread::park(), + } + } +} diff --git a/library/core/tests/lib.rs b/library/core/tests/lib.rs index 9ab98ba88865a..73a3a1fc3e0af 100644 --- a/library/core/tests/lib.rs +++ b/library/core/tests/lib.rs @@ -29,6 +29,8 @@ #![feature(flt2dec)] #![feature(fmt_internals)] #![feature(float_minimum_maximum)] +#![feature(future_join)] +#![feature(future_poll_fn)] #![feature(array_from_fn)] #![feature(hashmap_internals)] #![feature(try_find)] @@ -94,6 +96,7 @@ mod clone; mod cmp; mod const_ptr; mod fmt; +mod future; mod hash; mod intrinsics; mod iter;