5 files changed, 140 insertions, 58 deletions

diff --git a/embassy-executor/Cargo.toml b/embassy-executor/Cargo.toml
index 45b0955bf..6fa1dd7fb 100644
--- a/embassy-executor/Cargo.toml
+++ b/embassy-executor/Cargo.toml
@@ -29,9 +29,8 @@ flavors = [
 
 [features]
 default = []
-std = ["embassy-macros/std", "critical-section/std"]
-wasm = ["dep:wasm-bindgen", "dep:js-sys", "embassy-macros/wasm"]
-riscv = ["embassy-macros/riscv"]
+std = ["critical-section/std"]
+wasm = ["dep:wasm-bindgen", "dep:js-sys"]
 
 # Enable nightly-only features
 nightly = []
diff --git a/embassy-executor/src/lib.rs b/embassy-executor/src/lib.rs
index e4cbd04b9..4c7e2f4cd 100644
--- a/embassy-executor/src/lib.rs
+++ b/embassy-executor/src/lib.rs
@@ -8,18 +8,22 @@
 pub(crate) mod fmt;
 
 #[cfg(feature = "nightly")]
-pub use embassy_macros::{main, task};
+pub use embassy_macros::task;
 
 cfg_if::cfg_if! {
     if #[cfg(cortex_m)] {
         #[path="arch/cortex_m.rs"]
         mod arch;
         pub use arch::*;
+        #[cfg(feature = "nightly")]
+        pub use embassy_macros::main_cortex_m as main;
     }
     else if #[cfg(target_arch="riscv32")] {
         #[path="arch/riscv32.rs"]
         mod arch;
         pub use arch::*;
+        #[cfg(feature = "nightly")]
+        pub use embassy_macros::main_riscv as main;
     }
     else if #[cfg(all(target_arch="xtensa", feature = "nightly"))] {
         #[path="arch/xtensa.rs"]
@@ -30,11 +34,15 @@ cfg_if::cfg_if! {
         #[path="arch/wasm.rs"]
         mod arch;
         pub use arch::*;
+        #[cfg(feature = "nightly")]
+        pub use embassy_macros::main_wasm as main;
     }
     else if #[cfg(feature="std")] {
         #[path="arch/std.rs"]
         mod arch;
         pub use arch::*;
+        #[cfg(feature = "nightly")]
+        pub use embassy_macros::main_std as main;
     }
 }
 
diff --git a/embassy-macros/Cargo.toml b/embassy-macros/Cargo.toml
index 9b83771c8..5c612c997 100644
--- a/embassy-macros/Cargo.toml
+++ b/embassy-macros/Cargo.toml
@@ -21,9 +21,5 @@ proc-macro2 = "1.0.29"
 proc-macro = true
 
 [features]
-std = []
-wasm = []
-riscv = []
-
 # Enabling this cause interrupt::take! to require embassy-executor
 rtos-trace-interrupt = []
diff --git a/embassy-macros/src/lib.rs b/embassy-macros/src/lib.rs
index f5df2a269..d2c696c72 100644
--- a/embassy-macros/src/lib.rs
+++ b/embassy-macros/src/lib.rs
@@ -45,7 +45,7 @@ pub fn task(args: TokenStream, item: TokenStream) -> TokenStream {
     task::run(args, f).unwrap_or_else(|x| x).into()
 }
 
-/// Creates a new `executor` instance and declares an application entry point spawning the corresponding function body as an async task.
+/// Creates a new `executor` instance and declares an application entry point for Cortex-M spawning the corresponding function body as an async task.
 ///
 /// The following restrictions apply:
 ///
@@ -64,10 +64,85 @@ pub fn task(args: TokenStream, item: TokenStream) -> TokenStream {
 /// }
 /// ```
 #[proc_macro_attribute]
-pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
+pub fn main_cortex_m(args: TokenStream, item: TokenStream) -> TokenStream {
     let args = syn::parse_macro_input!(args as syn::AttributeArgs);
     let f = syn::parse_macro_input!(item as syn::ItemFn);
-    main::run(args, f).unwrap_or_else(|x| x).into()
+    main::run(args, f, main::cortex_m()).unwrap_or_else(|x| x).into()
+}
+
+/// Creates a new `executor` instance and declares an application entry point for RISC-V spawning the corresponding function body as an async task.
+///
+/// The following restrictions apply:
+///
+/// * The function must accept exactly 1 parameter, an `embassy_executor::Spawner` handle that it can use to spawn additional tasks.
+/// * The function must be declared `async`.
+/// * The function must not use generics.
+/// * Only a single `main` task may be declared.
+///
+/// ## Examples
+/// Spawning a task:
+///
+/// ``` rust
+/// #[embassy_executor::main]
+/// async fn main(_s: embassy_executor::Spawner) {
+///     // Function body
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn main_riscv(args: TokenStream, item: TokenStream) -> TokenStream {
+    let args = syn::parse_macro_input!(args as syn::AttributeArgs);
+    let f = syn::parse_macro_input!(item as syn::ItemFn);
+    main::run(args, f, main::riscv()).unwrap_or_else(|x| x).into()
+}
+
+/// Creates a new `executor` instance and declares an application entry point for STD spawning the corresponding function body as an async task.
+///
+/// The following restrictions apply:
+///
+/// * The function must accept exactly 1 parameter, an `embassy_executor::Spawner` handle that it can use to spawn additional tasks.
+/// * The function must be declared `async`.
+/// * The function must not use generics.
+/// * Only a single `main` task may be declared.
+///
+/// ## Examples
+/// Spawning a task:
+///
+/// ``` rust
+/// #[embassy_executor::main]
+/// async fn main(_s: embassy_executor::Spawner) {
+///     // Function body
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn main_std(args: TokenStream, item: TokenStream) -> TokenStream {
+    let args = syn::parse_macro_input!(args as syn::AttributeArgs);
+    let f = syn::parse_macro_input!(item as syn::ItemFn);
+    main::run(args, f, main::std()).unwrap_or_else(|x| x).into()
+}
+
+/// Creates a new `executor` instance and declares an application entry point for WASM spawning the corresponding function body as an async task.
+///
+/// The following restrictions apply:
+///
+/// * The function must accept exactly 1 parameter, an `embassy_executor::Spawner` handle that it can use to spawn additional tasks.
+/// * The function must be declared `async`.
+/// * The function must not use generics.
+/// * Only a single `main` task may be declared.
+///
+/// ## Examples
+/// Spawning a task:
+///
+/// ``` rust
+/// #[embassy_executor::main]
+/// async fn main(_s: embassy_executor::Spawner) {
+///     // Function body
+/// }
+/// ```
+#[proc_macro_attribute]
+pub fn main_wasm(args: TokenStream, item: TokenStream) -> TokenStream {
+    let args = syn::parse_macro_input!(args as syn::AttributeArgs);
+    let f = syn::parse_macro_input!(item as syn::ItemFn);
+    main::run(args, f, main::wasm()).unwrap_or_else(|x| x).into()
 }
 
 #[proc_macro_attribute]
diff --git a/embassy-macros/src/macros/main.rs b/embassy-macros/src/macros/main.rs
index 54806847c..18f7c36c4 100644
--- a/embassy-macros/src/macros/main.rs
+++ b/embassy-macros/src/macros/main.rs
@@ -7,31 +7,34 @@ use crate::util::ctxt::Ctxt;
 #[derive(Debug, FromMeta)]
 struct Args {}
 
-pub fn run(args: syn::AttributeArgs, f: syn::ItemFn) -> 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");
-    }
-    if !f.sig.generics.params.is_empty() {
-        ctxt.error_spanned_by(&f.sig, "main function must not be generic");
+pub fn riscv() -> TokenStream {
+    quote! {
+        #[riscv_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));
+            })
+        }
     }
+}
 
-    if fargs.len() != 1 {
-        ctxt.error_spanned_by(&f.sig, "main function must have 1 argument: the 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));
+            })
+        }
     }
+}
 
-    ctxt.check()?;
-
-    let f_body = f.block;
-
-    #[cfg(feature = "wasm")]
-    let main = quote! {
+pub fn wasm() -> TokenStream {
+    quote! {
         #[wasm_bindgen::prelude::wasm_bindgen(start)]
         pub fn main() -> Result<(), wasm_bindgen::JsValue> {
             static EXECUTOR: ::embassy_executor::_export::StaticCell<::embassy_executor::Executor> = ::embassy_executor::_export::StaticCell::new();
@@ -43,10 +46,11 @@ pub fn run(args: syn::AttributeArgs, f: syn::ItemFn) -> Result<TokenStream, Toke
 
             Ok(())
         }
-    };
+    }
+}
 
-    #[cfg(all(feature = "std", not(feature = "wasm"), not(feature = "riscv")))]
-    let main = quote! {
+pub fn std() -> TokenStream {
+    quote! {
         fn main() -> ! {
             let mut executor = ::embassy_executor::Executor::new();
             let executor = unsafe { __make_static(&mut executor) };
@@ -55,31 +59,31 @@ pub fn run(args: syn::AttributeArgs, f: syn::ItemFn) -> Result<TokenStream, Toke
                 spawner.must_spawn(__embassy_main(spawner));
             })
         }
-    };
+    }
+}
 
-    #[cfg(all(not(feature = "std"), not(feature = "wasm"), not(feature = "riscv")))]
-    let main = 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 run(args: syn::AttributeArgs, f: syn::ItemFn, main: TokenStream) -> Result<TokenStream, TokenStream> {
+    #[allow(unused_variables)]
+    let args = Args::from_list(&args).map_err(|e| e.write_errors())?;
 
-    #[cfg(all(not(feature = "std"), not(feature = "wasm"), feature = "riscv"))]
-    let main = quote! {
-        #[riscv_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));
-            })
-        }
-    };
+    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");
+    }
+    if !f.sig.generics.params.is_empty() {
+        ctxt.error_spanned_by(&f.sig, "main function must not be generic");
+    }
+
+    if fargs.len() != 1 {
+        ctxt.error_spanned_by(&f.sig, "main function must have 1 argument: the spawner.");
+    }
+
+    ctxt.check()?;
+
+    let f_body = f.block;
 
     let result = quote! {
         #[::embassy_executor::task()]