Allow for an optional user-defined entry macro when targeting RISC-V

2 files changed, 31 insertions, 5 deletions

diff --git a/embassy-macros/src/lib.rs b/embassy-macros/src/lib.rs
index d2c696c72..dc5b519ff 100644
--- a/embassy-macros/src/lib.rs
+++ b/embassy-macros/src/lib.rs
@@ -79,6 +79,8 @@ pub fn main_cortex_m(args: TokenStream, item: TokenStream) -> TokenStream {
 /// * The function must not use generics.
 /// * Only a single `main` task may be declared.
 ///
+/// A user-defined entry macro can be optionally provided via the `entry` argument to override the default of `riscv_rt::entry`.
+///
 /// ## Examples
 /// Spawning a task:
 ///
@@ -88,11 +90,21 @@ pub fn main_cortex_m(args: TokenStream, item: TokenStream) -> TokenStream {
 ///     // Function body
 /// }
 /// ```
+///
+/// Spawning a task using a custom entry macro:
+/// ``` rust
+/// #[embassy_executor::main(entry = "esp_riscv_rt::entry")]
+/// 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()
+    main::run(args.clone(), f, main::riscv(args))
+        .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.
diff --git a/embassy-macros/src/macros/main.rs b/embassy-macros/src/macros/main.rs
index 6ae77398d..5c099f68a 100644
--- a/embassy-macros/src/macros/main.rs
+++ b/embassy-macros/src/macros/main.rs
@@ -1,16 +1,30 @@
 use darling::FromMeta;
 use proc_macro2::TokenStream;
 use quote::quote;
-use syn::{ReturnType, Type};
+use syn::{Expr, ReturnType, Type};
 
 use crate::util::ctxt::Ctxt;
 
 #[derive(Debug, FromMeta)]
-struct Args {}
+struct Args {
+    #[darling(default)]
+    entry: Option<String>,
+}
+
+pub fn riscv(args: syn::AttributeArgs) -> TokenStream {
+    let maybe_entry = match Args::from_list(&args) {
+        Ok(args) => args.entry,
+        Err(e) => return e.write_errors(),
+    };
+
+    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(),
+    };
 
-pub fn riscv() -> TokenStream {
     quote! {
-        #[riscv_rt::entry]
+        #[#entry]
         fn main() -> ! {
             let mut executor = ::embassy_executor::Executor::new();
             let executor = unsafe { __make_static(&mut executor) };