cros_async/
select.rs

1// Copyright 2020 The ChromiumOS Authors
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5// Need non-snake case so the macro can re-use type names for variables.
6#![allow(non_snake_case)]
7
8use std::future::Future;
9use std::pin::Pin;
10use std::task::Context;
11use std::task::Poll;
12
13use futures::future::maybe_done;
14use futures::future::FutureExt;
15use futures::future::MaybeDone;
16
17pub enum SelectResult<F: Future> {
18    Pending(F),
19    Finished(F::Output),
20}
21
22// Macro-generate future combinators to allow for running different numbers of top-level futures in
23// this FutureList. Generates the implementation of `FutureList` for the select types. For an
24// explicit example this is modeled after, see `UnitFutures`.
25macro_rules! generate {
26    ($(
27        $(#[$doc:meta])*
28        ($Select:ident, <$($Fut:ident),*>),
29    )*) => ($(
30        paste::item! {
31            pub(crate) struct $Select<$($Fut: Future + Unpin),*> {
32                $($Fut: MaybeDone<$Fut>,)*
33            }
34        }
35
36        impl<$($Fut: Future + Unpin),*> $Select<$($Fut),*> {
37            paste::item! {
38                pub(crate) fn new($($Fut: $Fut),*) -> $Select<$($Fut),*> {
39                    $Select {
40                        $($Fut: maybe_done($Fut),)*
41                    }
42                }
43            }
44        }
45
46        impl<$($Fut: Future + Unpin),*> Future for $Select<$($Fut),*> {
47            type Output = ($(SelectResult<$Fut>),*);
48
49            fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
50                let mut complete = false;
51                $(
52                    let $Fut = Pin::new(&mut self.$Fut);
53                    // The future impls `Unpin`, use `poll_unpin` to avoid wrapping it in
54                    // `Pin` to call `poll`.
55                    complete |= self.$Fut.poll_unpin(cx).is_ready();
56                )*
57
58                if complete {
59                    Poll::Ready(($(
60                        match std::mem::replace(&mut self.$Fut, MaybeDone::Gone) {
61                            MaybeDone::Future(f) => SelectResult::Pending(f),
62                            MaybeDone::Done(o) => SelectResult::Finished(o),
63                            MaybeDone::Gone => unreachable!(),
64                        }
65                    ), *))
66                } else {
67                    Poll::Pending
68                }
69            }
70        }
71    )*)
72}
73
74generate! {
75    /// _Future for the [`select2`] function.
76    (Select2, <_Fut1, _Fut2>),
77
78    /// _Future for the [`select3`] function.
79    (Select3, <_Fut1, _Fut2, _Fut3>),
80
81    /// _Future for the [`select4`] function.
82    (Select4, <_Fut1, _Fut2, _Fut3, _Fut4>),
83
84    /// _Future for the [`select5`] function.
85    (Select5, <_Fut1, _Fut2, _Fut3, _Fut4, _Fut5>),
86
87    /// _Future for the [`select6`] function.
88    (Select6, <_Fut1, _Fut2, _Fut3, _Fut4, _Fut5, _Fut6>),
89
90    /// _Future for the [`select7`] function.
91    (Select7, <_Fut1, _Fut2, _Fut3, _Fut4, _Fut5, _Fut6, _Fut7>),
92
93    /// _Future for the [`select8`] function.
94    (Select8, <_Fut1, _Fut2, _Fut3, _Fut4, _Fut5, _Fut6, _Fut7, _Fut8>),
95
96    /// _Future for the [`select9`] function.
97    (Select9, <_Fut1, _Fut2, _Fut3, _Fut4, _Fut5, _Fut6, _Fut7, _Fut8, _Fut9>),
98
99    /// _Future for the [`select10`] function.
100    (Select10, <_Fut1, _Fut2, _Fut3, _Fut4, _Fut5, _Fut6, _Fut7, _Fut8, _Fut9, _Fut10>),
101
102    /// _Future for the [`select11`] function.
103    (Select11, <_Fut1, _Fut2, _Fut3, _Fut4, _Fut5, _Fut6, _Fut7, _Fut8, _Fut9, _Fut10, _Fut11>),
104
105    /// _Future for the [`select12`] function.
106    (Select12, <_Fut1, _Fut2, _Fut3, _Fut4, _Fut5, _Fut6, _Fut7, _Fut8, _Fut9, _Fut10, _Fut11, _Fut12>),
107}