1 files changed, 119 insertions, 121 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
 }