From bbc12185f8c8dfeb909de2d6cc03f52730e17040 Mon Sep 17 00:00:00 2001
From: Kaspar Schleiser <kaspar@schleiser.de>
Date: Tue, 20 Feb 2024 21:46:31 +0100
Subject: [PATCH] riot-rs-embassy: improve `SendCell` API

---
 src/riot-rs-embassy/src/lib.rs      |  2 +-
 src/riot-rs-embassy/src/network.rs  |  5 ++--
 src/riot-rs-embassy/src/sendcell.rs | 40 +++++++++++++++++++++--------
 3 files changed, 34 insertions(+), 13 deletions(-)

diff --git a/src/riot-rs-embassy/src/lib.rs b/src/riot-rs-embassy/src/lib.rs
index 56380d486..9c4fd2bac 100644
--- a/src/riot-rs-embassy/src/lib.rs
+++ b/src/riot-rs-embassy/src/lib.rs
@@ -173,7 +173,7 @@ async fn init_task(mut peripherals: arch::OptionalPeripherals) {
         spawner.spawn(network::net_task(stack)).unwrap();
 
         if STACK
-            .lock(|c| c.set(SendCell::new(stack, &spawner)))
+            .lock(|c| c.set(SendCell::new(stack, spawner)))
             .is_err()
         {
             unreachable!();
diff --git a/src/riot-rs-embassy/src/network.rs b/src/riot-rs-embassy/src/network.rs
index 46c98585e..19e68fd8e 100644
--- a/src/riot-rs-embassy/src/network.rs
+++ b/src/riot-rs-embassy/src/network.rs
@@ -1,5 +1,6 @@
 use core::cell::OnceCell;
 
+use embassy_executor::Spawner;
 use embassy_net::Stack;
 use embassy_sync::blocking_mutex::CriticalSectionMutex;
 
@@ -15,8 +16,8 @@ pub(crate) static STACK: CriticalSectionMutex<OnceCell<SendCell<&'static Network
     CriticalSectionMutex::new(OnceCell::new());
 
 pub async fn network_stack() -> Option<&'static NetworkStack> {
-    let spawner = crate::Spawner::for_current_executor().await;
-    STACK.lock(|cell| cell.get().map(|x| *x.get(&spawner).unwrap()))
+    let spawner = Spawner::for_current_executor().await;
+    STACK.lock(|cell| cell.get().map(|x| *x.get(spawner).unwrap()))
 }
 
 #[embassy_executor::task]
diff --git a/src/riot-rs-embassy/src/sendcell.rs b/src/riot-rs-embassy/src/sendcell.rs
index a740829f9..662e8a872 100644
--- a/src/riot-rs-embassy/src/sendcell.rs
+++ b/src/riot-rs-embassy/src/sendcell.rs
@@ -1,8 +1,9 @@
-//! Pass non-Send objects around on same executor
+//! Pass non-Send objects around on same executor.
 //!
 //! This module provides `SendCell`, a structure that allows passing around
 //! non-Send objects from one async task to another, if they are on the same
 //! executor. This is allowed because embassy executors are single threaded.
+//!
 //! `SendCell` checks for the correct executor *at runtime*.
 
 use embassy_executor::Spawner;
@@ -11,11 +12,18 @@ use embassy_executor::Spawner;
 // SendCell guarantees at runtime that its content stays on the same embassy
 // executor. Those are single threaded, so it is guaranteed that the content
 // stays on the same thread.
+// While `SendCell::get()` allows passing any `Spawner` object, those are `!Send`,
+// thus they are guaranteed to be for the current Executor.
 unsafe impl<T> Send for SendCell<T> {}
 
 /// A cell that allows sending of non-Send types *if they stay on the same executor*.
 ///
-/// This is checked *at runtime*.
+/// This is *checked at runtime*.
+///
+/// Both `new()` and `get()` have async versions (`new_async()` and `get_async()`) that get a
+/// handle for the current `Spawner` themselves. They internally call the non-async versions. Use
+/// the sync versions if a `Spawner` object is available or the async versions cannot be used,
+/// e.g., in closures. Otherwise, the async versions are also fine.
 #[derive(Debug)]
 pub struct SendCell<T> {
     executor_id: usize,
@@ -23,24 +31,36 @@ pub struct SendCell<T> {
 }
 
 impl<T> SendCell<T> {
-    /// Create a new `SendCell`
-    ///
-    /// The `spawner` argument *must* point to the current executor.
-    pub fn new(inner: T, spawner: &Spawner) -> Self {
+    /// Creates a new `SendCell`.
+    pub fn new(inner: T, spawner: Spawner) -> Self {
         Self {
             executor_id: spawner.executor_id(),
             inner,
         }
     }
 
-    /// Get content of this `SendCell`
-    ///
-    /// The `spawner` argument *must* point to the current executor.
-    pub fn get(&self, spawner: &Spawner) -> Option<&T> {
+    /// Gets the contents of this `SendCell`.
+    pub fn get(&self, spawner: Spawner) -> Option<&T> {
         if spawner.executor_id() == self.executor_id {
             Some(&self.inner)
         } else {
             None
         }
     }
+
+    /// Creates a new `SendCell` (async version).
+    ///
+    /// Despite being async, this function never blocks/yields, it returns instantly.
+    pub async fn new_async(inner: T) -> Self {
+        let spawner = Spawner::for_current_executor().await;
+        SendCell::new(inner, spawner)
+    }
+
+    /// Gets the contents of this `SendCell` (async version).
+    ///
+    /// Despite being async, this function never blocks/yields, it returns instantly.
+    pub async fn get_async(&self) -> Option<&T> {
+        let spawner = Spawner::for_current_executor().await;
+        self.get(spawner)
+    }
 }