1 files changed, 28 insertions, 16 deletions

diff --git a/embassy-executor-macros/src/macros/task.rs b/embassy-executor-macros/src/macros/task.rs
index 91d6beee8..1c5e3571d 100644
--- a/embassy-executor-macros/src/macros/task.rs
+++ b/embassy-executor-macros/src/macros/task.rs
@@ -51,7 +51,11 @@ pub fn run(args: TokenStream, item: TokenStream) -> TokenStream {
         .embassy_executor
         .unwrap_or(Expr::Verbatim(TokenStream::from_str("::embassy_executor").unwrap()));
 
-    if f.sig.asyncness.is_none() {
+    let returns_impl_trait = match &f.sig.output {
+        ReturnType::Type(_, ty) => matches!(**ty, Type::ImplTrait(_)),
+        _ => false,
+    };
+    if f.sig.asyncness.is_none() && !returns_impl_trait {
         error(&mut errors, &f.sig, "task functions must be async");
     }
     if !f.sig.generics.params.is_empty() {
@@ -66,17 +70,19 @@ pub fn run(args: TokenStream, item: TokenStream) -> TokenStream {
     if !f.sig.variadic.is_none() {
         error(&mut errors, &f.sig, "task functions must not be variadic");
     }
-    match &f.sig.output {
-        ReturnType::Default => {}
-        ReturnType::Type(_, ty) => match &**ty {
-            Type::Tuple(tuple) if tuple.elems.is_empty() => {}
-            Type::Never(_) => {}
-            _ => error(
-                &mut errors,
-                &f.sig,
-                "task functions must either not return a value, return `()` or return `!`",
-            ),
-        },
+    if f.sig.asyncness.is_some() {
+        match &f.sig.output {
+            ReturnType::Default => {}
+            ReturnType::Type(_, ty) => match &**ty {
+                Type::Tuple(tuple) if tuple.elems.is_empty() => {}
+                Type::Never(_) => {}
+                _ => error(
+                    &mut errors,
+                    &f.sig,
+                    "task functions must either not return a value, return `()` or return `!`",
+                ),
+            },
+        }
     }
 
     let mut args = Vec::new();
@@ -125,15 +131,21 @@ 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 {
-            type Fut: ::core::future::Future + 'static;
+            type Fut: ::core::future::Future<Output: #embassy_executor::_export::TaskReturnValue> + 'static;
             fn construct(#fargs) -> Self::Fut;
         }
 
         impl _EmbassyInternalTaskTrait for () {
-            type Fut = impl core::future::Future + 'static;
+            type Fut = impl core::future::Future<Output: #embassy_executor::_export::TaskReturnValue> + 'static;
             fn construct(#fargs) -> Self::Fut {
                 #task_inner_ident(#(#full_args,)*)
             }
@@ -141,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! {
@@ -158,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();