WIP

src/future/join/vec.rs

@@ -1,5 +1,5 @@
 use super::Join as JoinTrait;
-use crate::utils::{iter_pin_mut_vec, PollVec};
+use crate::utils::{iter_pin_mut_vec, PollVec, WakerVec};
 
 use core::fmt;
 use core::future::{Future, IntoFuture};
@@ -26,6 +26,7 @@ where
     consumed: bool,
     pending: usize,
     items: Vec<MaybeUninit<<Fut as Future>::Output>>,
+    wakers: WakerVec,
     state: PollVec,
     #[pin]
     futures: Vec<Fut>,
@@ -36,13 +37,15 @@ where
     Fut: Future,
 {
     pub(crate) fn new(futures: Vec<Fut>) -> Self {
+        let len = futures.len();
         Join {
             consumed: false,
-            pending: futures.len(),
+            pending: len,
             items: std::iter::repeat_with(MaybeUninit::uninit)
-                .take(futures.len())
+                .take(len)
                 .collect(),
-            state: PollVec::new(futures.len()),
+            wakers: WakerVec::new(len),
+            state: PollVec::new(len),
             futures,
         }
     }
@@ -84,12 +87,26 @@ where
             "Futures must not be polled after completing"
         );
 
+        let mut readiness = this.wakers.readiness().lock().unwrap();
+        readiness.set_waker(cx.waker());
+
+        if !readiness.any_ready() {
+            // Nothing is ready yet
+            return Poll::Pending;
+        }
+
+        // unlock readiness so we don't deadlock when polling
+        drop(readiness);
+
         // Poll all futures
         let futures = this.futures.as_mut();
         let states = &mut this.state[..];
         for (i, fut) in iter_pin_mut_vec(futures).enumerate() {
-            if states[i].is_pending() {
-                if let Poll::Ready(value) = fut.poll(cx) {
+            if states[i].is_pending() && this.wakers.readiness().lock().unwrap().clear_ready(i) {
+                // Obtain the intermediate waker.
+                let mut cx = Context::from_waker(this.wakers.get(i).unwrap());
+
+                if let Poll::Ready(value) = fut.poll(&mut cx) {
                     this.items[i] = MaybeUninit::new(value);
                     states[i].set_ready();
                     *this.pending -= 1;