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::FutureExt; 15 use futures::future::MaybeDone; 16 17 pub 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`. 25 macro_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 74 generate! { 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 } 108