From 9e897cbea92dc9b99a65802cc4e0661919001be1 Mon Sep 17 00:00:00 2001
From: Dario Nieuwenhuis <dirbaio@dirbaio.net>
Date: Tue, 26 Apr 2022 19:08:18 +0200
Subject: [PATCH] executor: Add `Spawner::for_current_executor`.

---
 embassy/src/executor/arch/cortex_m.rs         |  3 +
 embassy/src/executor/spawner.rs               | 62 +++++++++++++++++--
 .../src/bin/self_spawn_current_executor.rs    | 24 +++++++
 3 files changed, 85 insertions(+), 4 deletions(-)
 create mode 100644 examples/nrf/src/bin/self_spawn_current_executor.rs

diff --git a/embassy/src/executor/arch/cortex_m.rs b/embassy/src/executor/arch/cortex_m.rs
index 87c9c3c8..16f29008 100644
--- a/embassy/src/executor/arch/cortex_m.rs
+++ b/embassy/src/executor/arch/cortex_m.rs
@@ -115,6 +115,9 @@ impl<I: Interrupt> InterruptExecutor<I> {
     /// different "thread" (the interrupt), so spawning tasks on it is effectively
     /// sending them.
     ///
+    /// To obtain a [`Spawner`] for this executor, use [`Spawner::for_current_executor`] from
+    /// a task running in it.
+    ///
     /// This function requires `&'static mut self`. This means you have to store the
     /// Executor instance in a place where it'll live forever and grants you mutable
     /// access. There's a few ways to do this:
diff --git a/embassy/src/executor/spawner.rs b/embassy/src/executor/spawner.rs
index d1e4ff17..e6770e29 100644
--- a/embassy/src/executor/spawner.rs
+++ b/embassy/src/executor/spawner.rs
@@ -1,6 +1,8 @@
 use core::marker::PhantomData;
 use core::mem;
 use core::ptr::NonNull;
+use core::task::Poll;
+use futures::future::poll_fn;
 
 use super::raw;
 
@@ -75,6 +77,23 @@ impl Spawner {
         }
     }
 
+    /// Get a Spawner for the current executor.
+    ///
+    /// This function is `async` just to get access to the current async
+    /// context. It returns instantly, it does not block/yield.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the current executor is not an Embassy executor.
+    pub async fn for_current_executor() -> Self {
+        poll_fn(|cx| unsafe {
+            let task = raw::task_from_waker(cx.waker());
+            let executor = (&*task.as_ptr()).executor.get();
+            Poll::Ready(Self::new(&*executor))
+        })
+        .await
+    }
+
     /// Spawn a task into an executor.
     ///
     /// You obtain the `token` by calling a task function (i.e. one marked with `#[embassy::task]`).
@@ -91,10 +110,15 @@ impl Spawner {
         }
     }
 
-    /// Used by the `embassy_macros::main!` macro to throw an error when spawn
-    /// fails. This is here to allow conditional use of `defmt::unwrap!`
-    /// without introducing a `defmt` feature in the `embassy_macros` package,
-    /// which would require use of `-Z namespaced-features`.
+    // Used by the `embassy_macros::main!` macro to throw an error when spawn
+    // fails. This is here to allow conditional use of `defmt::unwrap!`
+    // without introducing a `defmt` feature in the `embassy_macros` package,
+    // which would require use of `-Z namespaced-features`.
+    /// Spawn a task into an executor, panicking on failure.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the spawning fails.
     pub fn must_spawn<F>(&self, token: SpawnToken<F>) {
         unwrap!(self.spawn(token));
     }
@@ -125,6 +149,27 @@ unsafe impl Send for SendSpawner {}
 unsafe impl Sync for SendSpawner {}
 
 impl SendSpawner {
+    pub(crate) fn new(executor: &'static raw::Executor) -> Self {
+        Self { executor }
+    }
+
+    /// Get a Spawner for the current executor.
+    ///
+    /// This function is `async` just to get access to the current async
+    /// context. It returns instantly, it does not block/yield.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the current executor is not an Embassy executor.
+    pub async fn for_current_executor() -> Self {
+        poll_fn(|cx| unsafe {
+            let task = raw::task_from_waker(cx.waker());
+            let executor = (&*task.as_ptr()).executor.get();
+            Poll::Ready(Self::new(&*executor))
+        })
+        .await
+    }
+
     /// Spawn a task into an executor.
     ///
     /// You obtain the `token` by calling a task function (i.e. one marked with `#[embassy::task]`).
@@ -140,4 +185,13 @@ impl SendSpawner {
             None => Err(SpawnError::Busy),
         }
     }
+
+    /// Spawn a task into an executor, panicking on failure.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the spawning fails.
+    pub fn must_spawn<F: Send>(&self, token: SpawnToken<F>) {
+        unwrap!(self.spawn(token));
+    }
 }
diff --git a/examples/nrf/src/bin/self_spawn_current_executor.rs b/examples/nrf/src/bin/self_spawn_current_executor.rs
new file mode 100644
index 00000000..4850d295
--- /dev/null
+++ b/examples/nrf/src/bin/self_spawn_current_executor.rs
@@ -0,0 +1,24 @@
+#![no_std]
+#![no_main]
+#![feature(type_alias_impl_trait)]
+
+use defmt::{info, unwrap};
+use embassy::executor::Spawner;
+use embassy::time::{Duration, Timer};
+use embassy_nrf::Peripherals;
+
+use defmt_rtt as _; // global logger
+use panic_probe as _;
+
+#[embassy::task(pool_size = 2)]
+async fn my_task(n: u32) {
+    Timer::after(Duration::from_secs(1)).await;
+    info!("Spawning self! {}", n);
+    unwrap!(Spawner::for_current_executor().await.spawn(my_task(n + 1)));
+}
+
+#[embassy::main]
+async fn main(spawner: Spawner, _p: Peripherals) {
+    info!("Hello World!");
+    unwrap!(spawner.spawn(my_task(0)));
+}