embassy_futures/
select.rs

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
17impl<A, B> Either<A, B> {
18    /// Did the first future complete first?
19    pub fn is_first(&self) -> bool {
20        matches!(self, Either::First(_))
21    }
22
23    /// Did the second future complete first?
24    pub fn is_second(&self) -> bool {
25        matches!(self, Either::Second(_))
26    }
27}
28
29/// Wait for one of two futures to complete.
30///
31/// This function returns a new future which polls all the futures.
32/// When one of them completes, it will complete with its result value.
33///
34/// The other future is dropped.
35pub fn select<A, B>(a: A, b: B) -> Select<A, B>
36where
37    A: Future,
38    B: Future,
39{
40    Select { a, b }
41}
42
43/// Future for the [`select`] function.
44#[derive(Debug)]
45#[must_use = "futures do nothing unless you `.await` or poll them"]
46pub struct Select<A, B> {
47    a: A,
48    b: B,
49}
50
51impl<A: Unpin, B: Unpin> Unpin for Select<A, B> {}
52
53impl<A, B> Future for Select<A, B>
54where
55    A: Future,
56    B: Future,
57{
58    type Output = Either<A::Output, B::Output>;
59
60    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
61        let this = unsafe { self.get_unchecked_mut() };
62        let a = unsafe { Pin::new_unchecked(&mut this.a) };
63        let b = unsafe { Pin::new_unchecked(&mut this.b) };
64        if let Poll::Ready(x) = a.poll(cx) {
65            return Poll::Ready(Either::First(x));
66        }
67        if let Poll::Ready(x) = b.poll(cx) {
68            return Poll::Ready(Either::Second(x));
69        }
70        Poll::Pending
71    }
72}
73
74// ====================================================================
75
76/// Result for [`select3`].
77#[derive(Debug, Clone)]
78#[cfg_attr(feature = "defmt", derive(defmt::Format))]
79pub enum Either3<A, B, C> {
80    /// First future finished first.
81    First(A),
82    /// Second future finished first.
83    Second(B),
84    /// Third future finished first.
85    Third(C),
86}
87
88impl<A, B, C> Either3<A, B, C> {
89    /// Did the first future complete first?
90    pub fn is_first(&self) -> bool {
91        matches!(self, Either3::First(_))
92    }
93
94    /// Did the second future complete first?
95    pub fn is_second(&self) -> bool {
96        matches!(self, Either3::Second(_))
97    }
98
99    /// Did the third future complete first?
100    pub fn is_third(&self) -> bool {
101        matches!(self, Either3::Third(_))
102    }
103}
104
105/// Same as [`select`], but with more futures.
106pub fn select3<A, B, C>(a: A, b: B, c: C) -> Select3<A, B, C>
107where
108    A: Future,
109    B: Future,
110    C: Future,
111{
112    Select3 { a, b, c }
113}
114
115/// Future for the [`select3`] function.
116#[derive(Debug)]
117#[must_use = "futures do nothing unless you `.await` or poll them"]
118pub struct Select3<A, B, C> {
119    a: A,
120    b: B,
121    c: C,
122}
123
124impl<A, B, C> Future for Select3<A, B, C>
125where
126    A: Future,
127    B: Future,
128    C: Future,
129{
130    type Output = Either3<A::Output, B::Output, C::Output>;
131
132    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
133        let this = unsafe { self.get_unchecked_mut() };
134        let a = unsafe { Pin::new_unchecked(&mut this.a) };
135        let b = unsafe { Pin::new_unchecked(&mut this.b) };
136        let c = unsafe { Pin::new_unchecked(&mut this.c) };
137        if let Poll::Ready(x) = a.poll(cx) {
138            return Poll::Ready(Either3::First(x));
139        }
140        if let Poll::Ready(x) = b.poll(cx) {
141            return Poll::Ready(Either3::Second(x));
142        }
143        if let Poll::Ready(x) = c.poll(cx) {
144            return Poll::Ready(Either3::Third(x));
145        }
146        Poll::Pending
147    }
148}
149
150// ====================================================================
151
152/// Result for [`select4`].
153#[derive(Debug, Clone)]
154#[cfg_attr(feature = "defmt", derive(defmt::Format))]
155pub enum Either4<A, B, C, D> {
156    /// First future finished first.
157    First(A),
158    /// Second future finished first.
159    Second(B),
160    /// Third future finished first.
161    Third(C),
162    /// Fourth future finished first.
163    Fourth(D),
164}
165
166impl<A, B, C, D> Either4<A, B, C, D> {
167    /// Did the first future complete first?
168    pub fn is_first(&self) -> bool {
169        matches!(self, Either4::First(_))
170    }
171
172    /// Did the second future complete first?
173    pub fn is_second(&self) -> bool {
174        matches!(self, Either4::Second(_))
175    }
176
177    /// Did the third future complete first?
178    pub fn is_third(&self) -> bool {
179        matches!(self, Either4::Third(_))
180    }
181
182    /// Did the fourth future complete first?
183    pub fn is_fourth(&self) -> bool {
184        matches!(self, Either4::Fourth(_))
185    }
186}
187
188/// Same as [`select`], but with more futures.
189pub fn select4<A, B, C, D>(a: A, b: B, c: C, d: D) -> Select4<A, B, C, D>
190where
191    A: Future,
192    B: Future,
193    C: Future,
194    D: Future,
195{
196    Select4 { a, b, c, d }
197}
198
199/// Future for the [`select4`] function.
200#[derive(Debug)]
201#[must_use = "futures do nothing unless you `.await` or poll them"]
202pub struct Select4<A, B, C, D> {
203    a: A,
204    b: B,
205    c: C,
206    d: D,
207}
208
209impl<A, B, C, D> Future for Select4<A, B, C, D>
210where
211    A: Future,
212    B: Future,
213    C: Future,
214    D: Future,
215{
216    type Output = Either4<A::Output, B::Output, C::Output, D::Output>;
217
218    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
219        let this = unsafe { self.get_unchecked_mut() };
220        let a = unsafe { Pin::new_unchecked(&mut this.a) };
221        let b = unsafe { Pin::new_unchecked(&mut this.b) };
222        let c = unsafe { Pin::new_unchecked(&mut this.c) };
223        let d = unsafe { Pin::new_unchecked(&mut this.d) };
224        if let Poll::Ready(x) = a.poll(cx) {
225            return Poll::Ready(Either4::First(x));
226        }
227        if let Poll::Ready(x) = b.poll(cx) {
228            return Poll::Ready(Either4::Second(x));
229        }
230        if let Poll::Ready(x) = c.poll(cx) {
231            return Poll::Ready(Either4::Third(x));
232        }
233        if let Poll::Ready(x) = d.poll(cx) {
234            return Poll::Ready(Either4::Fourth(x));
235        }
236        Poll::Pending
237    }
238}
239
240// ====================================================================
241
242/// Result for [`select5`].
243#[derive(Debug, Clone)]
244#[cfg_attr(feature = "defmt", derive(defmt::Format))]
245pub enum Either5<A, B, C, D, E> {
246    /// First future finished first.
247    First(A),
248    /// Second future finished first.
249    Second(B),
250    /// Third future finished first.
251    Third(C),
252    /// Fourth future finished first.
253    Fourth(D),
254    /// Fifth future finished first.
255    Fifth(E),
256}
257
258impl<A, B, C, D, E> Either5<A, B, C, D, E> {
259    /// Did the first future complete first?
260    pub fn is_first(&self) -> bool {
261        matches!(self, Either5::First(_))
262    }
263
264    /// Did the second future complete first?
265    pub fn is_second(&self) -> bool {
266        matches!(self, Either5::Second(_))
267    }
268
269    /// Did the third future complete first?
270    pub fn is_third(&self) -> bool {
271        matches!(self, Either5::Third(_))
272    }
273
274    /// Did the fourth future complete first?
275    pub fn is_fourth(&self) -> bool {
276        matches!(self, Either5::Fourth(_))
277    }
278
279    /// Did the fifth future complete first?
280    pub fn is_fifth(&self) -> bool {
281        matches!(self, Either5::Fifth(_))
282    }
283}
284
285/// Same as [`select`], but with more futures.
286pub fn select5<A, B, C, D, E>(a: A, b: B, c: C, d: D, e: E) -> Select5<A, B, C, D, E>
287where
288    A: Future,
289    B: Future,
290    C: Future,
291    D: Future,
292    E: Future,
293{
294    Select5 { a, b, c, d, e }
295}
296
297/// Future for the [`select5`] function.
298#[derive(Debug)]
299#[must_use = "futures do nothing unless you `.await` or poll them"]
300pub struct Select5<A, B, C, D, E> {
301    a: A,
302    b: B,
303    c: C,
304    d: D,
305    e: E,
306}
307
308impl<A, B, C, D, E> Future for Select5<A, B, C, D, E>
309where
310    A: Future,
311    B: Future,
312    C: Future,
313    D: Future,
314    E: Future,
315{
316    type Output = Either5<A::Output, B::Output, C::Output, D::Output, E::Output>;
317
318    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
319        let this = unsafe { self.get_unchecked_mut() };
320        let a = unsafe { Pin::new_unchecked(&mut this.a) };
321        let b = unsafe { Pin::new_unchecked(&mut this.b) };
322        let c = unsafe { Pin::new_unchecked(&mut this.c) };
323        let d = unsafe { Pin::new_unchecked(&mut this.d) };
324        let e = unsafe { Pin::new_unchecked(&mut this.e) };
325        if let Poll::Ready(x) = a.poll(cx) {
326            return Poll::Ready(Either5::First(x));
327        }
328        if let Poll::Ready(x) = b.poll(cx) {
329            return Poll::Ready(Either5::Second(x));
330        }
331        if let Poll::Ready(x) = c.poll(cx) {
332            return Poll::Ready(Either5::Third(x));
333        }
334        if let Poll::Ready(x) = d.poll(cx) {
335            return Poll::Ready(Either5::Fourth(x));
336        }
337        if let Poll::Ready(x) = e.poll(cx) {
338            return Poll::Ready(Either5::Fifth(x));
339        }
340        Poll::Pending
341    }
342}
343
344// ====================================================================
345
346/// Result for [`select6`].
347#[derive(Debug, Clone)]
348#[cfg_attr(feature = "defmt", derive(defmt::Format))]
349pub enum Either6<A, B, C, D, E, F> {
350    /// First future finished first.
351    First(A),
352    /// Second future finished first.
353    Second(B),
354    /// Third future finished first.
355    Third(C),
356    /// Fourth future finished first.
357    Fourth(D),
358    /// Fifth future finished first.
359    Fifth(E),
360    /// Sixth future finished first.
361    Sixth(F),
362}
363
364impl<A, B, C, D, E, F> Either6<A, B, C, D, E, F> {
365    /// Did the first future complete first?
366    pub fn is_first(&self) -> bool {
367        matches!(self, Either6::First(_))
368    }
369
370    /// Did the second future complete first?
371    pub fn is_second(&self) -> bool {
372        matches!(self, Either6::Second(_))
373    }
374
375    /// Did the third future complete first?
376    pub fn is_third(&self) -> bool {
377        matches!(self, Either6::Third(_))
378    }
379
380    /// Did the fourth future complete first?
381    pub fn is_fourth(&self) -> bool {
382        matches!(self, Either6::Fourth(_))
383    }
384
385    /// Did the fifth future complete first?
386    pub fn is_fifth(&self) -> bool {
387        matches!(self, Either6::Fifth(_))
388    }
389
390    /// Did the sixth future complete first?
391    pub fn is_sixth(&self) -> bool {
392        matches!(self, Either6::Sixth(_))
393    }
394}
395
396/// Same as [`select`], but with more futures.
397pub fn select6<A, B, C, D, E, F>(a: A, b: B, c: C, d: D, e: E, f: F) -> Select6<A, B, C, D, E, F>
398where
399    A: Future,
400    B: Future,
401    C: Future,
402    D: Future,
403    E: Future,
404    F: Future,
405{
406    Select6 { a, b, c, d, e, f }
407}
408
409/// Future for the [`select6`] function.
410#[derive(Debug)]
411#[must_use = "futures do nothing unless you `.await` or poll them"]
412pub struct Select6<A, B, C, D, E, F> {
413    a: A,
414    b: B,
415    c: C,
416    d: D,
417    e: E,
418    f: F,
419}
420
421impl<A, B, C, D, E, F> Future for Select6<A, B, C, D, E, F>
422where
423    A: Future,
424    B: Future,
425    C: Future,
426    D: Future,
427    E: Future,
428    F: Future,
429{
430    type Output = Either6<A::Output, B::Output, C::Output, D::Output, E::Output, F::Output>;
431
432    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
433        let this = unsafe { self.get_unchecked_mut() };
434        let a = unsafe { Pin::new_unchecked(&mut this.a) };
435        let b = unsafe { Pin::new_unchecked(&mut this.b) };
436        let c = unsafe { Pin::new_unchecked(&mut this.c) };
437        let d = unsafe { Pin::new_unchecked(&mut this.d) };
438        let e = unsafe { Pin::new_unchecked(&mut this.e) };
439        let f = unsafe { Pin::new_unchecked(&mut this.f) };
440        if let Poll::Ready(x) = a.poll(cx) {
441            return Poll::Ready(Either6::First(x));
442        }
443        if let Poll::Ready(x) = b.poll(cx) {
444            return Poll::Ready(Either6::Second(x));
445        }
446        if let Poll::Ready(x) = c.poll(cx) {
447            return Poll::Ready(Either6::Third(x));
448        }
449        if let Poll::Ready(x) = d.poll(cx) {
450            return Poll::Ready(Either6::Fourth(x));
451        }
452        if let Poll::Ready(x) = e.poll(cx) {
453            return Poll::Ready(Either6::Fifth(x));
454        }
455        if let Poll::Ready(x) = f.poll(cx) {
456            return Poll::Ready(Either6::Sixth(x));
457        }
458        Poll::Pending
459    }
460}
461
462// ====================================================================
463
464/// Future for the [`select_array`] function.
465#[derive(Debug)]
466#[must_use = "futures do nothing unless you `.await` or poll them"]
467pub struct SelectArray<Fut, const N: usize> {
468    inner: [Fut; N],
469}
470
471/// Creates a new future which will select over an array of futures.
472///
473/// The returned future will wait for any future to be ready. Upon
474/// completion the item resolved will be returned, along with the index of the
475/// future that was ready.
476///
477/// If the array is empty, the resulting future will be Pending forever.
478pub fn select_array<Fut: Future, const N: usize>(arr: [Fut; N]) -> SelectArray<Fut, N> {
479    SelectArray { inner: arr }
480}
481
482impl<Fut: Future, const N: usize> Future for SelectArray<Fut, N> {
483    type Output = (Fut::Output, usize);
484
485    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
486        // Safety: Since `self` is pinned, `inner` cannot move. Since `inner` cannot move,
487        // its elements also cannot move. Therefore it is safe to access `inner` and pin
488        // references to the contained futures.
489        let item = unsafe {
490            self.get_unchecked_mut()
491                .inner
492                .iter_mut()
493                .enumerate()
494                .find_map(|(i, f)| match Pin::new_unchecked(f).poll(cx) {
495                    Poll::Pending => None,
496                    Poll::Ready(e) => Some((i, e)),
497                })
498        };
499
500        match item {
501            Some((idx, res)) => Poll::Ready((res, idx)),
502            None => Poll::Pending,
503        }
504    }
505}
506
507// ====================================================================
508
509/// Future for the [`select_slice`] function.
510#[derive(Debug)]
511#[must_use = "futures do nothing unless you `.await` or poll them"]
512pub struct SelectSlice<'a, Fut> {
513    inner: Pin<&'a mut [Fut]>,
514}
515
516/// Creates a new future which will select over a slice of futures.
517///
518/// The returned future will wait for any future to be ready. Upon
519/// completion the item resolved will be returned, along with the index of the
520/// future that was ready.
521///
522/// If the slice is empty, the resulting future will be Pending forever.
523pub fn select_slice<'a, Fut: Future>(slice: Pin<&'a mut [Fut]>) -> SelectSlice<'a, Fut> {
524    SelectSlice { inner: slice }
525}
526
527impl<'a, Fut: Future> Future for SelectSlice<'a, Fut> {
528    type Output = (Fut::Output, usize);
529
530    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
531        // Safety: refer to
532        //   https://users.rust-lang.org/t/working-with-pinned-slices-are-there-any-structurally-pinning-vec-like-collection-types/50634/2
533        #[inline(always)]
534        fn pin_iter<T>(slice: Pin<&mut [T]>) -> impl Iterator<Item = Pin<&mut T>> {
535            unsafe { slice.get_unchecked_mut().iter_mut().map(|v| Pin::new_unchecked(v)) }
536        }
537        for (i, fut) in pin_iter(self.inner.as_mut()).enumerate() {
538            if let Poll::Ready(res) = fut.poll(cx) {
539                return Poll::Ready((res, i));
540            }
541        }
542
543        Poll::Pending
544    }
545}
embassy_futures/select.rs

embassy_futures/
select.rs
