diff --git a/embassy/src/util/mod.rs b/embassy/src/util/mod.rs index 7744778b..f5cc9623 100644 --- a/embassy/src/util/mod.rs +++ b/embassy/src/util/mod.rs @@ -1,9 +1,11 @@ //! Misc utilities mod forever; +mod select_all; mod yield_now; pub use forever::*; +pub use select_all::*; pub use yield_now::*; /// Unsafely unborrow an owned singleton out of a `&mut`. diff --git a/embassy/src/util/select_all.rs b/embassy/src/util/select_all.rs new file mode 100644 index 00000000..aef22d89 --- /dev/null +++ b/embassy/src/util/select_all.rs @@ -0,0 +1,58 @@ +use core::future::Future; +use core::pin::Pin; +use core::task::{Context, Poll}; + +/// Future for the [`select_all`] function. +#[derive(Debug)] +#[must_use = "futures do nothing unless you `.await` or poll them"] +pub struct SelectAll { + inner: [Fut; N], +} + +impl Unpin for SelectAll {} + +/// Creates a new future which will select over a list of futures. +/// +/// The returned future will wait for any future within `iter` to be ready. Upon +/// completion the item resolved will be returned, along with the index of the +/// future that was ready. +/// +/// # Panics +/// +/// This function will panic if the array specified contains no items. +pub fn select_all(arr: [Fut; N]) -> SelectAll { + assert!(N > 0); + SelectAll { inner: arr } +} + +impl SelectAll { + /// Consumes this combinator, returning the underlying futures. + pub fn into_inner(self) -> [Fut; N] { + self.inner + } +} + +impl Future for SelectAll { + type Output = (Fut::Output, usize); + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + // Safety: Since `self` is pinned, `inner` cannot move. Since `inner` cannot move, + // its elements also cannot move. Therefore it is safe to access `inner` and pin + // references to the contained futures. + let item = unsafe { + self.get_unchecked_mut() + .inner + .iter_mut() + .enumerate() + .find_map(|(i, f)| match Pin::new_unchecked(f).poll(cx) { + Poll::Pending => None, + Poll::Ready(e) => Some((i, e)), + }) + }; + + match item { + Some((idx, res)) => Poll::Ready((res, idx)), + None => Poll::Pending, + } + } +}