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}