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 {