async_exec/
utils.rs

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

async_exec/
utils.rs
