1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59

use std::future::Future;
use tokio::task::JoinSet;

pub trait IntoJoinSet: Iterator {
  fn into_join_set<T>(self) -> JoinSet<T>
  where
    Self: Sized,
    Self::Item: Future<Output = T> + Send + 'static,
    T: Send + 'static,
  {
    self.collect()
  }

  fn into_join_set_by<T, F, M>(self, f: M) -> JoinSet<T>
  where
    Self: Sized,
    Self::Item: Send + 'static,
    F: Future<Output = T> + Send + 'static,
    T: Send + 'static,
    M: FnMut(Self::Item) -> F,
  {
    self.map(f).into_join_set()
  }
}

impl<T> IntoJoinSet for T where T: Iterator + ?Sized {}

#[cfg(test)]
mod tests {
  use super::*;
  use std::future;

  #[tokio::test]
  async fn into_join_set() {
    let mut set = (0..10)
      .into_iter()
      .map(future::ready)
      .into_join_set();

    assert!(set.len() == 10);

    while let Some(result) = set.join_next().await {
      result.unwrap();
    }
  }

  #[tokio::test]
  async fn into_join_set_by() {
    let mut set = (0..10)
      .into_iter()
      .into_join_set_by(future::ready);

    assert!(set.len() == 10);

    while let Some(result) = set.join_next().await {
      result.unwrap();
    }
  }
}