1 files changed, 102 insertions, 49 deletions

diff --git a/embassy-executor-macros/src/macros/task.rs b/embassy-executor-macros/src/macros/task.rs
index 96c6267b2..0404dba64 100644
--- a/embassy-executor-macros/src/macros/task.rs
+++ b/embassy-executor-macros/src/macros/task.rs
@@ -2,47 +2,68 @@ use darling::export::NestedMeta;
 use darling::FromMeta;
 use proc_macro2::{Span, TokenStream};
 use quote::{format_ident, quote};
-use syn::{parse_quote, Expr, ExprLit, ItemFn, Lit, LitInt, ReturnType, Type};
+use syn::visit::Visit;
+use syn::{Expr, ExprLit, Lit, LitInt, ReturnType, Type};
 
-use crate::util::ctxt::Ctxt;
+use crate::util::*;
 
-#[derive(Debug, FromMeta)]
+#[derive(Debug, FromMeta, Default)]
 struct Args {
     #[darling(default)]
     pool_size: Option<syn::Expr>,
 }
 
-pub fn run(args: &[NestedMeta], f: syn::ItemFn) -> Result<TokenStream, TokenStream> {
-    let args = Args::from_list(args).map_err(|e| e.write_errors())?;
+pub fn run(args: TokenStream, item: TokenStream) -> TokenStream {
+    let mut errors = TokenStream::new();
+
+    // If any of the steps for this macro fail, we still want to expand to an item that is as close
+    // to the expected output as possible. This helps out IDEs such that completions and other
+    // related features keep working.
+    let f: ItemFn = match syn::parse2(item.clone()) {
+        Ok(x) => x,
+        Err(e) => return token_stream_with_error(item, e),
+    };
+
+    let args = match NestedMeta::parse_meta_list(args) {
+        Ok(x) => x,
+        Err(e) => return token_stream_with_error(item, e),
+    };
+
+    let args = match Args::from_list(&args) {
+        Ok(x) => x,
+        Err(e) => {
+            errors.extend(e.write_errors());
+            Args::default()
+        }
+    };
 
     let pool_size = args.pool_size.unwrap_or(Expr::Lit(ExprLit {
         attrs: vec![],
         lit: Lit::Int(LitInt::new("1", Span::call_site())),
     }));
 
-    let ctxt = Ctxt::new();
-
     if f.sig.asyncness.is_none() {
-        ctxt.error_spanned_by(&f.sig, "task functions must be async");
+        error(&mut errors, &f.sig, "task functions must be async");
     }
     if !f.sig.generics.params.is_empty() {
-        ctxt.error_spanned_by(&f.sig, "task functions must not be generic");
+        error(&mut errors, &f.sig, "task functions must not be generic");
     }
     if !f.sig.generics.where_clause.is_none() {
-        ctxt.error_spanned_by(&f.sig, "task functions must not have `where` clauses");
+        error(&mut errors, &f.sig, "task functions must not have `where` clauses");
     }
     if !f.sig.abi.is_none() {
-        ctxt.error_spanned_by(&f.sig, "task functions must not have an ABI qualifier");
+        error(&mut errors, &f.sig, "task functions must not have an ABI qualifier");
     }
     if !f.sig.variadic.is_none() {
-        ctxt.error_spanned_by(&f.sig, "task functions must not be variadic");
+        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(_) => {}
-            _ => ctxt.error_spanned_by(
+            _ => error(
+                &mut errors,
                 &f.sig,
                 "task functions must either not return a value, return `()` or return `!`",
             ),
@@ -55,26 +76,31 @@ pub fn run(args: &[NestedMeta], f: syn::ItemFn) -> Result<TokenStream, TokenStre
     for arg in fargs.iter_mut() {
         match arg {
             syn::FnArg::Receiver(_) => {
-                ctxt.error_spanned_by(arg, "task functions must not have receiver arguments");
+                error(&mut errors, arg, "task functions must not have receiver arguments");
             }
-            syn::FnArg::Typed(t) => match t.pat.as_mut() {
-                syn::Pat::Ident(id) => {
-                    id.mutability = None;
-                    args.push((id.clone(), t.attrs.clone()));
+            syn::FnArg::Typed(t) => {
+                check_arg_ty(&mut errors, &t.ty);
+                match t.pat.as_mut() {
+                    syn::Pat::Ident(id) => {
+                        id.mutability = None;
+                        args.push((id.clone(), t.attrs.clone()));
+                    }
+                    _ => {
+                        error(
+                            &mut errors,
+                            arg,
+                            "pattern matching in task arguments is not yet supported",
+                        );
+                    }
                 }
-                _ => {
-                    ctxt.error_spanned_by(arg, "pattern matching in task arguments is not yet supported");
-                }
-            },
+            }
         }
     }
 
-    ctxt.check()?;
-
     let task_ident = f.sig.ident.clone();
     let task_inner_ident = format_ident!("__{}_task", task_ident);
 
-    let mut task_inner = f;
+    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();
@@ -91,35 +117,43 @@ pub fn run(args: &[NestedMeta], f: syn::ItemFn) -> Result<TokenStream, TokenStre
     }
 
     #[cfg(feature = "nightly")]
-    let mut task_outer: ItemFn = parse_quote! {
-        #visibility fn #task_ident(#fargs) -> ::embassy_executor::SpawnToken<impl Sized> {
-            trait _EmbassyInternalTaskTrait {
-                type Fut: ::core::future::Future + 'static;
-                fn construct(#fargs) -> Self::Fut;
-            }
+    let mut task_outer_body = quote! {
+        trait _EmbassyInternalTaskTrait {
+            type Fut: ::core::future::Future + 'static;
+            fn construct(#fargs) -> Self::Fut;
+        }
 
-            impl _EmbassyInternalTaskTrait for () {
-                type Fut = impl core::future::Future + 'static;
-                fn construct(#fargs) -> Self::Fut {
-                    #task_inner_ident(#(#full_args,)*)
-                }
+        impl _EmbassyInternalTaskTrait for () {
+            type Fut = impl core::future::Future + 'static;
+            fn construct(#fargs) -> Self::Fut {
+                #task_inner_ident(#(#full_args,)*)
             }
-
-            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,)*)) }
         }
+
+        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,)*)) }
     };
     #[cfg(not(feature = "nightly"))]
-    let mut task_outer: ItemFn = parse_quote! {
-        #visibility fn #task_ident(#fargs) -> ::embassy_executor::SpawnToken<impl Sized> {
-            const POOL_SIZE: usize = #pool_size;
-            static POOL: ::embassy_executor::_export::TaskPoolRef = ::embassy_executor::_export::TaskPoolRef::new();
-            unsafe { POOL.get::<_, POOL_SIZE>()._spawn_async_fn(move || #task_inner_ident(#(#full_args,)*)) }
-        }
+    let mut task_outer_body = quote! {
+        const POOL_SIZE: usize = #pool_size;
+        static POOL: ::embassy_executor::_export::TaskPoolRef = ::embassy_executor::_export::TaskPoolRef::new();
+        unsafe { POOL.get::<_, POOL_SIZE>()._spawn_async_fn(move || #task_inner_ident(#(#full_args,)*)) }
     };
 
-    task_outer.attrs.append(&mut task_inner.attrs.clone());
+    let task_outer_attrs = task_inner.attrs.clone();
+
+    if !errors.is_empty() {
+        task_outer_body = quote! {
+            #![allow(unused_variables, unreachable_code)]
+            let _x: ::embassy_executor::SpawnToken<()> = ::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.
@@ -129,8 +163,27 @@ pub fn run(args: &[NestedMeta], f: syn::ItemFn) -> Result<TokenStream, TokenStre
         #[doc(hidden)]
         #task_inner
 
-        #task_outer
+        #(#task_outer_attrs)*
+        #visibility fn #task_ident #generics (#fargs) -> ::embassy_executor::SpawnToken<impl Sized> #where_clause{
+            #task_outer_body
+        }
+
+        #errors
     };
 
-    Ok(result)
+    result
+}
+
+fn check_arg_ty(errors: &mut TokenStream, ty: &Type) {
+    struct Visitor<'a> {
+        errors: &'a mut TokenStream,
+    }
+
+    impl<'a, 'ast> Visit<'ast> for Visitor<'a> {
+        fn visit_type_impl_trait(&mut self, i: &'ast syn::TypeImplTrait) {
+            error(self.errors, i, "`impl Trait` is not allowed in task arguments. It is syntax sugar for generics, and tasks can't be generic.");
+        }
+    }
+
+    Visit::visit_type(&mut Visitor { errors }, ty);
 }