1 | use super::assert_future; |
2 | use crate::future::{Either, FutureExt}; |
3 | use core::pin::Pin; |
4 | use futures_core::future::{FusedFuture, Future}; |
5 | use futures_core::task::{Context, Poll}; |
6 | |
7 | /// Future for the [`select()`] function. |
8 | #[must_use = "futures do nothing unless you `.await` or poll them" ] |
9 | #[derive(Debug)] |
10 | pub struct Select<A, B> { |
11 | inner: Option<(A, B)>, |
12 | } |
13 | |
14 | impl<A: Unpin, B: Unpin> Unpin for Select<A, B> {} |
15 | |
16 | /// Waits for either one of two differently-typed futures to complete. |
17 | /// |
18 | /// This function will return a new future which awaits for either one of both |
19 | /// futures to complete. The returned future will finish with both the value |
20 | /// resolved and a future representing the completion of the other work. |
21 | /// |
22 | /// Note that this function consumes the receiving futures and returns a |
23 | /// wrapped version of them. |
24 | /// |
25 | /// Also note that if both this and the second future have the same |
26 | /// output type you can use the `Either::factor_first` method to |
27 | /// conveniently extract out the value at the end. |
28 | /// |
29 | /// # Examples |
30 | /// |
31 | /// A simple example |
32 | /// |
33 | /// ``` |
34 | /// # futures::executor::block_on(async { |
35 | /// use futures::{ |
36 | /// pin_mut, |
37 | /// future::Either, |
38 | /// future::self, |
39 | /// }; |
40 | /// |
41 | /// // These two futures have different types even though their outputs have the same type. |
42 | /// let future1 = async { |
43 | /// future::pending::<()>().await; // will never finish |
44 | /// 1 |
45 | /// }; |
46 | /// let future2 = async { |
47 | /// future::ready(2).await |
48 | /// }; |
49 | /// |
50 | /// // 'select' requires Future + Unpin bounds |
51 | /// pin_mut!(future1); |
52 | /// pin_mut!(future2); |
53 | /// |
54 | /// let value = match future::select(future1, future2).await { |
55 | /// Either::Left((value1, _)) => value1, // `value1` is resolved from `future1` |
56 | /// // `_` represents `future2` |
57 | /// Either::Right((value2, _)) => value2, // `value2` is resolved from `future2` |
58 | /// // `_` represents `future1` |
59 | /// }; |
60 | /// |
61 | /// assert!(value == 2); |
62 | /// # }); |
63 | /// ``` |
64 | /// |
65 | /// A more complex example |
66 | /// |
67 | /// ``` |
68 | /// use futures::future::{self, Either, Future, FutureExt}; |
69 | /// |
70 | /// // A poor-man's join implemented on top of select |
71 | /// |
72 | /// fn join<A, B>(a: A, b: B) -> impl Future<Output=(A::Output, B::Output)> |
73 | /// where A: Future + Unpin, |
74 | /// B: Future + Unpin, |
75 | /// { |
76 | /// future::select(a, b).then(|either| { |
77 | /// match either { |
78 | /// Either::Left((x, b)) => b.map(move |y| (x, y)).left_future(), |
79 | /// Either::Right((y, a)) => a.map(move |x| (x, y)).right_future(), |
80 | /// } |
81 | /// }) |
82 | /// } |
83 | /// ``` |
84 | pub fn select<A, B>(future1: A, future2: B) -> Select<A, B> |
85 | where |
86 | A: Future + Unpin, |
87 | B: Future + Unpin, |
88 | { |
89 | assert_future::<Either<(A::Output, B), (B::Output, A)>, _>(Select { |
90 | inner: Some((future1, future2)), |
91 | }) |
92 | } |
93 | |
94 | impl<A, B> Future for Select<A, B> |
95 | where |
96 | A: Future + Unpin, |
97 | B: Future + Unpin, |
98 | { |
99 | type Output = Either<(A::Output, B), (B::Output, A)>; |
100 | |
101 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> { |
102 | /// When compiled with `-C opt-level=z`, this function will help the compiler eliminate the `None` branch, where |
103 | /// `Option::unwrap` does not. |
104 | #[inline (always)] |
105 | fn unwrap_option<T>(value: Option<T>) -> T { |
106 | match value { |
107 | None => unreachable!(), |
108 | Some(value) => value, |
109 | } |
110 | } |
111 | |
112 | let (a, b) = self.inner.as_mut().expect("cannot poll Select twice" ); |
113 | |
114 | if let Poll::Ready(val) = a.poll_unpin(cx) { |
115 | return Poll::Ready(Either::Left((val, unwrap_option(self.inner.take()).1))); |
116 | } |
117 | |
118 | if let Poll::Ready(val) = b.poll_unpin(cx) { |
119 | return Poll::Ready(Either::Right((val, unwrap_option(self.inner.take()).0))); |
120 | } |
121 | |
122 | Poll::Pending |
123 | } |
124 | } |
125 | |
126 | impl<A, B> FusedFuture for Select<A, B> |
127 | where |
128 | A: Future + Unpin, |
129 | B: Future + Unpin, |
130 | { |
131 | fn is_terminated(&self) -> bool { |
132 | self.inner.is_none() |
133 | } |
134 | } |
135 | |