xref: /aosp_15_r20/external/crosvm/cros_async/src/complete.rs (revision bb4ee6a4ae7042d18b07a98463b9c8b875e44b39)
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 
8 use std::future::Future;
9 use std::pin::Pin;
10 use std::task::Context;
11 use std::task::Poll;
12 
13 use futures::future::maybe_done;
14 use futures::future::MaybeDone;
15 use pin_utils::unsafe_pinned;
16 
17 // Macro-generate future combinators to allow for running different numbers of top-level futures in
18 // this FutureList. Generates the implementation of `FutureList` for the completion types. For an
19 // explicit example this is modeled after, see `UnitFutures`.
20 macro_rules! generate {
21     ($(
22         $(#[$doc:meta])*
23         ($Complete:ident, <$($Fut:ident),*>),
24     )*) => ($(
25         #[must_use = "Combinations of futures don't do anything unless run in an executor."]
26         pub(crate) struct $Complete<$($Fut: Future),*> {
27             $($Fut: MaybeDone<$Fut>,)*
28         }
29 
30         impl<$($Fut),*> $Complete<$($Fut),*>
31         where $(
32             $Fut: Future,
33         )*
34         {
35             // Safety:
36             // * No Drop impl
37             // * No Unpin impl
38             // * Not #[repr(packed)]
39             $(
40                 unsafe_pinned!($Fut: MaybeDone<$Fut>);
41             )*
42 
43             pub(crate) fn new($($Fut: $Fut),*) -> $Complete<$($Fut),*> {
44                 $(
45                     let $Fut = maybe_done($Fut);
46                 )*
47                 $Complete {
48                     $($Fut),*
49                 }
50             }
51         }
52 
53         impl<$($Fut),*> Future for $Complete<$($Fut),*>
54         where $(
55             $Fut: Future,
56         )*
57         {
58             type Output = ($($Fut::Output),*);
59 
60             fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll<Self::Output> {
61                 let mut complete = true;
62                 $(
63                     complete &= self.as_mut().$Fut().poll(cx).is_ready();
64                 )*
65 
66                 if complete {
67                     $(
68                         let $Fut = self.as_mut().$Fut().take_output().unwrap();
69                     )*
70                     Poll::Ready(($($Fut), *))
71                 } else {
72                     Poll::Pending
73                 }
74             }
75         }
76     )*)
77 }
78 
79 generate! {
80     /// _Future for the [`complete2`] function.
81     (Complete2, <_Fut1, _Fut2>),
82 
83     /// _Future for the [`complete3`] function.
84     (Complete3, <_Fut1, _Fut2, _Fut3>),
85 
86     /// _Future for the [`complete4`] function.
87     (Complete4, <_Fut1, _Fut2, _Fut3, _Fut4>),
88 
89     /// _Future for the [`complete5`] function.
90     (Complete5, <_Fut1, _Fut2, _Fut3, _Fut4, _Fut5>),
91 }
92