executor: fix unsoundness due to `impl Trait`, improve macro error handling. (#3425)

* executor-macros: don't parse function bodies.

* executor-macros: refactor for better recovery and ide-friendliness on errors.

* executor-macros: disallow `impl Trait` in task arguments.

Fixes #3420

* Fix example using `impl Trait` in tasks.

2 files changed, 221 insertions, 170 deletions

diff --git a/embassy-executor-macros/src/macros/main.rs b/embassy-executor-macros/src/macros/main.rs
index 66a3965d0..14b1d9de2 100644
--- a/embassy-executor-macros/src/macros/main.rs
+++ b/embassy-executor-macros/src/macros/main.rs
@@ -1,152 +1,107 @@
+use std::str::FromStr;
+
 use darling::export::NestedMeta;
-use darling::{Error, FromMeta};
+use darling::FromMeta;
 use proc_macro2::TokenStream;
 use quote::quote;
-use syn::{Expr, ReturnType, Type};
+use syn::{ReturnType, Type};
+
+use crate::util::*;
+
+enum Flavor {
+    Standard,
+    Wasm,
+}
 
-use crate::util::ctxt::Ctxt;
+pub(crate) struct Arch {
+    default_entry: Option<&'static str>,
+    flavor: Flavor,
+}
 
-#[derive(Debug, FromMeta)]
+pub static ARCH_AVR: Arch = Arch {
+    default_entry: Some("avr_device::entry"),
+    flavor: Flavor::Standard,
+};
+
+pub static ARCH_RISCV: Arch = Arch {
+    default_entry: Some("riscv_rt::entry"),
+    flavor: Flavor::Standard,
+};
+
+pub static ARCH_CORTEX_M: Arch = Arch {
+    default_entry: Some("cortex_m_rt::entry"),
+    flavor: Flavor::Standard,
+};
+
+pub static ARCH_SPIN: Arch = Arch {
+    default_entry: None,
+    flavor: Flavor::Standard,
+};
+
+pub static ARCH_STD: Arch = Arch {
+    default_entry: None,
+    flavor: Flavor::Standard,
+};
+
+pub static ARCH_WASM: Arch = Arch {
+    default_entry: Some("wasm_bindgen::prelude::wasm_bindgen(start)"),
+    flavor: Flavor::Wasm,
+};
+
+#[derive(Debug, FromMeta, Default)]
 struct Args {
     #[darling(default)]
     entry: Option<String>,
 }
 
-pub fn avr() -> TokenStream {
-    quote! {
-        #[avr_device::entry]
-        fn main() -> ! {
-            let mut executor = ::embassy_executor::Executor::new();
-            let executor = unsafe { __make_static(&mut executor) };
+pub fn run(args: TokenStream, item: TokenStream, arch: &Arch) -> TokenStream {
+    let mut errors = TokenStream::new();
 
-            executor.run(|spawner| {
-                spawner.must_spawn(__embassy_main(spawner));
-            })
-        }
-    }
-}
-
-pub fn riscv(args: &[NestedMeta]) -> TokenStream {
-    let maybe_entry = match Args::from_list(args) {
-        Ok(args) => args.entry,
-        Err(e) => return e.write_errors(),
+    // 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 entry = maybe_entry.unwrap_or("riscv_rt::entry".into());
-    let entry = match Expr::from_string(&entry) {
-        Ok(expr) => expr,
-        Err(e) => return e.write_errors(),
+    let args = match NestedMeta::parse_meta_list(args) {
+        Ok(x) => x,
+        Err(e) => return token_stream_with_error(item, e),
     };
 
-    quote! {
-        #[#entry]
-        fn main() -> ! {
-            let mut executor = ::embassy_executor::Executor::new();
-            let executor = unsafe { __make_static(&mut executor) };
-            executor.run(|spawner| {
-                spawner.must_spawn(__embassy_main(spawner));
-            })
+    let args = match Args::from_list(&args) {
+        Ok(x) => x,
+        Err(e) => {
+            errors.extend(e.write_errors());
+            Args::default()
         }
-    }
-}
-
-pub fn spin(args: &[NestedMeta]) -> TokenStream {
-    let maybe_entry = match Args::from_list(args) {
-        Ok(args) => args.entry,
-        Err(e) => return e.write_errors(),
-    };
-
-    let entry = match maybe_entry {
-        Some(str) => str,
-        None => return Error::missing_field("entry").write_errors(),
     };
-    let entry = match Expr::from_string(&entry) {
-        Ok(expr) => expr,
-        Err(e) => return e.write_errors(),
-    };
-
-    quote! {
-        #[#entry]
-        fn main() -> ! {
-            let mut executor = ::embassy_executor::Executor::new();
-            let executor = unsafe { __make_static(&mut executor) };
-            executor.run(|spawner| {
-                spawner.must_spawn(__embassy_main(spawner));
-            })
-        }
-    }
-}
-
-pub fn cortex_m() -> TokenStream {
-    quote! {
-        #[cortex_m_rt::entry]
-        fn main() -> ! {
-            let mut executor = ::embassy_executor::Executor::new();
-            let executor = unsafe { __make_static(&mut executor) };
-            executor.run(|spawner| {
-                spawner.must_spawn(__embassy_main(spawner));
-            })
-        }
-    }
-}
-
-pub fn wasm() -> TokenStream {
-    quote! {
-        #[wasm_bindgen::prelude::wasm_bindgen(start)]
-        pub fn main() -> Result<(), wasm_bindgen::JsValue> {
-            let executor = ::std::boxed::Box::leak(::std::boxed::Box::new(::embassy_executor::Executor::new()));
-
-            executor.start(|spawner| {
-                spawner.must_spawn(__embassy_main(spawner));
-            });
-
-            Ok(())
-        }
-    }
-}
-
-pub fn std() -> TokenStream {
-    quote! {
-        fn main() -> ! {
-            let mut executor = ::embassy_executor::Executor::new();
-            let executor = unsafe { __make_static(&mut executor) };
-
-            executor.run(|spawner| {
-                spawner.must_spawn(__embassy_main(spawner));
-            })
-        }
-    }
-}
-
-pub fn run(args: &[NestedMeta], f: syn::ItemFn, main: TokenStream) -> Result<TokenStream, TokenStream> {
-    #[allow(unused_variables)]
-    let args = Args::from_list(args).map_err(|e| e.write_errors())?;
 
     let fargs = f.sig.inputs.clone();
 
-    let ctxt = Ctxt::new();
-
     if f.sig.asyncness.is_none() {
-        ctxt.error_spanned_by(&f.sig, "main function must be async");
+        error(&mut errors, &f.sig, "main function must be async");
     }
     if !f.sig.generics.params.is_empty() {
-        ctxt.error_spanned_by(&f.sig, "main function must not be generic");
+        error(&mut errors, &f.sig, "main function must not be generic");
     }
     if !f.sig.generics.where_clause.is_none() {
-        ctxt.error_spanned_by(&f.sig, "main function must not have `where` clauses");
+        error(&mut errors, &f.sig, "main function must not have `where` clauses");
     }
     if !f.sig.abi.is_none() {
-        ctxt.error_spanned_by(&f.sig, "main function must not have an ABI qualifier");
+        error(&mut errors, &f.sig, "main function must not have an ABI qualifier");
     }
     if !f.sig.variadic.is_none() {
-        ctxt.error_spanned_by(&f.sig, "main function must not be variadic");
+        error(&mut errors, &f.sig, "main function 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,
                 "main function must either not return a value, return `()` or return `!`",
             ),
@@ -154,26 +109,69 @@ pub fn run(args: &[NestedMeta], f: syn::ItemFn, main: TokenStream) -> Result<Tok
     }
 
     if fargs.len() != 1 {
-        ctxt.error_spanned_by(&f.sig, "main function must have 1 argument: the spawner.");
+        error(&mut errors, &f.sig, "main function must have 1 argument: the spawner.");
     }
 
-    ctxt.check()?;
+    let entry = match args.entry.as_deref().or(arch.default_entry) {
+        None => TokenStream::new(),
+        Some(x) => match TokenStream::from_str(x) {
+            Ok(x) => quote!(#[#x]),
+            Err(e) => {
+                error(&mut errors, &f.sig, e);
+                TokenStream::new()
+            }
+        },
+    };
 
-    let f_body = f.block;
+    let f_body = f.body;
     let out = &f.sig.output;
 
+    let (main_ret, mut main_body) = match arch.flavor {
+        Flavor::Standard => (
+            quote!(!),
+            quote! {
+                unsafe fn __make_static<T>(t: &mut T) -> &'static mut T {
+                    ::core::mem::transmute(t)
+                }
+
+                let mut executor = ::embassy_executor::Executor::new();
+                let executor = unsafe { __make_static(&mut executor) };
+                executor.run(|spawner| {
+                    spawner.must_spawn(__embassy_main(spawner));
+                })
+            },
+        ),
+        Flavor::Wasm => (
+            quote!(Result<(), wasm_bindgen::JsValue>),
+            quote! {
+                let executor = ::std::boxed::Box::leak(::std::boxed::Box::new(::embassy_executor::Executor::new()));
+
+                executor.start(|spawner| {
+                    spawner.must_spawn(__embassy_main(spawner));
+                });
+
+                Ok(())
+            },
+        ),
+    };
+
+    if !errors.is_empty() {
+        main_body = quote! {loop{}};
+    }
+
     let result = quote! {
         #[::embassy_executor::task()]
         async fn __embassy_main(#fargs) #out {
             #f_body
         }
 
-        unsafe fn __make_static<T>(t: &mut T) -> &'static mut T {
-            ::core::mem::transmute(t)
+        #entry
+        fn main() -> #main_ret {
+            #main_body
         }
 
-        #main
+        #errors
     };
 
-    Ok(result)
+    result
 }
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);
 }