From 3763baf8fa4cf332a81214eb1610afdaf5173824 Mon Sep 17 00:00:00 2001
From: Dario Nieuwenhuis <dirbaio@dirbaio.net>
Date: Sun, 28 Aug 2022 21:15:30 +0200
Subject: [PATCH] futures: add select_slice, rename select_all to select_array.

---
 embassy-futures/README.md     |  3 +-
 embassy-futures/src/select.rs | 62 +++++++++++++++++++++++++++++++----
 2 files changed, 57 insertions(+), 8 deletions(-)

diff --git a/embassy-futures/README.md b/embassy-futures/README.md
index 971f4c835..3810bcafd 100644
--- a/embassy-futures/README.md
+++ b/embassy-futures/README.md
@@ -5,5 +5,6 @@ Utilities for working with futures:
 - [`select`](select::select) - waiting for one out of two futures to complete.
 - [`select3`](select::select3) - waiting for one out of three futures to complete.
 - [`select4`](select::select4) - waiting for one out of four futures to complete.
-- [`select_all`](select::select_all) - waiting for one future in a list of futures to complete.
+- [`select_array`](select::select_array) - waiting for one future in an array of futures to complete.
+- [`select_slice`](select::select_slice) - waiting for one future in a slice of futures to complete.
 - [`yield_now`](yield_now::yield_now) - yielding the current task.
diff --git a/embassy-futures/src/select.rs b/embassy-futures/src/select.rs
index 53fa1da61..facc2f60b 100644
--- a/embassy-futures/src/select.rs
+++ b/embassy-futures/src/select.rs
@@ -186,28 +186,76 @@ where
 
 // ====================================================================
 
-/// Future for the [`select_all`] function.
+/// Future for the [`select_array`] function.
 #[derive(Debug)]
 #[must_use = "futures do nothing unless you `.await` or poll them"]
-pub struct SelectAll<Fut, const N: usize> {
+pub struct SelectArray<Fut, const N: usize> {
     inner: [Fut; N],
 }
 
-/// Creates a new future which will select over a list of futures.
+/// Creates a new future which will select over an array of futures.
 ///
-/// The returned future will wait for any future within `iter` to be ready. Upon
+/// The returned future will wait for any future 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<Fut: Future, const N: usize>(arr: [Fut; N]) -> SelectAll<Fut, N> {
+pub fn select_array<Fut: Future, const N: usize>(arr: [Fut; N]) -> SelectArray<Fut, N> {
     assert!(N > 0);
-    SelectAll { inner: arr }
+    SelectArray { inner: arr }
 }
 
-impl<Fut: Future, const N: usize> Future for SelectAll<Fut, N> {
+impl<Fut: Future, const N: usize> Future for SelectArray<Fut, N> {
+    type Output = (Fut::Output, usize);
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        // 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,
+        }
+    }
+}
+
+// ====================================================================
+
+/// Future for the [`select_slice`] function.
+#[derive(Debug)]
+#[must_use = "futures do nothing unless you `.await` or poll them"]
+pub struct SelectSlice<'a, Fut> {
+    inner: &'a mut [Fut],
+}
+
+/// Creates a new future which will select over a slice of futures.
+///
+/// The returned future will wait for any future 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 slice specified contains no items.
+pub fn select_slice<'a, Fut: Future>(slice: &'a mut [Fut]) -> SelectSlice<'a, Fut> {
+    assert!(!slice.is_empty());
+    SelectSlice { inner: slice }
+}
+
+impl<'a, Fut: Future> Future for SelectSlice<'a, Fut> {
     type Output = (Fut::Output, usize);
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {