future_iter/
join_set.rs

1use std::future::Future;
2use tokio::task::JoinSet;
3
4pub mod prelude {
5  pub use super::{IntoJoinSet as _, IntoJoinSetBy as _, JoinSetFromIter as _};
6}
7
8pub trait JoinSetFromIter: Iterator {
9  fn join_set<T>(self) -> JoinSet<T>
10  where
11    Self: Sized,
12    Self::Item: Future<Output = T> + Send + 'static,
13    T: Send + 'static,
14  {
15    self.collect()
16  }
17
18  fn join_set_by<T, F, M>(self, f: M) -> JoinSet<T>
19  where
20    Self: Sized,
21    Self::Item: Send + 'static,
22    F: Future<Output = T> + Send + 'static,
23    T: Send + 'static,
24    M: FnMut(Self::Item) -> F,
25  {
26    self.map(f).join_set()
27  }
28}
29
30impl<T> JoinSetFromIter for T where T: Iterator + ?Sized {}
31
32pub trait IntoJoinSet<F, T>: IntoIterator
33where
34  Self: Sized,
35  <Self as IntoIterator>::Item: Future<Output = T> + Send + 'static,
36  F: Future<Output = T> + Send + 'static,
37  T: Send + 'static,
38{
39  fn into_join_set(self) -> JoinSet<T> {
40    self.into_iter().join_set()
41  }
42}
43
44impl<F, T> IntoJoinSet<F, T> for Vec<F>
45where
46  F: Future<Output = T> + Send + 'static,
47  T: Send + 'static,
48{
49}
50
51pub trait IntoJoinSetBy<F, T>: IntoIterator
52where
53  Self: Sized,
54  <Self as IntoIterator>::Item: Send + 'static,
55  F: Future<Output = T> + Send + 'static,
56  T: Send + 'static,
57{
58  fn into_join_set_by<M>(self, f: M) -> JoinSet<T>
59  where
60    M: FnMut(Self::Item) -> F,
61  {
62    self.into_iter().join_set_by(f)
63  }
64}
65
66impl<F, T, U> IntoJoinSetBy<F, T> for Vec<U>
67where
68  F: Future<Output = T> + Send + 'static,
69  T: Send + 'static,
70  U: Send + 'static,
71{
72}
73
74#[cfg(test)]
75mod tests {
76  use super::*;
77  use itertools::Itertools;
78  use std::future;
79
80  #[tokio::test]
81  async fn join_set_by() {
82    let mut set = (0..10).into_iter().join_set_by(future::ready);
83
84    assert!(set.len() == 10);
85
86    while let Some(result) = set.join_next().await {
87      result.unwrap();
88    }
89  }
90
91  #[tokio::test]
92  async fn into_join_set() {
93    let mut set = (0..10)
94      .into_iter()
95      .map(future::ready)
96      .collect_vec()
97      .into_join_set();
98
99    assert!(set.len() == 10);
100
101    while let Some(result) = set.join_next().await {
102      result.unwrap();
103    }
104  }
105
106  #[tokio::test]
107  async fn into_join_set_by() {
108    let mut set = (0..10)
109      .into_iter()
110      .collect_vec()
111      .into_join_set_by(future::ready);
112
113    assert!(set.len() == 10);
114
115    while let Some(result) = set.join_next().await {
116      result.unwrap();
117    }
118  }
119
120  #[tokio::test]
121  async fn into_join_set_by_with_different_type() {
122    let mut set = (0..10)
123      .into_iter()
124      .collect_vec()
125      .into_join_set_by(|it| future::ready(format!("{it}")));
126
127    assert!(set.len() == 10);
128
129    while let Some(result) = set.join_next().await {
130      result.unwrap();
131    }
132  }
133}