From b33dc11689358ac575c7a4e72aebcd55b2407d6b Mon Sep 17 00:00:00 2001 From: "Alexis (Poliorcetics) Bourget" <alexis.bourget@gmail.com> Date: Tue, 15 Nov 2022 21:12:34 +0100 Subject: [PATCH] feat: Remove `MaybeDone` for `tuple::join` (without additional deps) --- src/future/join/tuple.rs | 93 ++++++++++++++++++++++++++++------------ 1 file changed, 66 insertions(+), 27 deletions(-) diff --git a/src/future/join/tuple.rs b/src/future/join/tuple.rs index 9c91660..7704302 100644 --- a/src/future/join/tuple.rs +++ b/src/future/join/tuple.rs @@ -1,16 +1,30 @@ use super::Join as JoinTrait; -use crate::utils::MaybeDone; +use crate::utils::PollState; use core::fmt::{self, Debug}; use core::future::{Future, IntoFuture}; +use core::mem::MaybeUninit; use core::pin::Pin; use core::task::{Context, Poll}; use pin_project::pin_project; macro_rules! impl_merge_tuple { - ($StructName:ident $($F:ident)*) => { - /// Waits for two similarly-typed futures to complete. + ($mod_name:ident :: $StructName:ident $($F:ident)*) => { + mod $mod_name { + #[allow(unused_imports)] // Len == 0 + use super::*; + + pub(super) struct Outputs<$($F: Future),*> { + $(pub(super) $F: MaybeUninit<$F::Output>),* + } + + pub(super) struct States { + $(pub(super) $F: PollState),* + } + } + + /// Waits for similarly-typed futures to complete. /// /// This `struct` is created by the [`join`] method on the [`Join`] trait. See /// its documentation for more. @@ -21,8 +35,10 @@ macro_rules! impl_merge_tuple { #[must_use = "futures do nothing unless you `.await` or poll them"] #[allow(non_snake_case)] pub struct $StructName<$($F: Future),*> { - done: bool, - $(#[pin] $F: MaybeDone<$F>,)* + completed: usize, + $(#[pin] $F: $F,)* + outputs: $mod_name::Outputs<$($F),*>, + states: $mod_name::States, } impl<$($F),*> Debug for $StructName<$($F),*> @@ -46,17 +62,38 @@ macro_rules! impl_merge_tuple { fn poll( self: Pin<&mut Self>, cx: &mut Context<'_> ) -> Poll<Self::Output> { - let mut all_done = true; + const LEN: usize = { + // This will disappear after compilation, leaving only the usize constant but + // we need it to indicate a type for the slice content (for empty slices). + const ARRAY: &[&str] = &[$(stringify!($F)),*]; + ARRAY.len() + }; + let mut this = self.project(); - assert!(!*this.done, "Futures must not be polled after completing"); - $(all_done &= this.$F.as_mut().poll(cx).is_ready();)* + assert_ne!(*this.completed, LEN + 1, "Futures must not be polled after completing"); - if all_done { - *this.done = true; - Poll::Ready(($(this.$F.take().unwrap()),*)) - } else { + $( + if this.states.$F.is_pending() { + if let Poll::Ready(value) = this.$F.poll(cx) { + this.outputs.$F.write(value); + this.states.$F = PollState::Done; + *this.completed += 1; + } + } + )* + + if *this.completed < LEN { Poll::Pending + } else { + *this.completed += 1; // That way the 0-len case can be awaited once without problems + Poll::Ready(( $({ + let mut out = MaybeUninit::uninit(); + core::mem::swap(&mut this.outputs.$F, &mut out); + // Safety: we have reached `this.completed == LEN` so all futures have + // been successfully polled and the output written to `this.output`. + unsafe { out.assume_init() } + }),* )) } } } @@ -72,27 +109,29 @@ macro_rules! impl_merge_tuple { fn join(self) -> Self::Future { let ($($F,)*): ($($F,)*) = self; $StructName { - done: false, - $($F: MaybeDone::new($F.into_future())),* + completed: 0, + $($F: $F.into_future(),)* + outputs: $mod_name::Outputs { $($F: MaybeUninit::uninit()),* }, + states: $mod_name::States { $($F: PollState::default()),* }, } } } }; } -impl_merge_tuple! { Join0 } -impl_merge_tuple! { Join1 A } -impl_merge_tuple! { Join2 A B } -impl_merge_tuple! { Join3 A B C } -impl_merge_tuple! { Join4 A B C D } -impl_merge_tuple! { Join5 A B C D E } -impl_merge_tuple! { Join6 A B C D E F } -impl_merge_tuple! { Join7 A B C D E F G } -impl_merge_tuple! { Join8 A B C D E F G H } -impl_merge_tuple! { Join9 A B C D E F G H I } -impl_merge_tuple! { Join10 A B C D E F G H I J } -impl_merge_tuple! { Join11 A B C D E F G H I J K } -impl_merge_tuple! { Join12 A B C D E F G H I J K L } +impl_merge_tuple! { join_0 :: Join0 } +impl_merge_tuple! { join_1 :: Join1 A } +impl_merge_tuple! { join_2 :: Join2 A B } +impl_merge_tuple! { join_3 :: Join3 A B C } +impl_merge_tuple! { join_4 :: Join4 A B C D } +impl_merge_tuple! { join_5 :: Join5 A B C D E } +impl_merge_tuple! { join_6 :: Join6 A B C D E F } +impl_merge_tuple! { join_7 :: Join7 A B C D E F G } +impl_merge_tuple! { join_8 :: Join8 A B C D E F G H } +impl_merge_tuple! { join_9 :: Join9 A B C D E F G H I } +impl_merge_tuple! { join_10 :: Join10 A B C D E F G H I J } +impl_merge_tuple! { join_11 :: Join11 A B C D E F G H I J K } +impl_merge_tuple! { join_12 :: Join12 A B C D E F G H I J K L } #[cfg(test)] mod test {