From faf03f414b983ba063986e6412cea4b4ad05b8ea Mon Sep 17 00:00:00 2001
From: Emil Fresk <emil.fresk@gmail.com>
Date: Thu, 15 Feb 2024 17:36:27 +0100
Subject: [PATCH] .

---
 rtic-sync/src/bit_channel.rs | 124 +++++++++++++++++++++--------------
 1 file changed, 76 insertions(+), 48 deletions(-)

diff --git a/rtic-sync/src/bit_channel.rs b/rtic-sync/src/bit_channel.rs
index bd436d1658..b502e94afc 100644
--- a/rtic-sync/src/bit_channel.rs
+++ b/rtic-sync/src/bit_channel.rs
@@ -1,8 +1,11 @@
 //! A channel operating on bitflags.
 
+use core::{future::poll_fn, task::Poll};
+
+use atomic::AtomicType;
 use bitflags::{Bits, Flags};
-use core::sync::atomic::{AtomicU16, AtomicU8, Ordering};
-use portable_atomic::{AtomicI16, AtomicI32, AtomicI64, AtomicI8, AtomicU32, AtomicU64};
+use portable_atomic::Ordering;
+use rtic_common::waker_registration::CriticalSectionWakerRegistration as CSWaker;
 
 /// A channel for setting and clearing `bitflags` concurrently.
 pub struct BitChannel<T: Flags>
@@ -11,6 +14,7 @@ where
     T::Bits: AtomicType,
 {
     atomic: <T::Bits as AtomicType>::Atomic,
+    waker: CSWaker,
 }
 
 impl<T> BitChannel<T>
@@ -22,72 +26,96 @@ where
     pub const fn new() -> Self {
         BitChannel {
             atomic: T::Bits::ATOMIC_ZERO,
+            waker: CSWaker::new(),
         }
     }
 
     /// Set `bitflag`s.
     #[inline]
     pub fn send(&self, flags: T) {
-        T::Bits::fetch_or(&self.atomic, flags.bits(), Ordering::Relaxed);
+        T::Bits::fetch_or(&self.atomic, flags.bits(), Ordering::Acquire);
+
+        self.waker.wake();
     }
 
-    /// Receive the current value of the `bitflags` and reset all flags.
+    /// Receive the current value of the `bitflags` and reset all flags. This can be accessed
+    /// concurrently but not all receivers will see the flags, only the first one will.
     #[inline]
     pub fn recv(&self) -> T {
         <T as Flags>::from_bits_retain(T::Bits::fetch_and(
             &self.atomic,
             T::Bits::EMPTY,
-            Ordering::Relaxed,
+            Ordering::Release,
         ))
     }
+
+    /// Wait for new values, the return is guaranteed non-empty.
+    pub async fn wait(&self) -> T {
+        poll_fn(|cx| {
+            self.waker.register(cx.waker());
+            let val = self.recv();
+
+            if val.is_empty() {
+                Poll::Pending
+            } else {
+                Poll::Ready(val)
+            }
+        })
+        .await
+    }
 }
 
-/// Generic atomic trait, allows for taking any `bitflags::Bits` as an atomic.
-pub trait AtomicType: Sized {
-    /// The underlying atomic, e.g. `AtomicU8`, `AtomicU16`.
-    type Atomic: From<Self>;
-
-    /// The definition of that atomic with value 0.
-    const ATOMIC_ZERO: Self::Atomic;
-
-    /// The atomic's `fetch_and` implementation forwarded.
-    fn fetch_and(a: &Self::Atomic, b: Self, order: Ordering) -> Self;
-
-    /// The atomic's `fetch_or` implementation forwarded.
-    fn fetch_or(a: &Self::Atomic, b: Self, order: Ordering) -> Self;
-}
-
-macro_rules! atomic_type_impl {
-    ($atomic:ty, $integer:ty) => {
-        impl AtomicType for $integer {
-            type Atomic = $atomic;
-
-            const ATOMIC_ZERO: Self::Atomic = <$atomic>::new(0);
-
-            #[inline(always)]
-            fn fetch_and(a: &Self::Atomic, b: Self, order: Ordering) -> Self {
-                a.fetch_and(b, order)
-            }
-
-            #[inline(always)]
-            fn fetch_or(a: &Self::Atomic, b: Self, order: Ordering) -> Self {
-                a.fetch_or(b, order)
-            }
-        }
+mod atomic {
+    use portable_atomic::{
+        AtomicI16, AtomicI32, AtomicI64, AtomicI8, AtomicU16, AtomicU32, AtomicU64, AtomicU8,
+        Ordering,
     };
+
+    /// Generic atomic trait, allows for taking any `bitflags::Bits` as an atomic.
+    pub trait AtomicType: Sized {
+        /// The underlying atomic, e.g. `AtomicU8`, `AtomicU16`.
+        type Atomic: From<Self>;
+
+        /// The definition of that atomic with value 0.
+        const ATOMIC_ZERO: Self::Atomic;
+
+        /// The atomic's `fetch_and` implementation forwarded.
+        fn fetch_and(a: &Self::Atomic, b: Self, order: Ordering) -> Self;
+
+        /// The atomic's `fetch_or` implementation forwarded.
+        fn fetch_or(a: &Self::Atomic, b: Self, order: Ordering) -> Self;
+    }
+
+    macro_rules! atomic_type_impl {
+        ($atomic:ty, $integer:ty) => {
+            impl AtomicType for $integer {
+                type Atomic = $atomic;
+
+                const ATOMIC_ZERO: Self::Atomic = <$atomic>::new(0);
+
+                #[inline(always)]
+                fn fetch_and(a: &Self::Atomic, b: Self, order: Ordering) -> Self {
+                    a.fetch_and(b, order)
+                }
+
+                #[inline(always)]
+                fn fetch_or(a: &Self::Atomic, b: Self, order: Ordering) -> Self {
+                    a.fetch_or(b, order)
+                }
+            }
+        };
+    }
+
+    atomic_type_impl!(AtomicU8, u8);
+    atomic_type_impl!(AtomicU16, u16);
+    atomic_type_impl!(AtomicU32, u32);
+    atomic_type_impl!(AtomicU64, u64);
+    atomic_type_impl!(AtomicI8, i8);
+    atomic_type_impl!(AtomicI16, i16);
+    atomic_type_impl!(AtomicI32, i32);
+    atomic_type_impl!(AtomicI64, i64);
 }
 
-atomic_type_impl!(AtomicU8, u8);
-atomic_type_impl!(AtomicU16, u16);
-atomic_type_impl!(AtomicU32, u32);
-atomic_type_impl!(AtomicU64, u64);
-atomic_type_impl!(AtomicI8, i8);
-atomic_type_impl!(AtomicI16, i16);
-atomic_type_impl!(AtomicI32, i32);
-atomic_type_impl!(AtomicI64, i64);
-
-// etc...
-
 #[cfg(test)]
 mod tests {
     use super::*;