5 files changed, 381 insertions, 167 deletions

diff --git a/embassy-sync/src/blocking_mutex/raw_rwlock.rs b/embassy-sync/src/blocking_mutex/raw_rwlock.rs
deleted file mode 100644
index de4bd1dc5..000000000
--- a/embassy-sync/src/blocking_mutex/raw_rwlock.rs
+++ /dev/null
@@ -1,86 +0,0 @@
-use core::sync::atomic::{AtomicUsize, Ordering};
-use core::task::Waker;
-use core::cell::UnsafeCell;
-
-pub trait RawRwLock {
-    fn lock_read(&self);
-    fn try_lock_read(&self) -> bool;
-    fn unlock_read(&self);
-    fn lock_write(&self);
-    fn try_lock_write(&self) -> bool;
-    fn unlock_write(&self);
-}
-
-pub struct RawRwLockImpl {
-    state: AtomicUsize,
-    waker: UnsafeCell<Option<Waker>>,
-}
-
-impl RawRwLockImpl {
-    pub const fn new() -> Self {
-        Self {
-            state: AtomicUsize::new(0),
-            waker: UnsafeCell::new(None),
-        }
-    }
-}
-
-unsafe impl Send for RawRwLockImpl {}
-unsafe impl Sync for RawRwLockImpl {}
-
-impl RawRwLock for RawRwLockImpl {
-    fn lock_read(&self) {
-        loop {
-            let state = self.state.load(Ordering::Acquire);
-            if state & 1 == 0 {
-                if self.state.compare_and_swap(state, state + 2, Ordering::AcqRel) == state {
-                    break;
-                }
-            }
-        }
-    }
-
-    fn try_lock_read(&self) -> bool {
-        let state = self.state.load(Ordering::Acquire);
-        if state & 1 == 0 {
-            if self.state.compare_and_swap(state, state + 2, Ordering::AcqRel) == state {
-                return true;
-            }
-        }
-        false
-    }
-
-    fn unlock_read(&self) {
-        self.state.fetch_sub(2, Ordering::Release);
-        if self.state.load(Ordering::Acquire) == 0 {
-            if let Some(waker) = unsafe { &*self.waker.get() } {
-                waker.wake_by_ref();
-            }
-        }
-    }
-
-    fn lock_write(&self) {
-        loop {
-            let state = self.state.load(Ordering::Acquire);
-            if state == 0 {
-                if self.state.compare_and_swap(0, 1, Ordering::AcqRel) == 0 {
-                    break;
-                }
-            }
-        }
-    }
-
-    fn try_lock_write(&self) -> bool {
-        if self.state.compare_and_swap(0, 1, Ordering::AcqRel) == 0 {
-            return true;
-        }
-        false
-    }
-
-    fn unlock_write(&self) {
-        self.state.store(0, Ordering::Release);
-        if let Some(waker) = unsafe { &*self.waker.get() } {
-            waker.wake_by_ref();
-        }
-    }
-}
diff --git a/embassy-sync/src/blocking_rwlock/mod.rs b/embassy-sync/src/blocking_rwlock/mod.rs
new file mode 100644
index 000000000..a304b77d6
--- /dev/null
+++ b/embassy-sync/src/blocking_rwlock/mod.rs
@@ -0,0 +1,221 @@
+//! Blocking read-write lock.
+//!
+//! This module provides a blocking read-write lock that can be used to synchronize data.
+pub mod raw_rwlock;
+
+use core::cell::UnsafeCell;
+
+use self::raw_rwlock::RawRwLock;
+
+/// Blocking read-write lock (not async)
+///
+/// Provides a blocking read-write lock primitive backed by an implementation of [`raw_rwlock::RawRwLock`].
+///
+/// Which implementation you select depends on the context in which you're using the read-write lock, and you can choose which kind
+/// of interior mutability fits your use case.
+///
+/// Use [`CriticalSectionRwLock`] when data can be shared between threads and interrupts.
+///
+/// Use [`NoopRwLock`] when data is only shared between tasks running on the same executor.
+///
+/// Use [`ThreadModeRwLock`] when data is shared between tasks running on the same executor but you want a global singleton.
+///
+/// In all cases, the blocking read-write lock is intended to be short lived and not held across await points.
+/// Use the async [`RwLock`](crate::rwlock::RwLock) if you need a lock that is held across await points.
+pub struct RwLock<R, T: ?Sized> {
+    // NOTE: `raw` must be FIRST, so when using ThreadModeRwLock the "can't drop in non-thread-mode" gets
+    // to run BEFORE dropping `data`.
+    raw: R,
+    data: UnsafeCell<T>,
+}
+
+unsafe impl<R: RawRwLock + Send, T: ?Sized + Send> Send for RwLock<R, T> {}
+unsafe impl<R: RawRwLock + Sync, T: ?Sized + Send> Sync for RwLock<R, T> {}
+
+impl<R: RawRwLock, T> RwLock<R, T> {
+    /// Creates a new read-write lock in an unlocked state ready for use.
+    #[inline]
+    pub const fn new(val: T) -> RwLock<R, T> {
+        RwLock {
+            raw: R::INIT,
+            data: UnsafeCell::new(val),
+        }
+    }
+
+    /// Creates a critical section and grants temporary read access to the protected data.
+    pub fn read_lock<U>(&self, f: impl FnOnce(&T) -> U) -> U {
+        self.raw.read_lock(|| {
+            let ptr = self.data.get() as *const T;
+            let inner = unsafe { &*ptr };
+            f(inner)
+        })
+    }
+
+    /// Creates a critical section and grants temporary write access to the protected data.
+    pub fn write_lock<U>(&self, f: impl FnOnce(&mut T) -> U) -> U {
+        self.raw.write_lock(|| {
+            let ptr = self.data.get() as *mut T;
+            let inner = unsafe { &mut *ptr };
+            f(inner)
+        })
+    }
+}
+
+impl<R, T> RwLock<R, T> {
+    /// Creates a new read-write lock based on a pre-existing raw read-write lock.
+    ///
+    /// This allows creating a read-write lock in a constant context on stable Rust.
+    #[inline]
+    pub const fn const_new(raw_rwlock: R, val: T) -> RwLock<R, T> {
+        RwLock {
+            raw: raw_rwlock,
+            data: UnsafeCell::new(val),
+        }
+    }
+
+    /// Consumes this read-write lock, returning the underlying data.
+    #[inline]
+    pub fn into_inner(self) -> T {
+        self.data.into_inner()
+    }
+
+    /// Returns a mutable reference to the underlying data.
+    ///
+    /// Since this call borrows the `RwLock` mutably, no actual locking needs to
+    /// take place---the mutable borrow statically guarantees no locks exist.
+    #[inline]
+    pub fn get_mut(&mut self) -> &mut T {
+        unsafe { &mut *self.data.get() }
+    }
+}
+
+/// A read-write lock that allows borrowing data across executors and interrupts.
+///
+/// # Safety
+///
+/// This read-write lock is safe to share between different executors and interrupts.
+pub type CriticalSectionRwLock<T> = RwLock<raw_rwlock::CriticalSectionRawRwLock, T>;
+
+/// A read-write lock that allows borrowing data in the context of a single executor.
+///
+/// # Safety
+///
+/// **This Read-Write Lock is only safe within a single executor.**
+pub type NoopRwLock<T> = RwLock<raw_rwlock::NoopRawRwLock, T>;
+
+impl<T> RwLock<raw_rwlock::CriticalSectionRawRwLock, T> {
+    /// Borrows the data for the duration of the critical section
+    pub fn borrow<'cs>(&'cs self, _cs: critical_section::CriticalSection<'cs>) -> &'cs T {
+        let ptr = self.data.get() as *const T;
+        unsafe { &*ptr }
+    }
+}
+
+impl<T> RwLock<raw_rwlock::NoopRawRwLock, T> {
+    /// Borrows the data
+    #[allow(clippy::should_implement_trait)]
+    pub fn borrow(&self) -> &T {
+        let ptr = self.data.get() as *const T;
+        unsafe { &*ptr }
+    }
+}
+
+// ThreadModeRwLock does NOT use the generic read-write lock from above because it's special:
+// it's Send+Sync even if T: !Send. There's no way to do that without specialization (I think?).
+//
+// There's still a ThreadModeRawRwLock for use with the generic RwLock (handy with Channel, for example),
+// but that will require T: Send even though it shouldn't be needed.
+
+#[cfg(any(cortex_m, feature = "std"))]
+pub use thread_mode_rwlock::*;
+#[cfg(any(cortex_m, feature = "std"))]
+mod thread_mode_rwlock {
+    use super::*;
+
+    /// A "read-write lock" that only allows borrowing from thread mode.
+    ///
+    /// # Safety
+    ///
+    /// **This Read-Write Lock is only safe on single-core systems.**
+    ///
+    /// On multi-core systems, a `ThreadModeRwLock` **is not sufficient** to ensure exclusive access.
+    pub struct ThreadModeRwLock<T: ?Sized> {
+        inner: UnsafeCell<T>,
+    }
+
+    // NOTE: ThreadModeRwLock only allows borrowing from one execution context ever: thread mode.
+    // Therefore it cannot be used to send non-sendable stuff between execution contexts, so it can
+    // be Send+Sync even if T is not Send (unlike CriticalSectionRwLock)
+    unsafe impl<T: ?Sized> Sync for ThreadModeRwLock<T> {}
+    unsafe impl<T: ?Sized> Send for ThreadModeRwLock<T> {}
+
+    impl<T> ThreadModeRwLock<T> {
+        /// Creates a new read-write lock
+        pub const fn new(value: T) -> Self {
+            ThreadModeRwLock {
+                inner: UnsafeCell::new(value),
+            }
+        }
+    }
+
+    impl<T: ?Sized> ThreadModeRwLock<T> {
+        /// Lock the `ThreadModeRwLock` for reading, granting access to the data.
+        ///
+        /// # Panics
+        ///
+        /// This will panic if not currently running in thread mode.
+        pub fn read_lock<R>(&self, f: impl FnOnce(&T) -> R) -> R {
+            f(self.borrow())
+        }
+
+        /// Lock the `ThreadModeRwLock` for writing, granting access to the data.
+        ///
+        /// # Panics
+        ///
+        /// This will panic if not currently running in thread mode.
+        pub fn write_lock<R>(&self, f: impl FnOnce(&mut T) -> R) -> R {
+            f(self.borrow_mut())
+        }
+
+        /// Borrows the data
+        ///
+        /// # Panics
+        ///
+        /// This will panic if not currently running in thread mode.
+        pub fn borrow(&self) -> &T {
+            assert!(
+                raw_rwlock::in_thread_mode(),
+                "ThreadModeRwLock can only be borrowed from thread mode."
+            );
+            unsafe { &*self.inner.get() }
+        }
+
+        /// Mutably borrows the data
+        ///
+        /// # Panics
+        ///
+        /// This will panic if not currently running in thread mode.
+        pub fn borrow_mut(&self) -> &mut T {
+            assert!(
+                raw_rwlock::in_thread_mode(),
+                "ThreadModeRwLock can only be borrowed from thread mode."
+            );
+            unsafe { &mut *self.inner.get() }
+        }
+    }
+
+    impl<T: ?Sized> Drop for ThreadModeRwLock<T> {
+        fn drop(&mut self) {
+            // Only allow dropping from thread mode. Dropping calls drop on the inner `T`, so
+            // `drop` needs the same guarantees as `lock`. `ThreadModeRwLock<T>` is Send even if
+            // T isn't, so without this check a user could create a ThreadModeRwLock in thread mode,
+            // send it to interrupt context and drop it there, which would "send" a T even if T is not Send.
+            assert!(
+                raw_rwlock::in_thread_mode(),
+                "ThreadModeRwLock can only be dropped from thread mode."
+            );
+
+            // Drop of the inner `T` happens after this.
+        }
+    }
+}
\ No newline at end of file
diff --git a/embassy-sync/src/blocking_rwlock/raw.rs b/embassy-sync/src/blocking_rwlock/raw.rs
new file mode 100644
index 000000000..85e8374b5
--- /dev/null
+++ b/embassy-sync/src/blocking_rwlock/raw.rs
@@ -0,0 +1,159 @@
+//! Read-Write Lock primitives.
+//!
+//! This module provides a trait for read-write locks that can be used in different contexts.
+use core::marker::PhantomData;
+
+/// Raw read-write lock trait.
+///
+/// This read-write lock is "raw", which means it does not actually contain the protected data, it
+/// just implements the read-write lock mechanism. For most uses you should use [`super::RwLock`] instead,
+/// which is generic over a RawRwLock and contains the protected data.
+///
+/// Note that, unlike other read-write locks, implementations only guarantee no
+/// concurrent access from other threads: concurrent access from the current
+/// thread is allowed. For example, it's possible to lock the same read-write lock multiple times reentrantly.
+///
+/// Therefore, locking a `RawRwLock` is only enough to guarantee safe shared (`&`) access
+/// to the data, it is not enough to guarantee exclusive (`&mut`) access.
+///
+/// # Safety
+///
+/// RawRwLock implementations must ensure that, while locked, no other thread can lock
+/// the RawRwLock concurrently.
+///
+/// Unsafe code is allowed to rely on this fact, so incorrect implementations will cause undefined behavior.
+pub unsafe trait RawRwLock {
+    /// Create a new `RawRwLock` instance.
+    ///
+    /// This is a const instead of a method to allow creating instances in const context.
+    const INIT: Self;
+
+    /// Lock this `RawRwLock` for reading.
+    fn read_lock<R>(&self, f: impl FnOnce() -> R) -> R;
+
+    /// Lock this `RawRwLock` for writing.
+    fn write_lock<R>(&self, f: impl FnOnce() -> R) -> R;
+}
+
+/// A read-write lock that allows borrowing data across executors and interrupts.
+///
+/// # Safety
+///
+/// This read-write lock is safe to share between different executors and interrupts.
+pub struct CriticalSectionRawRwLock {
+    _phantom: PhantomData<()>,
+}
+unsafe impl Send for CriticalSectionRawRwLock {}
+unsafe impl Sync for CriticalSectionRawRwLock {}
+
+impl CriticalSectionRawRwLock {
+    /// Create a new `CriticalSectionRawRwLock`.
+    pub const fn new() -> Self {
+        Self { _phantom: PhantomData }
+    }
+}
+
+unsafe impl RawRwLock for CriticalSectionRawRwLock {
+    const INIT: Self = Self::new();
+
+    fn read_lock<R>(&self, f: impl FnOnce() -> R) -> R {
+        critical_section::with(|_| f())
+    }
+
+    fn write_lock<R>(&self, f: impl FnOnce() -> R) -> R {
+        critical_section::with(|_| f())
+    }
+}
+
+// ================
+
+/// A read-write lock that allows borrowing data in the context of a single executor.
+///
+/// # Safety
+///
+/// **This Read-Write Lock is only safe within a single executor.**
+pub struct NoopRawRwLock {
+    _phantom: PhantomData<*mut ()>,
+}
+
+unsafe impl Send for NoopRawRwLock {}
+
+impl NoopRawRwLock {
+    /// Create a new `NoopRawRwLock`.
+    pub const fn new() -> Self {
+        Self { _phantom: PhantomData }
+    }
+}
+
+unsafe impl RawRwLock for NoopRawRwLock {
+    const INIT: Self = Self::new();
+    fn read_lock<R>(&self, f: impl FnOnce() -> R) -> R {
+        f()
+    }
+
+    fn write_lock<R>(&self, f: impl FnOnce() -> R) -> R {
+        f()
+    }
+}
+
+// ================
+
+#[cfg(any(cortex_m, feature = "std"))]
+mod thread_mode {
+    use super::*;
+
+    /// A "read-write lock" that only allows borrowing from thread mode.
+    ///
+    /// # Safety
+    ///
+    /// **This Read-Write Lock is only safe on single-core systems.**
+    ///
+    /// On multi-core systems, a `ThreadModeRawRwLock` **is not sufficient** to ensure exclusive access.
+    pub struct ThreadModeRawRwLock {
+        _phantom: PhantomData<()>,
+    }
+
+    unsafe impl Send for ThreadModeRawRwLock {}
+    unsafe impl Sync for ThreadModeRawRwLock {}
+
+    impl ThreadModeRawRwLock {
+        /// Create a new `ThreadModeRawRwLock`.
+        pub const fn new() -> Self {
+            Self { _phantom: PhantomData }
+        }
+    }
+
+    unsafe impl RawRwLock for ThreadModeRawRwLock {
+        const INIT: Self = Self::new();
+        fn read_lock<R>(&self, f: impl FnOnce() -> R) -> R {
+            assert!(in_thread_mode(), "ThreadModeRwLock can only be locked from thread mode.");
+
+            f()
+        }
+
+        fn write_lock<R>(&self, f: impl FnOnce() -> R) -> R {
+            assert!(in_thread_mode(), "ThreadModeRwLock can only be locked from thread mode.");
+
+            f()
+        }
+    }
+
+    impl Drop for ThreadModeRawRwLock {
+        fn drop(&mut self) {
+            assert!(
+                in_thread_mode(),
+                "ThreadModeRwLock can only be dropped from thread mode."
+            );
+        }
+    }
+
+    pub(crate) fn in_thread_mode() -> bool {
+        #[cfg(feature = "std")]
+        return Some("main") == std::thread::current().name();
+
+        #[cfg(not(feature = "std"))]
+        return unsafe { (0xE000ED04 as *const u32).read_volatile() } & 0x1FF == 0;
+    }
+}
+#[cfg(any(cortex_m, feature = "std"))]
+pub use thread_mode::*;
\ No newline at end of file
diff --git a/embassy-sync/src/lib.rs b/embassy-sync/src/lib.rs
index 5d91b4d9c..b9ead5f79 100644
--- a/embassy-sync/src/lib.rs
+++ b/embassy-sync/src/lib.rs
@@ -11,6 +11,7 @@ pub(crate) mod fmt;
 mod ring_buffer;
 
 pub mod blocking_mutex;
+pub mod blocking_rwlock;
 pub mod channel;
 pub mod lazy_lock;
 pub mod mutex;
diff --git a/embassy-sync/tests/rwlock.rs b/embassy-sync/tests/rwlock.rs
deleted file mode 100644
index 301cb0492..000000000
--- a/embassy-sync/tests/rwlock.rs
+++ /dev/null
@@ -1,81 +0,0 @@
-use embassy_sync::blocking_mutex::raw::NoopRawMutex;
-use embassy_sync::rwlock::RwLock;
-use futures_executor::block_on;
-
-#[futures_test::test]
-async fn test_rwlock_read() {
-    let lock = RwLock::<NoopRawMutex, _>::new(5);
-
-    {
-        let read_guard = lock.read().await;
-        assert_eq!(*read_guard, 5);
-    }
-
-    {
-        let read_guard = lock.read().await;
-        assert_eq!(*read_guard, 5);
-    }
-}
-
-#[futures_test::test]
-async fn test_rwlock_write() {
-    let lock = RwLock::<NoopRawMutex, _>::new(5);
-
-    {
-        let mut write_guard = lock.write().await;
-        *write_guard = 10;
-    }
-
-    {
-        let read_guard = lock.read().await;
-        assert_eq!(*read_guard, 10);
-    }
-}
-
-#[futures_test::test]
-async fn test_rwlock_try_read() {
-    let lock = RwLock::<NoopRawMutex, _>::new(5);
-
-    {
-        let read_guard = lock.try_read().unwrap();
-        assert_eq!(*read_guard, 5);
-    }
-
-    {
-        let read_guard = lock.try_read().unwrap();
-        assert_eq!(*read_guard, 5);
-    }
-}
-
-#[futures_test::test]
-async fn test_rwlock_try_write() {
-    let lock = RwLock::<NoopRawMutex, _>::new(5);
-
-    {
-        let mut write_guard = lock.try_write().unwrap();
-        *write_guard = 10;
-    }
-
-    {
-        let read_guard = lock.try_read().unwrap();
-        assert_eq!(*read_guard, 10);
-    }
-}
-
-#[futures_test::test]
-async fn test_rwlock_fairness() {
-    let lock = RwLock::<NoopRawMutex, _>::new(5);
-
-    let read1 = lock.read().await;
-    let read2 = lock.read().await;
-
-    let write_fut = lock.write();
-    futures_util::pin_mut!(write_fut);
-
-    assert!(futures_util::poll!(write_fut.as_mut()).is_pending());
-
-    drop(read1);
-    drop(read2);
-
-    assert!(futures_util::poll!(write_fut.as_mut()).is_ready());
-}