From 5478f439e13b55c7b9b858f76986786371b97c8f Mon Sep 17 00:00:00 2001
From: Ibraheem Ahmed <ibrah1440@gmail.com>
Date: Wed, 8 Dec 2021 17:08:23 -0500
Subject: [PATCH] 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 bed9f3dd51c..aadff103eba 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<F: Future> {
+    Future(F),
+    Done(F::Output),
+    Took,
+}
+
+#[unstable(feature = "future_join", issue = "91642")]
+impl<F: Future> MaybeDone<F> {
+    pub fn take_output(&mut self) -> Option<F::Output> {
+        match &*self {
+            MaybeDone::Done(_) => match mem::replace(self, Self::Took) {
+                MaybeDone::Done(val) => Some(val),
+                _ => unreachable!(),
+            },
+            _ => None,
+        }
+    }
+}
+
+#[unstable(feature = "future_join", issue = "91642")]
+impl<F: Future> Future for MaybeDone<F> {
+    type Output = ();
+
+    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        // 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(())
+    }
 }