futures_utils_lite/
zip_array.rs

1use core::{
2    pin::Pin,
3    task::{Context, Poll},
4};
5
6use alloc::{boxed::Box, vec::Vec};
7use futures_lite::{Future, FutureExt};
8use pin_project_lite::pin_project;
9
10pin_project! {
11    pub struct ZipArray<T> {
12        arr: Vec<(Pin<Box<dyn Future<Output = T> + Send>>, Option<T>)>,
13    }
14}
15
16impl<T> Future for ZipArray<T> {
17    type Output = Vec<T>;
18
19    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
20        let this = self.project();
21
22        let mut is_pending = true;
23
24        for fu in &mut this.arr.iter_mut() {
25            if fu.1.is_none() {
26                if let Poll::Ready(output) = fu.0.poll(cx) {
27                    fu.1 = Some(output);
28                } else {
29                    is_pending = true;
30                }
31            }
32        }
33
34        if is_pending {
35            Poll::Pending
36        } else {
37            let mut res = Vec::with_capacity(this.arr.len());
38
39            for fu in this.arr {
40                let opt = core::mem::take(&mut fu.1);
41
42                let output = opt.expect("Logic error!!!!!!");
43
44                res.push(output)
45            }
46
47            Poll::Ready(res)
48        }
49    }
50}
51
52pub fn zip_array<T, F>(fs: Vec<F>) -> ZipArray<T>
53where
54    F: Future<Output = T> + Send + 'static,
55{
56    let mut arr = Vec::with_capacity(fs.len());
57
58    for f in fs {
59        let fu: Pin<Box<dyn Future<Output = T> + Send>> = Box::pin(f);
60        arr.push((fu, None));
61    }
62
63    ZipArray { arr }
64}