Merge #694

694: Add select, select3, select4. r=Dirbaio a=Dirbaio

The difference with those from the `futures` crate is they don't return the other partially-run futures, so they can work with `!Unpin` futures which makes them much easier to use.

Co-authored-by: Dario Nieuwenhuis <[email protected]>

5 files changed, 287 insertions, 125 deletions

diff --git a/embassy/src/util/mod.rs b/embassy/src/util/mod.rs
index f5cc96238..928edf0e2 100644
--- a/embassy/src/util/mod.rs
+++ b/embassy/src/util/mod.rs
@@ -1,74 +1,13 @@
 //! Misc utilities
 
 mod forever;
-mod select_all;
+mod select;
+mod steal;
+mod unborrow;
 mod yield_now;
 
 pub use forever::*;
-pub use select_all::*;
+pub use select::*;
+pub use steal::*;
+pub use unborrow::*;
 pub use yield_now::*;
-
-/// Unsafely unborrow an owned singleton out of a `&mut`.
-///
-/// It is intended to be implemented for owned peripheral singletons, such as `USART3` or `AnyPin`.
-/// Unborrowing an owned `T` yields the same `T`. Unborrowing a `&mut T` yields a copy of the T.
-///
-/// This allows writing HAL drivers that either own or borrow their peripherals, but that don't have
-/// to store pointers in the borrowed case.
-///
-/// Safety: this trait can be used to copy non-Copy types. Implementors must not cause
-/// immediate UB when copied, and must not cause UB when copies are later used, provided they
-/// are only used according the [`Self::unborrow`] safety contract.
-///
-pub unsafe trait Unborrow {
-    /// Unborrow result type
-    type Target;
-
-    /// Unborrow a value.
-    ///
-    /// Safety: This returns a copy of a singleton that's normally not
-    /// copiable. The returned copy must ONLY be used while the lifetime of `self` is
-    /// valid, as if it were accessed through `self` every time.
-    unsafe fn unborrow(self) -> Self::Target;
-}
-
-unsafe impl<'a, T: Unborrow> Unborrow for &'a mut T {
-    type Target = T::Target;
-    unsafe fn unborrow(self) -> Self::Target {
-        T::unborrow(core::ptr::read(self))
-    }
-}
-
-pub trait Steal {
-    unsafe fn steal() -> Self;
-}
-
-macro_rules! unsafe_impl_unborrow_tuples {
-    ($($t:ident),+) => {
-        unsafe impl<$($t),+> Unborrow for ($($t),+)
-        where
-            $(
-                $t: Unborrow<Target = $t>
-            ),+
-        {
-            type Target = ($($t),+);
-            unsafe fn unborrow(self) -> Self::Target {
-                self
-            }
-        }
-
-
-    };
-}
-
-unsafe_impl_unborrow_tuples!(A, B);
-unsafe_impl_unborrow_tuples!(A, B, C);
-unsafe_impl_unborrow_tuples!(A, B, C, D);
-unsafe_impl_unborrow_tuples!(A, B, C, D, E);
-unsafe_impl_unborrow_tuples!(A, B, C, D, E, F);
-unsafe_impl_unborrow_tuples!(A, B, C, D, E, F, G);
-unsafe_impl_unborrow_tuples!(A, B, C, D, E, F, G, H);
-unsafe_impl_unborrow_tuples!(A, B, C, D, E, F, G, H, I);
-unsafe_impl_unborrow_tuples!(A, B, C, D, E, F, G, H, I, J);
-unsafe_impl_unborrow_tuples!(A, B, C, D, E, F, G, H, I, J, K);
-unsafe_impl_unborrow_tuples!(A, B, C, D, E, F, G, H, I, J, K, L);
diff --git a/embassy/src/util/select.rs b/embassy/src/util/select.rs
new file mode 100644
index 000000000..ccc50f117
--- /dev/null
+++ b/embassy/src/util/select.rs
@@ -0,0 +1,218 @@
+use core::future::Future;
+use core::pin::Pin;
+use core::task::{Context, Poll};
+
+#[derive(Debug, Clone)]
+pub enum Either<A, B> {
+    First(A),
+    Second(B),
+}
+
+/// Wait for one of two futures to complete.
+///
+/// This function returns a new future which polls all the futures.
+/// When one of them completes, it will complete with its result value.
+///
+/// The other future is dropped.
+pub fn select<A, B>(a: A, b: B) -> Select<A, B>
+where
+    A: Future,
+    B: Future,
+{
+    Select { a, b }
+}
+
+/// Future for the [`select`] function.
+#[derive(Debug)]
+#[must_use = "futures do nothing unless you `.await` or poll them"]
+pub struct Select<A, B> {
+    a: A,
+    b: B,
+}
+
+impl<A: Unpin, B: Unpin> Unpin for Select<A, B> {}
+
+impl<A, B> Future for Select<A, B>
+where
+    A: Future,
+    B: Future,
+{
+    type Output = Either<A::Output, B::Output>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = unsafe { self.get_unchecked_mut() };
+        let a = unsafe { Pin::new_unchecked(&mut this.a) };
+        let b = unsafe { Pin::new_unchecked(&mut this.b) };
+        if let Poll::Ready(x) = a.poll(cx) {
+            return Poll::Ready(Either::First(x));
+        }
+        if let Poll::Ready(x) = b.poll(cx) {
+            return Poll::Ready(Either::Second(x));
+        }
+        Poll::Pending
+    }
+}
+
+// ====================================================================
+
+#[derive(Debug, Clone)]
+pub enum Either3<A, B, C> {
+    First(A),
+    Second(B),
+    Third(C),
+}
+
+/// Same as [`select`], but with more futures.
+pub fn select3<A, B, C>(a: A, b: B, c: C) -> Select3<A, B, C>
+where
+    A: Future,
+    B: Future,
+    C: Future,
+{
+    Select3 { a, b, c }
+}
+
+/// Future for the [`select3`] function.
+#[derive(Debug)]
+#[must_use = "futures do nothing unless you `.await` or poll them"]
+pub struct Select3<A, B, C> {
+    a: A,
+    b: B,
+    c: C,
+}
+
+impl<A, B, C> Future for Select3<A, B, C>
+where
+    A: Future,
+    B: Future,
+    C: Future,
+{
+    type Output = Either3<A::Output, B::Output, C::Output>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = unsafe { self.get_unchecked_mut() };
+        let a = unsafe { Pin::new_unchecked(&mut this.a) };
+        let b = unsafe { Pin::new_unchecked(&mut this.b) };
+        let c = unsafe { Pin::new_unchecked(&mut this.c) };
+        if let Poll::Ready(x) = a.poll(cx) {
+            return Poll::Ready(Either3::First(x));
+        }
+        if let Poll::Ready(x) = b.poll(cx) {
+            return Poll::Ready(Either3::Second(x));
+        }
+        if let Poll::Ready(x) = c.poll(cx) {
+            return Poll::Ready(Either3::Third(x));
+        }
+        Poll::Pending
+    }
+}
+
+// ====================================================================
+
+#[derive(Debug, Clone)]
+pub enum Either4<A, B, C, D> {
+    First(A),
+    Second(B),
+    Third(C),
+    Fourth(D),
+}
+
+/// Same as [`select`], but with more futures.
+pub fn select4<A, B, C, D>(a: A, b: B, c: C, d: D) -> Select4<A, B, C, D>
+where
+    A: Future,
+    B: Future,
+    C: Future,
+    D: Future,
+{
+    Select4 { a, b, c, d }
+}
+
+/// Future for the [`select4`] function.
+#[derive(Debug)]
+#[must_use = "futures do nothing unless you `.await` or poll them"]
+pub struct Select4<A, B, C, D> {
+    a: A,
+    b: B,
+    c: C,
+    d: D,
+}
+
+impl<A, B, C, D> Future for Select4<A, B, C, D>
+where
+    A: Future,
+    B: Future,
+    C: Future,
+    D: Future,
+{
+    type Output = Either4<A::Output, B::Output, C::Output, D::Output>;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        let this = unsafe { self.get_unchecked_mut() };
+        let a = unsafe { Pin::new_unchecked(&mut this.a) };
+        let b = unsafe { Pin::new_unchecked(&mut this.b) };
+        let c = unsafe { Pin::new_unchecked(&mut this.c) };
+        let d = unsafe { Pin::new_unchecked(&mut this.d) };
+        if let Poll::Ready(x) = a.poll(cx) {
+            return Poll::Ready(Either4::First(x));
+        }
+        if let Poll::Ready(x) = b.poll(cx) {
+            return Poll::Ready(Either4::Second(x));
+        }
+        if let Poll::Ready(x) = c.poll(cx) {
+            return Poll::Ready(Either4::Third(x));
+        }
+        if let Poll::Ready(x) = d.poll(cx) {
+            return Poll::Ready(Either4::Fourth(x));
+        }
+        Poll::Pending
+    }
+}
+
+// ====================================================================
+
+/// Future for the [`select_all`] function.
+#[derive(Debug)]
+#[must_use = "futures do nothing unless you `.await` or poll them"]
+pub struct SelectAll<Fut, const N: usize> {
+    inner: [Fut; N],
+}
+
+/// 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<Fut: Future, const N: usize>(arr: [Fut; N]) -> SelectAll<Fut, N> {
+    assert!(N > 0);
+    SelectAll { inner: arr }
+}
+
+impl<Fut: Future, const N: usize> Future for SelectAll<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,
+        }
+    }
+}
diff --git a/embassy/src/util/select_all.rs b/embassy/src/util/select_all.rs
deleted file mode 100644
index aef22d894..000000000
--- a/embassy/src/util/select_all.rs
+++ /dev/null
@@ -1,58 +0,0 @@
-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<Fut, const N: usize> {
-    inner: [Fut; N],
-}
-
-impl<Fut: Unpin, const N: usize> Unpin for SelectAll<Fut, N> {}
-
-/// 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<Fut: Future, const N: usize>(arr: [Fut; N]) -> SelectAll<Fut, N> {
-    assert!(N > 0);
-    SelectAll { inner: arr }
-}
-
-impl<Fut, const N: usize> SelectAll<Fut, N> {
-    /// Consumes this combinator, returning the underlying futures.
-    pub fn into_inner(self) -> [Fut; N] {
-        self.inner
-    }
-}
-
-impl<Fut: Future, const N: usize> Future for SelectAll<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,
-        }
-    }
-}
diff --git a/embassy/src/util/steal.rs b/embassy/src/util/steal.rs
new file mode 100644
index 000000000..a0d5f1359
--- /dev/null
+++ b/embassy/src/util/steal.rs
@@ -0,0 +1,3 @@
+pub trait Steal {
+    unsafe fn steal() -> Self;
+}
diff --git a/embassy/src/util/unborrow.rs b/embassy/src/util/unborrow.rs
new file mode 100644
index 000000000..dacfa3d42
--- /dev/null
+++ b/embassy/src/util/unborrow.rs
@@ -0,0 +1,60 @@
+/// Unsafely unborrow an owned singleton out of a `&mut`.
+///
+/// It is intended to be implemented for owned peripheral singletons, such as `USART3` or `AnyPin`.
+/// Unborrowing an owned `T` yields the same `T`. Unborrowing a `&mut T` yields a copy of the T.
+///
+/// This allows writing HAL drivers that either own or borrow their peripherals, but that don't have
+/// to store pointers in the borrowed case.
+///
+/// Safety: this trait can be used to copy non-Copy types. Implementors must not cause
+/// immediate UB when copied, and must not cause UB when copies are later used, provided they
+/// are only used according the [`Self::unborrow`] safety contract.
+///
+pub unsafe trait Unborrow {
+    /// Unborrow result type
+    type Target;
+
+    /// Unborrow a value.
+    ///
+    /// Safety: This returns a copy of a singleton that's normally not
+    /// copiable. The returned copy must ONLY be used while the lifetime of `self` is
+    /// valid, as if it were accessed through `self` every time.
+    unsafe fn unborrow(self) -> Self::Target;
+}
+
+unsafe impl<'a, T: Unborrow> Unborrow for &'a mut T {
+    type Target = T::Target;
+    unsafe fn unborrow(self) -> Self::Target {
+        T::unborrow(core::ptr::read(self))
+    }
+}
+
+macro_rules! unsafe_impl_unborrow_tuples {
+    ($($t:ident),+) => {
+        unsafe impl<$($t),+> Unborrow for ($($t),+)
+        where
+            $(
+                $t: Unborrow<Target = $t>
+            ),+
+        {
+            type Target = ($($t),+);
+            unsafe fn unborrow(self) -> Self::Target {
+                self
+            }
+        }
+
+
+    };
+}
+
+unsafe_impl_unborrow_tuples!(A, B);
+unsafe_impl_unborrow_tuples!(A, B, C);
+unsafe_impl_unborrow_tuples!(A, B, C, D);
+unsafe_impl_unborrow_tuples!(A, B, C, D, E);
+unsafe_impl_unborrow_tuples!(A, B, C, D, E, F);
+unsafe_impl_unborrow_tuples!(A, B, C, D, E, F, G);
+unsafe_impl_unborrow_tuples!(A, B, C, D, E, F, G, H);
+unsafe_impl_unborrow_tuples!(A, B, C, D, E, F, G, H, I);
+unsafe_impl_unborrow_tuples!(A, B, C, D, E, F, G, H, I, J);
+unsafe_impl_unborrow_tuples!(A, B, C, D, E, F, G, H, I, J, K);
+unsafe_impl_unborrow_tuples!(A, B, C, D, E, F, G, H, I, J, K, L);