2 files changed, 116 insertions, 34 deletions

diff --git a/embassy-executor-macros/src/macros/main.rs b/embassy-executor-macros/src/macros/main.rs
index 95722b19b..dc470e51c 100644
--- a/embassy-executor-macros/src/macros/main.rs
+++ b/embassy-executor-macros/src/macros/main.rs
@@ -37,6 +37,12 @@ pub static ARCH_CORTEX_M: Arch = Arch {
     executor_required: false,
 };
 
+pub static ARCH_CORTEX_AR: Arch = Arch {
+    default_entry: None,
+    flavor: Flavor::Standard,
+    executor_required: false,
+};
+
 pub static ARCH_SPIN: Arch = Arch {
     default_entry: None,
     flavor: Flavor::Standard,
@@ -164,6 +170,14 @@ For example: `#[embassy_executor::main(entry = ..., executor = \"some_crate::Exe
     let f_body = f.body;
     let out = &f.sig.output;
 
+    let name_main_task = if cfg!(feature = "metadata-name") {
+        quote!(
+            main_task.metadata().set_name("main\0");
+        )
+    } else {
+        quote!()
+    };
+
     let (main_ret, mut main_body) = match arch.flavor {
         Flavor::Standard => (
             quote!(!),
@@ -175,7 +189,9 @@ For example: `#[embassy_executor::main(entry = ..., executor = \"some_crate::Exe
                 let mut executor = #executor::new();
                 let executor = unsafe { __make_static(&mut executor) };
                 executor.run(|spawner| {
-                    spawner.must_spawn(__embassy_main(spawner));
+                    let main_task = __embassy_main(spawner).unwrap();
+                    #name_main_task
+                    spawner.spawn(main_task);
                 })
             },
         ),
@@ -185,7 +201,9 @@ For example: `#[embassy_executor::main(entry = ..., executor = \"some_crate::Exe
                 let executor = ::std::boxed::Box::leak(::std::boxed::Box::new(#executor::new()));
 
                 executor.start(|spawner| {
-                    spawner.must_spawn(__embassy_main(spawner));
+                    let main_task = __embassy_main(spawner).unwrap();
+                    #name_main_task
+                    spawner.spawn(main_task);
                 });
 
                 Ok(())
diff --git a/embassy-executor-macros/src/macros/task.rs b/embassy-executor-macros/src/macros/task.rs
index 91d6beee8..755948882 100644
--- a/embassy-executor-macros/src/macros/task.rs
+++ b/embassy-executor-macros/src/macros/task.rs
@@ -5,7 +5,7 @@ use darling::FromMeta;
 use proc_macro2::{Span, TokenStream};
 use quote::{format_ident, quote};
 use syn::visit::{self, Visit};
-use syn::{Expr, ExprLit, Lit, LitInt, ReturnType, Type};
+use syn::{Expr, ExprLit, Lit, LitInt, ReturnType, Type, Visibility};
 
 use crate::util::*;
 
@@ -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();
@@ -106,13 +112,11 @@ pub fn run(args: TokenStream, item: TokenStream) -> TokenStream {
         }
     }
 
-    let task_ident = f.sig.ident.clone();
-    let task_inner_ident = format_ident!("__{}_task", task_ident);
-
-    let mut task_inner = f.clone();
-    let visibility = task_inner.vis.clone();
-    task_inner.vis = syn::Visibility::Inherited;
-    task_inner.sig.ident = task_inner_ident.clone();
+    // Copy the generics + where clause to avoid more spurious errors.
+    let generics = &f.sig.generics;
+    let where_clause = &f.sig.generics.where_clause;
+    let unsafety = &f.sig.unsafety;
+    let visibility = &f.vis;
 
     // assemble the original input arguments,
     // including any attributes that may have
@@ -125,15 +129,79 @@ pub fn run(args: TokenStream, item: TokenStream) -> TokenStream {
         ));
     }
 
+    let task_ident = f.sig.ident.clone();
+    let task_inner_ident = format_ident!("__{}_task", task_ident);
+
+    let task_inner_future_output = match &f.sig.output {
+        ReturnType::Default => quote! {-> impl ::core::future::Future<Output = ()>},
+        // Special case the never type since we can't stuff it into a `impl Future<Output = !>`
+        ReturnType::Type(arrow, maybe_never)
+            if f.sig.asyncness.is_some() && matches!(**maybe_never, Type::Never(_)) =>
+        {
+            quote! {
+                #arrow impl ::core::future::Future<Output=#embassy_executor::_export::Never>
+            }
+        }
+        ReturnType::Type(arrow, maybe_never) if matches!(**maybe_never, Type::Never(_)) => quote! {
+            #arrow #maybe_never
+        },
+        // Grab the arrow span, why not
+        ReturnType::Type(arrow, typ) if f.sig.asyncness.is_some() => quote! {
+            #arrow impl ::core::future::Future<Output = #typ>
+        },
+        // We assume that if `f` isn't async, it must return `-> impl Future<...>`
+        // This is checked using traits later
+        ReturnType::Type(arrow, typ) => quote! {
+            #arrow #typ
+        },
+    };
+
+    // We have to rename the function since it might be recursive;
+    let mut task_inner_function = f.clone();
+    let task_inner_function_ident = format_ident!("__{}_task_inner_function", task_ident);
+    task_inner_function.sig.ident = task_inner_function_ident.clone();
+    task_inner_function.vis = Visibility::Inherited;
+
+    let task_inner_body = if errors.is_empty() {
+        quote! {
+            #task_inner_function
+
+            // SAFETY: All the preconditions to `#task_ident` apply to
+            //         all contexts `#task_inner_ident` is called in
+            #unsafety {
+                #task_inner_function_ident(#(#full_args,)*)
+            }
+        }
+    } else {
+        quote! {
+            async {::core::todo!()}
+        }
+    };
+
+    let task_inner = quote! {
+        #visibility fn #task_inner_ident #generics (#fargs)
+        #task_inner_future_output
+        #where_clause
+        {
+            #task_inner_body
+        }
+    };
+
+    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 +209,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,23 +226,19 @@ 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();
+    let task_outer_attrs = &f.attrs;
 
     if !errors.is_empty() {
         task_outer_body = quote! {
             #![allow(unused_variables, unreachable_code)]
-            let _x: #embassy_executor::SpawnToken<()> = ::core::todo!();
+            let _x: ::core::result::Result<#embassy_executor::SpawnToken<()>, #embassy_executor::SpawnError> = ::core::todo!();
             _x
         };
     }
 
-    // Copy the generics + where clause to avoid more spurious errors.
-    let generics = &f.sig.generics;
-    let where_clause = &f.sig.generics.where_clause;
-
     let result = quote! {
         // This is the user's task function, renamed.
         // We put it outside the #task_ident fn below, because otherwise
@@ -184,7 +248,7 @@ pub fn run(args: TokenStream, item: TokenStream) -> TokenStream {
         #task_inner
 
         #(#task_outer_attrs)*
-        #visibility fn #task_ident #generics (#fargs) -> #embassy_executor::SpawnToken<impl Sized> #where_clause{
+        #visibility #unsafety fn #task_ident #generics (#fargs) -> ::core::result::Result<#embassy_executor::SpawnToken<impl Sized>, #embassy_executor::SpawnError> #where_clause{
             #task_outer_body
         }
 
@@ -201,7 +265,7 @@ fn check_arg_ty(errors: &mut TokenStream, ty: &Type) {
 
     impl<'a, 'ast> Visit<'ast> for Visitor<'a> {
         fn visit_type_reference(&mut self, i: &'ast syn::TypeReference) {
-            // only check for elided lifetime here. If not elided, it's checked by `visit_lifetime`.
+            // Only check for elided lifetime here. If not elided, it's checked by `visit_lifetime`.
             if i.lifetime.is_none() {
                 error(
                     self.errors,