From 293f54d13406850d24d1226eb77989f4fa8db9f4 Mon Sep 17 00:00:00 2001
From: Dario Nieuwenhuis <dirbaio@dirbaio.net>
Date: Wed, 27 Apr 2022 03:23:54 +0200
Subject: [PATCH] executor: add raw::TaskPool.

This simplifies the macro code a bit.
---
 embassy-macros/src/macros/task.rs |  9 ++---
 embassy/src/executor/raw/mod.rs   | 56 +++++++++++++++++++++----------
 2 files changed, 42 insertions(+), 23 deletions(-)

diff --git a/embassy-macros/src/macros/task.rs b/embassy-macros/src/macros/task.rs
index 396ce18f..c450982c 100644
--- a/embassy-macros/src/macros/task.rs
+++ b/embassy-macros/src/macros/task.rs
@@ -76,14 +76,11 @@ pub fn run(args: syn::AttributeArgs, f: syn::ItemFn) -> Result<TokenStream, Toke
         #visibility fn #task_ident(#fargs) -> #embassy_path::executor::SpawnToken<impl ::core::future::Future + 'static> {
             use ::core::future::Future;
             use #embassy_path::executor::SpawnToken;
-            use #embassy_path::executor::raw::TaskStorage;
+            use #embassy_path::executor::raw::TaskPool;
 
             type Fut = impl Future + 'static;
 
-            #[allow(clippy::declare_interior_mutable_const)]
-            const NEW_TS: TaskStorage<Fut> = TaskStorage::new();
-
-            static POOL: [TaskStorage<Fut>; #pool_size] = [NEW_TS; #pool_size];
+            static POOL: TaskPool<Fut, #pool_size> = TaskPool::new();
 
             // Opaque type laundering, to obscure its origin!
             // Workaround for "opaque type's hidden type cannot be another opaque type from the same scope"
@@ -92,7 +89,7 @@ pub fn run(args: syn::AttributeArgs, f: syn::ItemFn) -> Result<TokenStream, Toke
                 token
             }
 
-            launder_tait(unsafe { TaskStorage::spawn_pool(&POOL, move || #task_inner_ident(#(#arg_names,)*)) })
+            launder_tait(POOL.spawn(move || #task_inner_ident(#(#arg_names,)*)))
         }
     };
 
diff --git a/embassy/src/executor/raw/mod.rs b/embassy/src/executor/raw/mod.rs
index 3ae52ae3..5af35d86 100644
--- a/embassy/src/executor/raw/mod.rs
+++ b/embassy/src/executor/raw/mod.rs
@@ -121,7 +121,7 @@ impl TaskHeader {
 /// the memory for the task is allocated: on the stack, or on the heap with e.g. `Box::leak`, etc.
 
 // repr(C) is needed to guarantee that the Task is located at offset 0
-// This makes it safe to cast between Task and Task pointers.
+// This makes it safe to cast between TaskHeader and TaskStorage pointers.
 #[repr(C)]
 pub struct TaskStorage<F: Future + 'static> {
     raw: TaskHeader,
@@ -129,6 +129,9 @@ pub struct TaskStorage<F: Future + 'static> {
 }
 
 impl<F: Future + 'static> TaskStorage<F> {
+    #[cfg(feature = "nightly")]
+    const NEW: Self = Self::new();
+
     /// Create a new TaskStorage, in not-spawned state.
     #[cfg(feature = "nightly")]
     pub const fn new() -> Self {
@@ -147,22 +150,6 @@ impl<F: Future + 'static> TaskStorage<F> {
         }
     }
 
-    /// Try to spawn a task in a pool.
-    ///
-    /// See [`Self::spawn()`] for details.
-    ///
-    /// This will loop over the pool and spawn the task in the first storage that
-    /// is currently free. If none is free,
-    pub fn spawn_pool(pool: &'static [Self], future: impl FnOnce() -> F) -> SpawnToken<F> {
-        for task in pool {
-            if task.spawn_allocate() {
-                return unsafe { task.spawn_initialize(future) };
-            }
-        }
-
-        SpawnToken::new_failed()
-    }
-
     /// Try to spawn the task.
     ///
     /// The `future` closure constructs the future. It's only called if spawning is
@@ -222,6 +209,41 @@ impl<F: Future + 'static> TaskStorage<F> {
 
 unsafe impl<F: Future + 'static> Sync for TaskStorage<F> {}
 
+/// Raw storage that can hold up to N tasks of the same type.
+///
+/// This is essentially a `[TaskStorage<F>; N]`.
+#[cfg(feature = "nightly")]
+pub struct TaskPool<F: Future + 'static, const N: usize> {
+    pool: [TaskStorage<F>; N],
+}
+
+#[cfg(feature = "nightly")]
+impl<F: Future + 'static, const N: usize> TaskPool<F, N> {
+    /// Create a new TaskPool, with all tasks in non-spawned state.
+    pub const fn new() -> Self {
+        Self {
+            pool: [TaskStorage::NEW; N],
+        }
+    }
+
+    /// Try to spawn a task in the pool.
+    ///
+    /// See [`TaskStorage::spawn()`] for details.
+    ///
+    /// This will loop over the pool and spawn the task in the first storage that
+    /// is currently free. If none is free, a "poisoned" SpawnToken is returned,
+    /// which will cause [`Executor::spawn()`] to return the error.
+    pub fn spawn(&'static self, future: impl FnOnce() -> F) -> SpawnToken<F> {
+        for task in &self.pool {
+            if task.spawn_allocate() {
+                return unsafe { task.spawn_initialize(future) };
+            }
+        }
+
+        SpawnToken::new_failed()
+    }
+}
+
 /// Raw executor.
 ///
 /// This is the core of the Embassy executor. It is low-level, requiring manual