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
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 = unsafe { self.get_unchecked_mut() };
50        let a = unsafe { Pin::new_unchecked(&mut this.a) };
51        let b = unsafe { Pin::new_unchecked(&mut this.b) };
52        if let Poll::Ready(x) = a.poll(cx) {
53            return Poll::Ready(Either::First(x));
54        }
55        if let Poll::Ready(x) = 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 = unsafe { self.get_unchecked_mut() };
105        let a = unsafe { Pin::new_unchecked(&mut this.a) };
106        let b = unsafe { Pin::new_unchecked(&mut this.b) };
107        let c = unsafe { Pin::new_unchecked(&mut this.c) };
108        if let Poll::Ready(x) = a.poll(cx) {
109            return Poll::Ready(Either3::First(x));
110        }
111        if let Poll::Ready(x) = b.poll(cx) {
112            return Poll::Ready(Either3::Second(x));
113        }
114        if let Poll::Ready(x) = 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 = unsafe { self.get_unchecked_mut() };
169        let a = unsafe { Pin::new_unchecked(&mut this.a) };
170        let b = unsafe { Pin::new_unchecked(&mut this.b) };
171        let c = unsafe { Pin::new_unchecked(&mut this.c) };
172        let d = unsafe { Pin::new_unchecked(&mut this.d) };
173        if let Poll::Ready(x) = a.poll(cx) {
174            return Poll::Ready(Either4::First(x));
175        }
176        if let Poll::Ready(x) = b.poll(cx) {
177            return Poll::Ready(Either4::Second(x));
178        }
179        if let Poll::Ready(x) = c.poll(cx) {
180            return Poll::Ready(Either4::Third(x));
181        }
182        if let Poll::Ready(x) = 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 = unsafe {
217            self.get_unchecked_mut()
218                .inner
219                .iter_mut()
220                .enumerate()
221                .find_map(|(i, f)| match Pin::new_unchecked(f).poll(cx) {
222                    Poll::Pending => None,
223                    Poll::Ready(e) => Some((i, e)),
224                })
225        };
226
227        match item {
228            Some((idx, res)) => 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 = unsafe {
262            self.get_unchecked_mut()
263                .inner
264                .iter_mut()
265                .enumerate()
266                .find_map(|(i, f)| match Pin::new_unchecked(f).poll(cx) {
267                    Poll::Pending => None,
268                    Poll::Ready(e) => Some((i, e)),
269                })
270        };
271
272        match item {
273            Some((idx, res)) => Poll::Ready((res, idx)),
274            None => Poll::Pending,
275        }
276    }
277}
embassy_futures/select.rs

embassy_futures/
select.rs
