1use std::task::{Poll, Context};
4use std::pin::{Pin, pin};
5use std::future::Future;
6
7pub type Fut<'a, T> = &'a mut (dyn Future<Output = T> + Unpin);
9
10pub struct Race<'a, const N: usize, T> {
12 contenders: [Fut<'a, T>; N],
13}
14
15pub fn race<'a, const N: usize, T>(contenders: [Fut<'a, T>; N]) -> Race<'a, N, T> {
16 Race {
17 contenders,
18 }
19}
20
21impl<'a, const N: usize, T> Future for Race<'a, N, T> {
22 type Output = T;
23
24 fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<Self::Output> {
25 for contender in self.contenders.iter_mut() {
26 let pinned = pin!(contender);
27 if let Poll::Ready(ret) = pinned.poll(ctx) {
28 return Poll::Ready(ret);
29 }
30 }
31
32 Poll::Pending
33 }
34}
35
36pub enum Sel2<T1, T2> {
37 A(T1),
38 B(T2),
39}
40
41pub enum Sel3<T1, T2, T3> {
42 A(T1),
43 B(T2),
44 C(T3),
45}
46
47pub enum Sel4<T1, T2, T3, T4> {
48 A(T1),
49 B(T2),
50 C(T3),
51 D(T4),
52}
53
54pub async fn select2<T1, T2, A, B>(a: A, b: B) -> Sel2<T1, T2>
55where
56 A: Future<Output = T1>,
57 B: Future<Output = T2>,
58{
59 let mut wrapper_a = pin!(async { Sel2::A(a.await) });
60 let mut wrapper_b = pin!(async { Sel2::B(b.await) });
61 race([&mut wrapper_a, &mut wrapper_b]).await
62}
63
64pub async fn select3<T1, T2, T3, A, B, C>(a: A, b: B, c: C) -> Sel3<T1, T2, T3>
65where
66 A: Future<Output = T1>,
67 B: Future<Output = T2>,
68 C: Future<Output = T3>,
69{
70 let mut wrapper_a = pin!(async { Sel3::A(a.await) });
71 let mut wrapper_b = pin!(async { Sel3::B(b.await) });
72 let mut wrapper_c = pin!(async { Sel3::C(c.await) });
73 race([&mut wrapper_a, &mut wrapper_b, &mut wrapper_c]).await
74}
75
76pub async fn select4<T1, T2, T3, T4, A, B, C, D>(a: A, b: B, c: C, d: D) -> Sel4<T1, T2, T3, T4>
77where
78 A: Future<Output = T1>,
79 B: Future<Output = T2>,
80 C: Future<Output = T3>,
81 D: Future<Output = T4>,
82{
83 let mut wrapper_a = pin!(async { Sel4::A(a.await) });
84 let mut wrapper_b = pin!(async { Sel4::B(b.await) });
85 let mut wrapper_c = pin!(async { Sel4::C(c.await) });
86 let mut wrapper_d = pin!(async { Sel4::D(d.await) });
87 race([&mut wrapper_a, &mut wrapper_b, &mut wrapper_c, &mut wrapper_d]).await
88}
89
90#[test]
91fn test_select() {
92 let task = async {
93 let task_1 = ();
94 let task_2 = ();
95
96 let opt_1 = sel_opt(task_1);
97 let opt_2 = sel_opt(task_2);
98
99 let ret = match select(task_1, task_2) {
100 Sel::A(ret_1) => self.handle_1(ret_1),
101 Sel::B(ret_2) => self.handle_2(ret_2),
102 };
103 };
104}