Merge pull request #140 from lulf/reexport-macro-param

Add embassy_prefix attribute parameter to task and main macros

5 files changed, 122 insertions, 29 deletions

diff --git a/embassy-macros/src/chip/nrf.rs b/embassy-macros/src/chip/nrf.rs
index b2d1c322f..b7ea8d389 100644
--- a/embassy-macros/src/chip/nrf.rs
+++ b/embassy-macros/src/chip/nrf.rs
@@ -1,3 +1,4 @@
+use crate::path::ModulePrefix;
 use darling::FromMeta;
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
@@ -30,32 +31,37 @@ impl Default for LfclkSource {
     }
 }
 
-#[derive(Debug, FromMeta)]
+#[derive(Debug, FromMeta, Default)]
 pub struct Args {
     #[darling(default)]
+    pub embassy_prefix: ModulePrefix,
+    #[darling(default)]
     pub hfclk_source: HfclkSource,
     #[darling(default)]
     pub lfclk_source: LfclkSource,
 }
 
-pub fn generate(args: Args) -> TokenStream {
+pub fn generate(args: &Args) -> TokenStream {
     let hfclk_source = format_ident!("{}", format!("{:?}", args.hfclk_source));
     let lfclk_source = format_ident!("{}", format!("{:?}", args.lfclk_source));
 
+    let embassy_path = args.embassy_prefix.append("embassy").path();
+    let embassy_nrf_path = args.embassy_prefix.append("embassy_nrf").path();
+
     quote!(
-        use embassy_nrf::{interrupt, peripherals, rtc};
+        use #embassy_nrf_path::{interrupt, peripherals, rtc};
 
-        let mut config = embassy_nrf::system::Config::default();
-        config.hfclk_source = embassy_nrf::system::HfclkSource::#hfclk_source;
-        config.lfclk_source = embassy_nrf::system::LfclkSource::#lfclk_source;
-        unsafe { embassy_nrf::system::configure(config) };
+        let mut config = #embassy_nrf_path::system::Config::default();
+        config.hfclk_source = #embassy_nrf_path::system::HfclkSource::#hfclk_source;
+        config.lfclk_source = #embassy_nrf_path::system::LfclkSource::#lfclk_source;
+        unsafe { #embassy_nrf_path::system::configure(config) };
 
-        let mut rtc = rtc::RTC::new(unsafe { <peripherals::RTC1 as embassy::util::Steal>::steal() }, interrupt::take!(RTC1));
+        let mut rtc = rtc::RTC::new(unsafe { <peripherals::RTC1 as #embassy_path::util::Steal>::steal() }, interrupt::take!(RTC1));
         let rtc = unsafe { make_static(&mut rtc) };
         rtc.start();
         let mut alarm = rtc.alarm0();
 
-        unsafe { embassy::time::set_clock(rtc) };
+        unsafe { #embassy_path::time::set_clock(rtc) };
 
         let alarm = unsafe { make_static(&mut alarm) };
         executor.set_alarm(alarm);
diff --git a/embassy-macros/src/chip/rp.rs b/embassy-macros/src/chip/rp.rs
index 39a418ff0..6093dc9d3 100644
--- a/embassy-macros/src/chip/rp.rs
+++ b/embassy-macros/src/chip/rp.rs
@@ -1,15 +1,20 @@
+use crate::path::ModulePrefix;
 use darling::FromMeta;
 use proc_macro2::TokenStream;
 use quote::quote;
 
-#[derive(Debug, FromMeta)]
-pub struct Args {}
+#[derive(Debug, FromMeta, Default)]
+pub struct Args {
+    #[darling(default)]
+    pub embassy_prefix: ModulePrefix,
+}
 
-pub fn generate(_args: Args) -> TokenStream {
+pub fn generate(args: &Args) -> TokenStream {
+    let embassy_rp_path = args.embassy_prefix.append("embassy_rp").path();
     quote!(
-        use embassy_rp::{interrupt, peripherals};
+        use #embassy_rp_path::{interrupt, peripherals};
 
-        let mut config = embassy_rp::system::Config::default();
-        unsafe { embassy_rp::system::configure(config) };
+        let mut config = #embassy_rp_path::system::Config::default();
+        unsafe { #embassy_rp_path::system::configure(config) };
     )
 }
diff --git a/embassy-macros/src/chip/stm32.rs b/embassy-macros/src/chip/stm32.rs
index 2476a9623..486a9aa65 100644
--- a/embassy-macros/src/chip/stm32.rs
+++ b/embassy-macros/src/chip/stm32.rs
@@ -1,11 +1,14 @@
+use crate::path::ModulePrefix;
 use darling::FromMeta;
 use proc_macro2::TokenStream;
 use quote::{format_ident, quote};
 use syn::spanned::Spanned;
 
-#[derive(Debug, FromMeta)]
+#[derive(Debug, FromMeta, Default)]
 pub struct Args {
     #[darling(default)]
+    pub embassy_prefix: ModulePrefix,
+    #[darling(default)]
     pub use_hse: Option<u32>,
     #[darling(default)]
     pub sysclk: Option<u32>,
@@ -15,7 +18,10 @@ pub struct Args {
     pub require_pll48clk: bool,
 }
 
-pub fn generate(args: Args) -> TokenStream {
+pub fn generate(args: &Args) -> TokenStream {
+    let embassy_path = args.embassy_prefix.append("embassy").path();
+    let embassy_stm32_path = args.embassy_prefix.append("embassy_stm32").path();
+
     let mut clock_cfg_args = quote! {};
     if args.use_hse.is_some() {
         let mhz = args.use_hse.unwrap();
@@ -37,7 +43,7 @@ pub fn generate(args: Args) -> TokenStream {
     }
 
     quote!(
-        use embassy_stm32::{rtc, interrupt, Peripherals, pac, hal::rcc::RccExt, hal::time::U32Ext};
+        use #embassy_stm32_path::{rtc, interrupt, Peripherals, pac, hal::rcc::RccExt, hal::time::U32Ext};
 
         let dp = pac::Peripherals::take().unwrap();
         let rcc = dp.RCC.constrain();
@@ -50,7 +56,7 @@ pub fn generate(args: Args) -> TokenStream {
         rtc.start();
         let mut alarm = rtc.alarm1();
 
-        unsafe { embassy::time::set_clock(rtc) };
+        unsafe { #embassy_path::time::set_clock(rtc) };
 
         let alarm = unsafe { make_static(&mut alarm) };
         executor.set_alarm(alarm);
diff --git a/embassy-macros/src/lib.rs b/embassy-macros/src/lib.rs
index 8f35cc9ef..402f0e08a 100644
--- a/embassy-macros/src/lib.rs
+++ b/embassy-macros/src/lib.rs
@@ -7,12 +7,18 @@ use proc_macro::{Span, TokenStream};
 use quote::{format_ident, quote};
 use syn::spanned::Spanned;
 
+mod path;
+
+use path::ModulePrefix;
+
 #[derive(Debug, FromMeta)]
-struct MacroArgs {
+struct TaskArgs {
     #[darling(default)]
     pool_size: Option<usize>,
     #[darling(default)]
     send: bool,
+    #[darling(default)]
+    embassy_prefix: ModulePrefix,
 }
 
 #[proc_macro_attribute]
@@ -20,13 +26,16 @@ pub fn task(args: TokenStream, item: TokenStream) -> TokenStream {
     let macro_args = syn::parse_macro_input!(args as syn::AttributeArgs);
     let mut task_fn = syn::parse_macro_input!(item as syn::ItemFn);
 
-    let macro_args = match MacroArgs::from_list(&macro_args) {
+    let macro_args = match TaskArgs::from_list(&macro_args) {
         Ok(v) => v,
         Err(e) => {
             return TokenStream::from(e.write_errors());
         }
     };
 
+    let embassy_prefix = macro_args.embassy_prefix.append("embassy");
+    let embassy_path = embassy_prefix.path();
+
     let pool_size: usize = macro_args.pool_size.unwrap_or(1);
 
     let mut fail = false;
@@ -99,8 +108,8 @@ pub fn task(args: TokenStream, item: TokenStream) -> TokenStream {
     };
 
     let result = quote! {
-        #visibility fn #name(#args) -> ::embassy::executor::SpawnToken<#impl_ty> {
-            use ::embassy::executor::raw::Task;
+        #visibility fn #name(#args) -> #embassy_path::executor::SpawnToken<#impl_ty> {
+            use #embassy_path::executor::raw::Task;
             #task_fn
             type F = #impl_ty;
             const NEW_TASK: Task<F> = Task::new();
@@ -203,6 +212,15 @@ mod chip;
 #[path = "chip/rp.rs"]
 mod chip;
 
+#[cfg(feature = "std")]
+mod chip {
+    #[derive(Debug, darling::FromMeta, Default)]
+    pub struct Args {
+        #[darling(default)]
+        pub embassy_prefix: crate::path::ModulePrefix,
+    }
+}
+
 #[cfg(any(feature = "nrf", feature = "stm32", feature = "rp"))]
 #[proc_macro_attribute]
 pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
@@ -252,11 +270,13 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
         return TokenStream::new();
     }
 
+    let embassy_prefix_lit = macro_args.embassy_prefix.literal();
+    let embassy_path = macro_args.embassy_prefix.append("embassy").path();
     let task_fn_body = task_fn.block.clone();
-    let chip_setup = chip::generate(macro_args);
+    let chip_setup = chip::generate(&macro_args);
 
     let result = quote! {
-        #[embassy::task]
+        #[#embassy_path::task(embassy_prefix = #embassy_prefix_lit)]
         async fn __embassy_main(#args) {
             #task_fn_body
         }
@@ -267,7 +287,7 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
                 ::core::mem::transmute(t)
             }
 
-            let mut executor = ::embassy::executor::Executor::new();
+            let mut executor = #embassy_path::executor::Executor::new();
             let executor = unsafe { make_static(&mut executor) };
 
             #chip_setup
@@ -283,9 +303,20 @@ pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
 
 #[cfg(feature = "std")]
 #[proc_macro_attribute]
-pub fn main(_: TokenStream, item: TokenStream) -> TokenStream {
+pub fn main(args: TokenStream, item: TokenStream) -> TokenStream {
+    let macro_args = syn::parse_macro_input!(args as syn::AttributeArgs);
     let task_fn = syn::parse_macro_input!(item as syn::ItemFn);
 
+    let macro_args = match chip::Args::from_list(&macro_args) {
+        Ok(v) => v,
+        Err(e) => {
+            return TokenStream::from(e.write_errors());
+        }
+    };
+
+    let embassy_path = macro_args.embassy_prefix.append("embassy");
+    let embassy_std_path = macro_args.embassy_prefix.append("embassy_std");
+
     let mut fail = false;
     if task_fn.sig.asyncness.is_none() {
         task_fn
@@ -324,8 +355,12 @@ pub fn main(_: TokenStream, item: TokenStream) -> TokenStream {
 
     let task_fn_body = task_fn.block.clone();
 
+    let embassy_path = embassy_path.path();
+    let embassy_std_path = embassy_std_path.path();
+    let embassy_prefix_lit = macro_args.embassy_prefix.literal();
+
     let result = quote! {
-        #[embassy::task]
+        #[#embassy_path::task(embassy_prefix = #embassy_prefix_lit)]
         async fn __embassy_main(#args) {
             #task_fn_body
         }
@@ -335,7 +370,7 @@ pub fn main(_: TokenStream, item: TokenStream) -> TokenStream {
                 ::core::mem::transmute(t)
             }
 
-            let mut executor = ::embassy_std::Executor::new();
+            let mut executor = #embassy_std_path::Executor::new();
             let executor = unsafe { make_static(&mut executor) };
 
             executor.run(|spawner| {
diff --git a/embassy-macros/src/path.rs b/embassy-macros/src/path.rs
new file mode 100644
index 000000000..00fca7bdb
--- /dev/null
+++ b/embassy-macros/src/path.rs
@@ -0,0 +1,41 @@
+use darling::{FromMeta, Result};
+use proc_macro2::Span;
+use syn::{LitStr, Path};
+
+#[derive(Debug)]
+pub struct ModulePrefix {
+    literal: LitStr,
+}
+
+impl ModulePrefix {
+    pub fn new(path: &str) -> Self {
+        let literal = LitStr::new(path, Span::call_site());
+        Self { literal }
+    }
+
+    pub fn append(&self, component: &str) -> ModulePrefix {
+        let mut lit = self.literal().value();
+        lit.push_str(component);
+        Self::new(lit.as_str())
+    }
+
+    pub fn path(&self) -> Path {
+        self.literal.parse().unwrap()
+    }
+
+    pub fn literal(&self) -> &LitStr {
+        &self.literal
+    }
+}
+
+impl FromMeta for ModulePrefix {
+    fn from_string(value: &str) -> Result<Self> {
+        Ok(ModulePrefix::new(value))
+    }
+}
+
+impl Default for ModulePrefix {
+    fn default() -> ModulePrefix {
+        ModulePrefix::new("::")
+    }
+}