wtx/misc/
join_array.rs

1use core::{
2  pin::Pin,
3  task::{Context, Poll},
4};
5
6/// Joins the result of an array of futures, waiting for them all to complete.
7///
8/// You should `Box` this structure if size is a concern.
9#[must_use = "Futures do nothing unless you await them"]
10#[derive(Debug)]
11pub struct JoinArray<F, const N: usize>
12where
13  F: Future,
14{
15  futures: [F; N],
16  outputs: Option<[Option<F::Output>; N]>,
17}
18
19impl<F, const N: usize> JoinArray<F, N>
20where
21  F: Future,
22{
23  /// Creates a new instance
24  #[inline]
25  pub const fn new(futures: [F; N]) -> Self {
26    Self { futures, outputs: Some([const { None }; N]) }
27  }
28}
29
30impl<F: Future, const N: usize> Future for JoinArray<F, N>
31where
32  F: Future,
33{
34  type Output = [F::Output; N];
35
36  #[inline]
37  fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
38    // SAFETY: No fields are moved
39    let JoinArray { futures, outputs } = unsafe { self.get_unchecked_mut() };
40    let Some(outputs_mut) = outputs else {
41      #[expect(
42        clippy::panic,
43        reason = "Compiler will probably remove this branch as long as the user don't mess up"
44      )]
45      {
46        panic!("Can't poll `JoinArray` again after completion");
47      }
48    };
49    let mut is_finished = true;
50    for (future, result) in futures.iter_mut().zip(outputs_mut) {
51      if result.is_some() {
52        continue;
53      }
54      // SAFETY: No future is moved
55      let pinned = unsafe { Pin::new_unchecked(future) };
56      if let Poll::Ready(output) = pinned.poll(cx) {
57        *result = Some(output);
58      } else {
59        is_finished = false;
60      }
61    }
62    if is_finished && let Some(array) = outputs.take() {
63      #[expect(clippy::unwrap_used, reason = "Compiler removes this branch")]
64      return Poll::Ready(array.map(|el| el.unwrap()));
65    }
66    Poll::Pending
67  }
68}
69
70#[cfg(test)]
71mod tests {
72  use crate::misc::JoinArray;
73  use core::future::ready;
74
75  #[wtx::test]
76  async fn polls_array() {
77    assert_eq!(JoinArray::new([ready(1), ready(2)]).await, [1, 2]);
78  }
79}
wtx/misc/join_array.rs

wtx/misc/
join_array.rs
