7 files changed, 109 insertions, 6 deletions

diff --git a/embassy-executor-macros/src/macros/task.rs b/embassy-executor-macros/src/macros/task.rs
index 91bf8e940..1c5e3571d 100644
--- a/embassy-executor-macros/src/macros/task.rs
+++ b/embassy-executor-macros/src/macros/task.rs
@@ -131,6 +131,12 @@ pub fn run(args: TokenStream, item: TokenStream) -> TokenStream {
         ));
     }
 
+    let spawn = if returns_impl_trait {
+        quote!(spawn)
+    } else {
+        quote!(_spawn_async_fn)
+    };
+
     #[cfg(feature = "nightly")]
     let mut task_outer_body = quote! {
         trait _EmbassyInternalTaskTrait {
@@ -147,7 +153,7 @@ pub fn run(args: TokenStream, item: TokenStream) -> TokenStream {
 
         const POOL_SIZE: usize = #pool_size;
         static POOL: #embassy_executor::raw::TaskPool<<() as _EmbassyInternalTaskTrait>::Fut, POOL_SIZE> = #embassy_executor::raw::TaskPool::new();
-        unsafe { POOL._spawn_async_fn(move || <() as _EmbassyInternalTaskTrait>::construct(#(#full_args,)*)) }
+        unsafe { POOL.#spawn(move || <() as _EmbassyInternalTaskTrait>::construct(#(#full_args,)*)) }
     };
     #[cfg(not(feature = "nightly"))]
     let mut task_outer_body = quote! {
@@ -164,7 +170,7 @@ pub fn run(args: TokenStream, item: TokenStream) -> TokenStream {
             {#embassy_executor::_export::task_pool_size::<_, _, _, POOL_SIZE>(#task_inner_ident)},
             {#embassy_executor::_export::task_pool_align::<_, _, _, POOL_SIZE>(#task_inner_ident)},
         > = unsafe { ::core::mem::transmute(#embassy_executor::_export::task_pool_new::<_, _, _, POOL_SIZE>(#task_inner_ident)) };
-        unsafe { __task_pool_get(#task_inner_ident)._spawn_async_fn(move || #task_inner_ident(#(#full_args,)*)) }
+        unsafe { __task_pool_get(#task_inner_ident).#spawn(move || #task_inner_ident(#(#full_args,)*)) }
     };
 
     let task_outer_attrs = task_inner.attrs.clone();
diff --git a/embassy-executor/tests/ui.rs b/embassy-executor/tests/ui.rs
index c4a1a601c..7757775ee 100644
--- a/embassy-executor/tests/ui.rs
+++ b/embassy-executor/tests/ui.rs
@@ -17,6 +17,8 @@ fn ui() {
     t.compile_fail("tests/ui/nonstatic_struct_elided.rs");
     t.compile_fail("tests/ui/nonstatic_struct_generic.rs");
     t.compile_fail("tests/ui/not_async.rs");
+    t.compile_fail("tests/ui/spawn_nonsend.rs");
+    t.compile_fail("tests/ui/return_impl_future_nonsend.rs");
     if rustversion::cfg!(stable) {
         // output is slightly different on nightly
         t.compile_fail("tests/ui/bad_return_impl_future.rs");
diff --git a/embassy-executor/tests/ui/return_impl_future_nonsend.rs b/embassy-executor/tests/ui/return_impl_future_nonsend.rs
new file mode 100644
index 000000000..b8c184b21
--- /dev/null
+++ b/embassy-executor/tests/ui/return_impl_future_nonsend.rs
@@ -0,0 +1,21 @@
+#![cfg_attr(feature = "nightly", feature(impl_trait_in_assoc_type))]
+
+use core::future::Future;
+
+use embassy_executor::SendSpawner;
+
+#[embassy_executor::task]
+fn task() -> impl Future<Output = ()> {
+    // runs in spawning thread
+    let non_send: *mut () = core::ptr::null_mut();
+    async move {
+        // runs in executor thread
+        println!("{}", non_send as usize);
+    }
+}
+
+fn send_spawn(s: SendSpawner) {
+    s.spawn(task()).unwrap();
+}
+
+fn main() {}
diff --git a/embassy-executor/tests/ui/return_impl_future_nonsend.stderr b/embassy-executor/tests/ui/return_impl_future_nonsend.stderr
new file mode 100644
index 000000000..8aeb9738a
--- /dev/null
+++ b/embassy-executor/tests/ui/return_impl_future_nonsend.stderr
@@ -0,0 +1,17 @@
+error: future cannot be sent between threads safely
+  --> tests/ui/return_impl_future_nonsend.rs:18:13
+   |
+18 |     s.spawn(task()).unwrap();
+   |             ^^^^^^ future created by async block is not `Send`
+   |
+   = help: within `impl Sized`, the trait `Send` is not implemented for `*mut ()`
+note: captured value is not `Send`
+  --> tests/ui/return_impl_future_nonsend.rs:13:24
+   |
+13 |         println!("{}", non_send as usize);
+   |                        ^^^^^^^^ has type `*mut ()` which is not `Send`
+note: required by a bound in `SendSpawner::spawn`
+  --> src/spawner.rs
+   |
+   |     pub fn spawn<S: Send>(&self, token: SpawnToken<S>) -> Result<(), SpawnError> {
+   |                     ^^^^ required by this bound in `SendSpawner::spawn`
diff --git a/embassy-executor/tests/ui/return_impl_send.stderr b/embassy-executor/tests/ui/return_impl_send.stderr
index cd693af2b..759be1cde 100644
--- a/embassy-executor/tests/ui/return_impl_send.stderr
+++ b/embassy-executor/tests/ui/return_impl_send.stderr
@@ -91,14 +91,14 @@ error[E0277]: `impl Send` is not a future
   | ^^^^^^^^^^^^^^^^^^^^^^^^^ `impl Send` is not a future
   |
   = help: the trait `Future` is not implemented for `impl Send`
-note: required by a bound in `TaskPool::<F, N>::_spawn_async_fn`
+note: required by a bound in `TaskPool::<F, N>::spawn`
  --> src/raw/mod.rs
   |
   | impl<F: Future + 'static, const N: usize> TaskPool<F, N> {
-  |         ^^^^^^ required by this bound in `TaskPool::<F, N>::_spawn_async_fn`
+  |         ^^^^^^ required by this bound in `TaskPool::<F, N>::spawn`
 ...
-  |     pub unsafe fn _spawn_async_fn<FutFn>(&'static self, future: FutFn) -> SpawnToken<impl Sized>
-  |                   --------------- required by a bound in this associated function
+  |     pub fn spawn(&'static self, future: impl FnOnce() -> F) -> SpawnToken<impl Sized> {
+  |            ----- required by a bound in this associated function
   = note: this error originates in the attribute macro `embassy_executor::task` (in Nightly builds, run with -Z macro-backtrace for more info)
 
 error[E0277]: task futures must resolve to `()` or `!`
diff --git a/embassy-executor/tests/ui/spawn_nonsend.rs b/embassy-executor/tests/ui/spawn_nonsend.rs
new file mode 100644
index 000000000..4c4cc7697
--- /dev/null
+++ b/embassy-executor/tests/ui/spawn_nonsend.rs
@@ -0,0 +1,16 @@
+#![cfg_attr(feature = "nightly", feature(impl_trait_in_assoc_type))]
+
+use core::future::Future;
+
+use embassy_executor::SendSpawner;
+
+#[embassy_executor::task]
+async fn task(non_send: *mut ()) {
+    println!("{}", non_send as usize);
+}
+
+fn send_spawn(s: SendSpawner) {
+    s.spawn(task(core::ptr::null_mut())).unwrap();
+}
+
+fn main() {}
diff --git a/embassy-executor/tests/ui/spawn_nonsend.stderr b/embassy-executor/tests/ui/spawn_nonsend.stderr
new file mode 100644
index 000000000..2a06c8b94
--- /dev/null
+++ b/embassy-executor/tests/ui/spawn_nonsend.stderr
@@ -0,0 +1,41 @@
+warning: unused import: `core::future::Future`
+ --> tests/ui/spawn_nonsend.rs:3:5
+  |
+3 | use core::future::Future;
+  |     ^^^^^^^^^^^^^^^^^^^^
+  |
+  = note: `#[warn(unused_imports)]` on by default
+
+error[E0277]: `*mut ()` cannot be sent between threads safely
+  --> tests/ui/spawn_nonsend.rs:13:13
+   |
+7  | #[embassy_executor::task]
+   | ------------------------- within this `impl Sized`
+...
+13 |     s.spawn(task(core::ptr::null_mut())).unwrap();
+   |       ----- ^^^^^^^^^^^^^^^^^^^^^^^^^^^ `*mut ()` cannot be sent between threads safely
+   |       |
+   |       required by a bound introduced by this call
+   |
+   = help: within `impl Sized`, the trait `Send` is not implemented for `*mut ()`
+note: required because it's used within this closure
+  --> tests/ui/spawn_nonsend.rs:7:1
+   |
+7  | #[embassy_executor::task]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^
+note: required because it appears within the type `impl Sized`
+  --> src/raw/mod.rs
+   |
+   |     pub unsafe fn _spawn_async_fn<FutFn>(&'static self, future: FutFn) -> SpawnToken<impl Sized>
+   |                                                                                      ^^^^^^^^^^
+note: required because it appears within the type `impl Sized`
+  --> tests/ui/spawn_nonsend.rs:7:1
+   |
+7  | #[embassy_executor::task]
+   | ^^^^^^^^^^^^^^^^^^^^^^^^^
+note: required by a bound in `SendSpawner::spawn`
+  --> src/spawner.rs
+   |
+   |     pub fn spawn<S: Send>(&self, token: SpawnToken<S>) -> Result<(), SpawnError> {
+   |                     ^^^^ required by this bound in `SendSpawner::spawn`
+   = note: this error originates in the attribute macro `embassy_executor::task` (in Nightly builds, run with -Z macro-backtrace for more info)