From 8c2da193b88488f0995b6d83d57b7d5d61e14ffe Mon Sep 17 00:00:00 2001
From: Dario Nieuwenhuis <dirbaio@dirbaio.net>
Date: Thu, 18 Mar 2021 00:38:27 +0100
Subject: [PATCH] Add spawn/spawn_pool APIs to Task

---
 embassy-macros/src/lib.rs           |  2 +-
 embassy/src/executor/mod.rs         |  5 ---
 embassy/src/executor/raw.rs         | 49 +++++++++++++++++++----------
 embassy/src/executor/timer_queue.rs |  2 +-
 embassy/src/interrupt.rs            |  2 +-
 5 files changed, 36 insertions(+), 24 deletions(-)

diff --git a/embassy-macros/src/lib.rs b/embassy-macros/src/lib.rs
index f207497d5..0bc5f0035 100644
--- a/embassy-macros/src/lib.rs
+++ b/embassy-macros/src/lib.rs
@@ -105,7 +105,7 @@ pub fn task(args: TokenStream, item: TokenStream) -> TokenStream {
             type F = #impl_ty;
             const NEW_TASK: Task<F> = Task::new();
             static POOL: [Task<F>; #pool_size] = [NEW_TASK; #pool_size];
-            unsafe { Task::spawn(&POOL, move || task(#arg_names)) }
+            unsafe { Task::spawn_pool(&POOL, move || task(#arg_names)) }
         }
     };
     result.into()
diff --git a/embassy/src/executor/mod.rs b/embassy/src/executor/mod.rs
index 8b23264de..fedff73e2 100644
--- a/embassy/src/executor/mod.rs
+++ b/embassy/src/executor/mod.rs
@@ -1,11 +1,7 @@
 pub use embassy_macros::task;
 
-use atomic_polyfill::Ordering;
-use core::future::Future;
 use core::marker::PhantomData;
-use core::pin::Pin;
 use core::ptr::NonNull;
-use core::task::{Context, Poll};
 use core::{mem, ptr};
 
 pub mod raw;
@@ -15,7 +11,6 @@ mod timer_queue;
 mod util;
 mod waker;
 
-use self::util::UninitCell;
 use crate::fmt::panic;
 use crate::interrupt::{Interrupt, InterruptExt};
 use crate::time::Alarm;
diff --git a/embassy/src/executor/raw.rs b/embassy/src/executor/raw.rs
index edc6d8053..7e981b084 100644
--- a/embassy/src/executor/raw.rs
+++ b/embassy/src/executor/raw.rs
@@ -87,23 +87,10 @@ impl<F: Future + 'static> Task<F> {
         }
     }
 
-    pub unsafe fn spawn(pool: &'static [Self], future: impl FnOnce() -> F) -> SpawnToken<F> {
+    pub fn spawn_pool(pool: &'static [Self], future: impl FnOnce() -> F) -> SpawnToken<F> {
         for task in pool {
-            let state = STATE_SPAWNED | STATE_RUN_QUEUED;
-            if task
-                .raw
-                .state
-                .compare_exchange(0, state, Ordering::AcqRel, Ordering::Acquire)
-                .is_ok()
-            {
-                // Initialize the task
-                task.raw.poll_fn.write(Self::poll);
-                task.future.write(future());
-
-                return SpawnToken {
-                    raw_task: Some(NonNull::new_unchecked(&task.raw as *const TaskHeader as _)),
-                    phantom: PhantomData,
-                };
+            if task.spawn_allocate() {
+                return unsafe { task.spawn_initialize(future) };
             }
         }
 
@@ -113,6 +100,36 @@ impl<F: Future + 'static> Task<F> {
         }
     }
 
+    pub fn spawn(&'static self, future: impl FnOnce() -> F) -> SpawnToken<F> {
+        if self.spawn_allocate() {
+            unsafe { self.spawn_initialize(future) }
+        } else {
+            SpawnToken {
+                raw_task: None,
+                phantom: PhantomData,
+            }
+        }
+    }
+
+    fn spawn_allocate(&'static self) -> bool {
+        let state = STATE_SPAWNED | STATE_RUN_QUEUED;
+        self.raw
+            .state
+            .compare_exchange(0, state, Ordering::AcqRel, Ordering::Acquire)
+            .is_ok()
+    }
+
+    unsafe fn spawn_initialize(&'static self, future: impl FnOnce() -> F) -> SpawnToken<F> {
+        // Initialize the task
+        self.raw.poll_fn.write(Self::poll);
+        self.future.write(future());
+
+        return SpawnToken {
+            raw_task: Some(NonNull::new_unchecked(&self.raw as *const TaskHeader as _)),
+            phantom: PhantomData,
+        };
+    }
+
     unsafe fn poll(p: NonNull<TaskHeader>) {
         let this = &*(p.as_ptr() as *const Task<F>);
 
diff --git a/embassy/src/executor/timer_queue.rs b/embassy/src/executor/timer_queue.rs
index a6939f110..76bc27ad4 100644
--- a/embassy/src/executor/timer_queue.rs
+++ b/embassy/src/executor/timer_queue.rs
@@ -1,4 +1,4 @@
-use atomic_polyfill::{AtomicPtr, Ordering};
+use atomic_polyfill::Ordering;
 use core::cell::Cell;
 use core::cmp::min;
 use core::ptr;
diff --git a/embassy/src/interrupt.rs b/embassy/src/interrupt.rs
index 013e722e6..a4285a9fe 100644
--- a/embassy/src/interrupt.rs
+++ b/embassy/src/interrupt.rs
@@ -1,7 +1,7 @@
 use core::ptr;
 use cortex_m::peripheral::NVIC;
 
-use atomic_polyfill::{AtomicBool, AtomicPtr, Ordering};
+use atomic_polyfill::{AtomicPtr, Ordering};
 
 pub use embassy_macros::interrupt_declare as declare;
 pub use embassy_macros::interrupt_take as take;