1//! Wait for the first of several futures to complete.
2
3use core::future::Future;
4use core::pin::Pin;
5use core::task::{Context, Poll};
6
7/// Result for [`select`].
8#[derive(Debug, Clone)]
9#[cfg_attr(feature = "defmt", derive(defmt::Format))]
10pub enum Either<A, B> {
11 /// First future finished first.
12 First(A),
13 /// Second future finished first.
14 Second(B),
15}
16
17/// Wait for one of two futures to complete.
18///
19/// This function returns a new future which polls all the futures.
20/// When one of them completes, it will complete with its result value.
21///
22/// The other future is dropped.
23pub fn select<A, B>(a: A, b: B) -> Select<A, B>
24where
25 A: Future,
26 B: Future,
27{
28 Select { a, b }
29}
30
31/// Future for the [`select`] function.
32#[derive(Debug)]
33#[must_use = "futures do nothing unless you `.await` or poll them"]
34pub struct Select<A, B> {
35 a: A,
36 b: B,
37}
38
39impl<A: Unpin, B: Unpin> Unpin for Select<A, B> {}
40
41impl<A, B> Future for Select<A, B>
42where
43 A: Future,
44 B: Future,
45{
46 type Output = Either<A::Output, B::Output>;
47
48 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
49 let this: &mut Select = unsafe { self.get_unchecked_mut() };
50 let a: Pin<&mut A> = unsafe { Pin::new_unchecked(&mut this.a) };
51 let b: Pin<&mut B> = unsafe { Pin::new_unchecked(&mut this.b) };
52 if let Poll::Ready(x: ::Output) = a.poll(cx) {
53 return Poll::Ready(Either::First(x));
54 }
55 if let Poll::Ready(x: ::Output) = b.poll(cx) {
56 return Poll::Ready(Either::Second(x));
57 }
58 Poll::Pending
59 }
60}
61
62// ====================================================================
63
64/// Result for [`select3`].
65#[derive(Debug, Clone)]
66#[cfg_attr(feature = "defmt", derive(defmt::Format))]
67pub enum Either3<A, B, C> {
68 /// First future finished first.
69 First(A),
70 /// Second future finished first.
71 Second(B),
72 /// Third future finished first.
73 Third(C),
74}
75
76/// Same as [`select`], but with more futures.
77pub fn select3<A, B, C>(a: A, b: B, c: C) -> Select3<A, B, C>
78where
79 A: Future,
80 B: Future,
81 C: Future,
82{
83 Select3 { a, b, c }
84}
85
86/// Future for the [`select3`] function.
87#[derive(Debug)]
88#[must_use = "futures do nothing unless you `.await` or poll them"]
89pub struct Select3<A, B, C> {
90 a: A,
91 b: B,
92 c: C,
93}
94
95impl<A, B, C> Future for Select3<A, B, C>
96where
97 A: Future,
98 B: Future,
99 C: Future,
100{
101 type Output = Either3<A::Output, B::Output, C::Output>;
102
103 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
104 let this: &mut Select3 = unsafe { self.get_unchecked_mut() };
105 let a: Pin<&mut A> = unsafe { Pin::new_unchecked(&mut this.a) };
106 let b: Pin<&mut B> = unsafe { Pin::new_unchecked(&mut this.b) };
107 let c: Pin<&mut C> = unsafe { Pin::new_unchecked(&mut this.c) };
108 if let Poll::Ready(x: ::Output) = a.poll(cx) {
109 return Poll::Ready(Either3::First(x));
110 }
111 if let Poll::Ready(x: ::Output) = b.poll(cx) {
112 return Poll::Ready(Either3::Second(x));
113 }
114 if let Poll::Ready(x: ::Output) = c.poll(cx) {
115 return Poll::Ready(Either3::Third(x));
116 }
117 Poll::Pending
118 }
119}
120
121// ====================================================================
122
123/// Result for [`select4`].
124#[derive(Debug, Clone)]
125#[cfg_attr(feature = "defmt", derive(defmt::Format))]
126pub enum Either4<A, B, C, D> {
127 /// First future finished first.
128 First(A),
129 /// Second future finished first.
130 Second(B),
131 /// Third future finished first.
132 Third(C),
133 /// Fourth future finished first.
134 Fourth(D),
135}
136
137/// Same as [`select`], but with more futures.
138pub fn select4<A, B, C, D>(a: A, b: B, c: C, d: D) -> Select4<A, B, C, D>
139where
140 A: Future,
141 B: Future,
142 C: Future,
143 D: Future,
144{
145 Select4 { a, b, c, d }
146}
147
148/// Future for the [`select4`] function.
149#[derive(Debug)]
150#[must_use = "futures do nothing unless you `.await` or poll them"]
151pub struct Select4<A, B, C, D> {
152 a: A,
153 b: B,
154 c: C,
155 d: D,
156}
157
158impl<A, B, C, D> Future for Select4<A, B, C, D>
159where
160 A: Future,
161 B: Future,
162 C: Future,
163 D: Future,
164{
165 type Output = Either4<A::Output, B::Output, C::Output, D::Output>;
166
167 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
168 let this: &mut Select4 = unsafe { self.get_unchecked_mut() };
169 let a: Pin<&mut A> = unsafe { Pin::new_unchecked(&mut this.a) };
170 let b: Pin<&mut B> = unsafe { Pin::new_unchecked(&mut this.b) };
171 let c: Pin<&mut C> = unsafe { Pin::new_unchecked(&mut this.c) };
172 let d: Pin<&mut D> = unsafe { Pin::new_unchecked(&mut this.d) };
173 if let Poll::Ready(x: ::Output) = a.poll(cx) {
174 return Poll::Ready(Either4::First(x));
175 }
176 if let Poll::Ready(x: ::Output) = b.poll(cx) {
177 return Poll::Ready(Either4::Second(x));
178 }
179 if let Poll::Ready(x: ::Output) = c.poll(cx) {
180 return Poll::Ready(Either4::Third(x));
181 }
182 if let Poll::Ready(x: ::Output) = d.poll(cx) {
183 return Poll::Ready(Either4::Fourth(x));
184 }
185 Poll::Pending
186 }
187}
188
189// ====================================================================
190
191/// Future for the [`select_array`] function.
192#[derive(Debug)]
193#[must_use = "futures do nothing unless you `.await` or poll them"]
194pub struct SelectArray<Fut, const N: usize> {
195 inner: [Fut; N],
196}
197
198/// Creates a new future which will select over an array of futures.
199///
200/// The returned future will wait for any future to be ready. Upon
201/// completion the item resolved will be returned, along with the index of the
202/// future that was ready.
203///
204/// If the array is empty, the resulting future will be Pending forever.
205pub fn select_array<Fut: Future, const N: usize>(arr: [Fut; N]) -> SelectArray<Fut, N> {
206 SelectArray { inner: arr }
207}
208
209impl<Fut: Future, const N: usize> Future for SelectArray<Fut, N> {
210 type Output = (Fut::Output, usize);
211
212 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
213 // Safety: Since `self` is pinned, `inner` cannot move. Since `inner` cannot move,
214 // its elements also cannot move. Therefore it is safe to access `inner` and pin
215 // references to the contained futures.
216 let item: Option<(usize, ::Output)> = unsafe {
217 self.get_unchecked_mut()
218 .inner
219 .iter_mut()
220 .enumerate()
221 .find_map(|(i: usize, f: &mut Fut)| match Pin::new_unchecked(pointer:f).poll(cx) {
222 Poll::Pending => None,
223 Poll::Ready(e: ::Output) => Some((i, e)),
224 })
225 };
226
227 match item {
228 Some((idx: usize, res: ::Output)) => Poll::Ready((res, idx)),
229 None => Poll::Pending,
230 }
231 }
232}
233
234// ====================================================================
235
236/// Future for the [`select_slice`] function.
237#[derive(Debug)]
238#[must_use = "futures do nothing unless you `.await` or poll them"]
239pub struct SelectSlice<'a, Fut> {
240 inner: &'a mut [Fut],
241}
242
243/// Creates a new future which will select over a slice of futures.
244///
245/// The returned future will wait for any future to be ready. Upon
246/// completion the item resolved will be returned, along with the index of the
247/// future that was ready.
248///
249/// If the slice is empty, the resulting future will be Pending forever.
250pub fn select_slice<'a, Fut: Future>(slice: &'a mut [Fut]) -> SelectSlice<'a, Fut> {
251 SelectSlice { inner: slice }
252}
253
254impl<'a, Fut: Future> Future for SelectSlice<'a, Fut> {
255 type Output = (Fut::Output, usize);
256
257 fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
258 // Safety: Since `self` is pinned, `inner` cannot move. Since `inner` cannot move,
259 // its elements also cannot move. Therefore it is safe to access `inner` and pin
260 // references to the contained futures.
261 let item: Option<(usize, ::Output)> = unsafe {
262 self.get_unchecked_mut()
263 .inner
264 .iter_mut()
265 .enumerate()
266 .find_map(|(i: usize, f: &mut Fut)| match Pin::new_unchecked(pointer:f).poll(cx) {
267 Poll::Pending => None,
268 Poll::Ready(e: ::Output) => Some((i, e)),
269 })
270 };
271
272 match item {
273 Some((idx: usize, res: ::Output)) => Poll::Ready((res, idx)),
274 None => Poll::Pending,
275 }
276 }
277}
278