1 // Copyright 2020 The Chromium OS Authors. All rights reserved. 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::{ 9 future::Future, 10 pin::Pin, 11 task::{Context, Poll}, 12 }; 13 14 use futures::future::{maybe_done, FutureExt, MaybeDone}; 15 16 pub enum SelectResult<F: Future> { 17 Pending(F), 18 Finished(F::Output), 19 } 20 21 // Macro-generate future combinators to allow for running different numbers of top-level futures in 22 // this FutureList. Generates the implementation of `FutureList` for the select types. For an 23 // explicit example this is modeled after, see `UnitFutures`. 24 macro_rules! generate { 25 ($( 26 $(#[$doc:meta])* 27 ($Select:ident, <$($Fut:ident),*>), 28 )*) => ($( 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