async_exec/
utils.rs

1//! Future Utilities
2
3use std::task::{Poll, Context};
4use std::pin::{Pin, pin};
5use std::future::Future;
6
7/// Borrowed Future
8pub type Fut<'a, T> = &'a mut (dyn Future<Output = T> + Unpin);
9
10/// Future wrapper which races inner futures against each other
11pub 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}