From d761e84968de093ac2e0850293af1057b6fb44ee Mon Sep 17 00:00:00 2001 From: Ibraheem Ahmed Date: Tue, 7 Dec 2021 20:59:16 -0500 Subject: [PATCH 1/5] implement `core::future::join` --- library/core/src/future/join.rs | 147 ++++++++++++++++++++++++++++++++ library/core/src/future/mod.rs | 4 + 2 files changed, 151 insertions(+) create mode 100644 library/core/src/future/join.rs diff --git a/library/core/src/future/join.rs b/library/core/src/future/join.rs new file mode 100644 index 0000000000000..752a3ea92ba41 --- /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::pin::Pin; +use crate::task::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()); +/// 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()); +/// 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)* }, + // The future currently being expanded, and the rest + @rest: ($current:expr, $($rest:tt)*) + ) => { + join! { + @count: ($($count)* _), // Add to the count + @futures: { $($fut)* $current => ($($count)*), }, // Add the future from @rest with it's position + @rest: ($($rest)*) // And leave the rest + } + }, + // Now generate the output future + ( + @count: ($($count:tt)*), + @futures: { + $( $fut:expr => ( $($pos:tt)* ), )* + }, + @rest: () + ) => {{ + let mut futures = ( $( MaybeDone::Future($fut), )* ); + + poll_fn(move |cx| { + let mut done = true; + + $( + // Extract the future from the tuple + let ( $($pos,)* fut, .. ) = &mut futures; + + // SAFETY: the futures are never moved + done &= unsafe { Pin::new_unchecked(fut).poll(cx).is_ready() }; + )* + + if done { + Poll::Ready(($({ + let ( $($pos,)* fut, .. ) = &mut futures; + + // SAFETY: the futures are never moved + unsafe { Pin::new_unchecked(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. +#[allow(dead_code)] +#[unstable(feature = "future_join", issue = "none")] +enum MaybeDone { + Future(F), + Done(F::Output), + Took, +} + +#[unstable(feature = "future_join", issue = "none")] +impl Unpin for MaybeDone {} + +#[unstable(feature = "future_join", issue = "none")] +impl MaybeDone { + #[allow(dead_code)] + fn take_output(self: Pin<&mut Self>) -> Option { + unsafe { + match &*self { + MaybeDone::Done(_) => match mem::replace(self.get_unchecked_mut(), Self::Took) { + MaybeDone::Done(val) => Some(val), + _ => unreachable!(), + }, + _ => None, + } + } + } +} + +#[unstable(feature = "future_join", issue = "none")] +impl Future for MaybeDone { + type Output = (); + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + 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; From 08dca1933b68456486e9d2f18aa3409cce49c37c Mon Sep 17 00:00:00 2001 From: Ibraheem Ahmed Date: Tue, 7 Dec 2021 21:01:10 -0500 Subject: [PATCH 2/5] generate `MaybeDone` futures inline `join` --- library/core/src/future/join.rs | 110 ++++++++++++++------------------ 1 file changed, 47 insertions(+), 63 deletions(-) diff --git a/library/core/src/future/join.rs b/library/core/src/future/join.rs index 752a3ea92ba41..03d106c969bd8 100644 --- a/library/core/src/future/join.rs +++ b/library/core/src/future/join.rs @@ -4,6 +4,7 @@ use crate::cell::UnsafeCell; use crate::future::{poll_fn, Future}; use crate::pin::Pin; use crate::task::Poll; + /// Polls multiple futures simultaneously, returning a tuple /// of all results once complete. /// @@ -53,42 +54,74 @@ pub macro join { @count: ($($count:tt)*), // Futures and their positions in the tuple: "{ a => (_), b => (_ _)) }" @futures: { $($fut:tt)* }, - // The future currently being expanded, and the rest + // Take a future from @rest to expand @rest: ($current:expr, $($rest:tt)*) ) => { join! { - @count: ($($count)* _), // Add to the count - @futures: { $($fut)* $current => ($($count)*), }, // Add the future from @rest with it's position - @rest: ($($rest)*) // And leave the rest + @count: ($($count)* _), + @futures: { $($fut)* $current => ($($count)*), }, + @rest: ($($rest)*) } }, // Now generate the output future ( @count: ($($count:tt)*), @futures: { - $( $fut:expr => ( $($pos:tt)* ), )* + $( $(@$f:tt)? $fut:expr => ( $($pos:tt)* ), )* }, @rest: () ) => {{ - let mut futures = ( $( MaybeDone::Future($fut), )* ); + // The futures and whether they have completed + let mut state = ( $( UnsafeCell::new(($fut, false)), )* ); + + // Make sure the futures don't panic + // if polled after completion, and + // store their output separately + let mut futures = ($( + ({ + let ( $($pos,)* state, .. ) = &state; + + poll_fn(move |cx| { + // SAFETY: each future borrows a distinct element + // of the tuple + let (fut, done) = unsafe { &mut *state.get() }; + + if *done { + return Poll::Ready(None) + } + + // SAFETY: The futures are never moved + match unsafe { Pin::new_unchecked(fut).poll(cx) } { + Poll::Ready(val) => { + *done = true; + Poll::Ready(Some(val)) + } + Poll::Pending => Poll::Pending + } + }) + }, None), + )*); poll_fn(move |cx| { let mut done = true; $( - // Extract the future from the tuple - let ( $($pos,)* fut, .. ) = &mut futures; + let ( $($pos,)* (fut, out), .. ) = &mut futures; - // SAFETY: the futures are never moved - done &= unsafe { Pin::new_unchecked(fut).poll(cx).is_ready() }; + // SAFETY: The futures are never moved + match unsafe { Pin::new_unchecked(fut).poll(cx) } { + Poll::Ready(Some(val)) => *out = Some(val), + // the future was already done + Poll::Ready(None) => {}, + Poll::Pending => done = false, + } )* if done { + // Extract all the outputs Poll::Ready(($({ - let ( $($pos,)* fut, .. ) = &mut futures; - - // SAFETY: the futures are never moved - unsafe { Pin::new_unchecked(fut).take_output().unwrap() } + let ( $($pos,)* (_, val), .. ) = &mut futures; + val.unwrap() }),*)) } else { Poll::Pending @@ -96,52 +129,3 @@ pub macro join { }).await }} } - -/// Future used by `join!` that stores it's output to -/// be later taken and doesn't panic when polled after ready. -#[allow(dead_code)] -#[unstable(feature = "future_join", issue = "none")] -enum MaybeDone { - Future(F), - Done(F::Output), - Took, -} - -#[unstable(feature = "future_join", issue = "none")] -impl Unpin for MaybeDone {} - -#[unstable(feature = "future_join", issue = "none")] -impl MaybeDone { - #[allow(dead_code)] - fn take_output(self: Pin<&mut Self>) -> Option { - unsafe { - match &*self { - MaybeDone::Done(_) => match mem::replace(self.get_unchecked_mut(), Self::Took) { - MaybeDone::Done(val) => Some(val), - _ => unreachable!(), - }, - _ => None, - } - } - } -} - -#[unstable(feature = "future_join", issue = "none")] -impl Future for MaybeDone { - type Output = (); - - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - 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(()) - } -} From d07cef22b0397eee2649807cd1dddf6d983e215f Mon Sep 17 00:00:00 2001 From: Ibraheem Ahmed Date: Tue, 7 Dec 2021 21:10:39 -0500 Subject: [PATCH 3/5] add tests for `core::future::join` --- library/core/tests/future.rs | 76 ++++++++++++++++++++++++++++++++++++ library/core/tests/lib.rs | 3 ++ 2 files changed, 79 insertions(+) create mode 100644 library/core/tests/future.rs diff --git a/library/core/tests/future.rs b/library/core/tests/future.rs new file mode 100644 index 0000000000000..f47dcc70434cf --- /dev/null +++ b/library/core/tests/future.rs @@ -0,0 +1,76 @@ +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 }); + assert_eq!(x, 0); + + let x = join!(async { 0 }, async { 1 }); + assert_eq!(x, (0, 1)); + + let x = join!(async { 0 }, async { 1 }, async { 2 }); + 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) + ); + assert_eq!(x, (0, 1, 2, 3, 4, 5, 6, 7)); + }); +} + +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; From a8c931410020840584a2efa5f77239a9c5fcb85c Mon Sep 17 00:00:00 2001 From: Ibraheem Ahmed Date: Wed, 8 Dec 2021 16:44:48 -0500 Subject: [PATCH 4/5] remove implicit .await from `core::future::join` --- library/core/src/future/join.rs | 100 ++++++++++++++++---------------- library/core/tests/future.rs | 17 ++++-- 2 files changed, 64 insertions(+), 53 deletions(-) diff --git a/library/core/src/future/join.rs b/library/core/src/future/join.rs index 03d106c969bd8..bed9f3dd51cc5 100644 --- a/library/core/src/future/join.rs +++ b/library/core/src/future/join.rs @@ -22,7 +22,7 @@ use crate::task::Poll; /// async fn two() -> usize { 2 } /// /// # let _ = async { -/// let x = join!(one(), two()); +/// let x = join!(one(), two()).await; /// assert_eq!(x, (1, 2)); /// # }; /// ``` @@ -39,7 +39,7 @@ use crate::task::Poll; /// async fn three() -> usize { 3 } /// /// # let _ = async { -/// let x = join!(one(), two(), three()); +/// let x = join!(one(), two(), three()).await; /// assert_eq!(x, (1, 2, 3)); /// # }; /// ``` @@ -71,61 +71,63 @@ pub macro join { }, @rest: () ) => {{ - // The futures and whether they have completed - let mut state = ( $( UnsafeCell::new(($fut, false)), )* ); + async move { + // The futures and whether they have completed + let mut state = ( $( UnsafeCell::new(($fut, false)), )* ); - // Make sure the futures don't panic - // if polled after completion, and - // store their output separately - let mut futures = ($( - ({ - let ( $($pos,)* state, .. ) = &state; + // Make sure the futures don't panic + // if polled after completion, and + // store their output separately + let mut futures = ($( + ({ + let ( $($pos,)* state, .. ) = &state; - poll_fn(move |cx| { - // SAFETY: each future borrows a distinct element - // of the tuple - let (fut, done) = unsafe { &mut *state.get() }; + poll_fn(move |cx| { + // SAFETY: each future borrows a distinct element + // of the tuple + let (fut, done) = unsafe { &mut *state.get() }; - if *done { - return Poll::Ready(None) - } + if *done { + return Poll::Ready(None) + } - // SAFETY: The futures are never moved - match unsafe { Pin::new_unchecked(fut).poll(cx) } { - Poll::Ready(val) => { - *done = true; - Poll::Ready(Some(val)) + // SAFETY: The futures are never moved + match unsafe { Pin::new_unchecked(fut).poll(cx) } { + Poll::Ready(val) => { + *done = true; + Poll::Ready(Some(val)) + } + Poll::Pending => Poll::Pending } - Poll::Pending => Poll::Pending - } - }) - }, None), - )*); + }) + }, None), + )*); - poll_fn(move |cx| { - let mut done = true; + poll_fn(move |cx| { + let mut done = true; - $( - let ( $($pos,)* (fut, out), .. ) = &mut futures; + $( + let ( $($pos,)* (fut, out), .. ) = &mut futures; - // SAFETY: The futures are never moved - match unsafe { Pin::new_unchecked(fut).poll(cx) } { - Poll::Ready(Some(val)) => *out = Some(val), - // the future was already done - Poll::Ready(None) => {}, - Poll::Pending => done = false, - } - )* + // SAFETY: The futures are never moved + match unsafe { Pin::new_unchecked(fut).poll(cx) } { + Poll::Ready(Some(val)) => *out = Some(val), + // the future was already done + Poll::Ready(None) => {}, + Poll::Pending => done = false, + } + )* - if done { - // Extract all the outputs - Poll::Ready(($({ - let ( $($pos,)* (_, val), .. ) = &mut futures; - val.unwrap() - }),*)) - } else { - Poll::Pending - } - }).await + if done { + // Extract all the outputs + Poll::Ready(($({ + let ( $($pos,)* (_, val), .. ) = &mut futures; + val.unwrap() + }),*)) + } else { + Poll::Pending + } + }).await + } }} } diff --git a/library/core/tests/future.rs b/library/core/tests/future.rs index f47dcc70434cf..73249b1b8a435 100644 --- a/library/core/tests/future.rs +++ b/library/core/tests/future.rs @@ -32,13 +32,13 @@ fn poll_n(val: usize, num: usize) -> PollN { #[test] fn test_join() { block_on(async move { - let x = join!(async { 0 }); + let x = join!(async { 0 }).await; assert_eq!(x, 0); - let x = join!(async { 0 }, async { 1 }); + let x = join!(async { 0 }, async { 1 }).await; assert_eq!(x, (0, 1)); - let x = join!(async { 0 }, async { 1 }, async { 2 }); + let x = join!(async { 0 }, async { 1 }, async { 2 }).await; assert_eq!(x, (0, 1, 2)); let x = join!( @@ -50,8 +50,17 @@ fn test_join() { 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); }); } From 5478f439e13b55c7b9b858f76986786371b97c8f Mon Sep 17 00:00:00 2001 From: Ibraheem Ahmed Date: Wed, 8 Dec 2021 17:08:23 -0500 Subject: [PATCH 5/5] trim down expansion of `core::future::join` --- library/core/src/future/join.rs | 98 +++++++++++++++++++-------------- 1 file changed, 56 insertions(+), 42 deletions(-) diff --git a/library/core/src/future/join.rs b/library/core/src/future/join.rs index bed9f3dd51cc5..aadff103ebab4 100644 --- a/library/core/src/future/join.rs +++ b/library/core/src/future/join.rs @@ -2,8 +2,9 @@ use crate::cell::UnsafeCell; use crate::future::{poll_fn, Future}; +use crate::mem; use crate::pin::Pin; -use crate::task::Poll; +use crate::task::{Context, Poll}; /// Polls multiple futures simultaneously, returning a tuple /// of all results once complete. @@ -70,64 +71,77 @@ pub macro join { $( $(@$f:tt)? $fut:expr => ( $($pos:tt)* ), )* }, @rest: () - ) => {{ + ) => { async move { - // The futures and whether they have completed - let mut state = ( $( UnsafeCell::new(($fut, false)), )* ); - - // Make sure the futures don't panic - // if polled after completion, and - // store their output separately - let mut futures = ($( - ({ - let ( $($pos,)* state, .. ) = &state; - - poll_fn(move |cx| { - // SAFETY: each future borrows a distinct element - // of the tuple - let (fut, done) = unsafe { &mut *state.get() }; - - if *done { - return Poll::Ready(None) - } - - // SAFETY: The futures are never moved - match unsafe { Pin::new_unchecked(fut).poll(cx) } { - Poll::Ready(val) => { - *done = true; - Poll::Ready(Some(val)) - } - Poll::Pending => Poll::Pending - } - }) - }, None), - )*); + let mut futures = ( $( MaybeDone::Future($fut), )* ); poll_fn(move |cx| { let mut done = true; $( - let ( $($pos,)* (fut, out), .. ) = &mut futures; + let ( $($pos,)* fut, .. ) = &mut futures; // SAFETY: The futures are never moved - match unsafe { Pin::new_unchecked(fut).poll(cx) } { - Poll::Ready(Some(val)) => *out = Some(val), - // the future was already done - Poll::Ready(None) => {}, - Poll::Pending => done = false, - } + done &= unsafe { Pin::new_unchecked(fut).poll(cx).is_ready() }; )* if done { // Extract all the outputs Poll::Ready(($({ - let ( $($pos,)* (_, val), .. ) = &mut futures; - val.unwrap() + 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(()) + } }