Extend the task macro to allow cfging arguments away

embassy-executor-macros/src/macros/task.rs

@@ -49,7 +49,7 @@ pub fn run(args: &[NestedMeta], f: syn::ItemFn) -> Result<TokenStream, TokenStre
         },
     }
 
-    let mut arg_names = Vec::new();
+    let mut args = Vec::new();
     let mut fargs = f.sig.inputs.clone();
 
     for arg in fargs.iter_mut() {
@@ -59,8 +59,8 @@ pub fn run(args: &[NestedMeta], f: syn::ItemFn) -> Result<TokenStream, TokenStre
             }
             syn::FnArg::Typed(t) => match t.pat.as_mut() {
                 syn::Pat::Ident(id) => {
-                    arg_names.push(id.ident.clone());
                     id.mutability = None;
+                    args.push((id.clone(), t.attrs.clone()));
                 }
                 _ => {
                     ctxt.error_spanned_by(arg, "pattern matching in task arguments is not yet supported");
@@ -79,21 +79,32 @@ pub fn run(args: &[NestedMeta], f: syn::ItemFn) -> Result<TokenStream, TokenStre
     task_inner.vis = syn::Visibility::Inherited;
     task_inner.sig.ident = task_inner_ident.clone();
 
+    // assemble the original input arguments,
+    // including any attributes that may have
+    // been applied previously
+    let mut full_args = Vec::new();
+    for (arg, cfgs) in args {
+        full_args.push(quote!(
+            #(#cfgs)*
+            #arg
+        ));
+    }
+
     #[cfg(feature = "nightly")]
     let mut task_outer: ItemFn = parse_quote! {
         #visibility fn #task_ident(#fargs) -> ::embassy_executor::SpawnToken<impl Sized> {
             type Fut = impl ::core::future::Future + 'static;
             const POOL_SIZE: usize = #pool_size;
             static POOL: ::embassy_executor::raw::TaskPool<Fut, POOL_SIZE> = ::embassy_executor::raw::TaskPool::new();
-            unsafe { POOL._spawn_async_fn(move || #task_inner_ident(#(#arg_names,)*)) }
+            unsafe { POOL._spawn_async_fn(move || #task_inner_ident(#(#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(#(#arg_names,)*)) }
+            unsafe { POOL.get::<_, POOL_SIZE>()._spawn_async_fn(move || #task_inner_ident(#(#full_args,)*)) }
         }
     };
 

embassy-executor/tests/test.rs

@@ -135,3 +135,17 @@ fn executor_task_self_wake_twice() {
         ]
     )
 }
+
+#[test]
+fn executor_task_cfg_args() {
+    // simulate cfg'ing away argument c
+    #[task]
+    async fn task1(a: u32, b: u32, #[cfg(any())] c: u32) {
+        let (_, _) = (a, b);
+    }
+
+    #[task]
+    async fn task2(a: u32, b: u32, #[cfg(all())] c: u32) {
+        let (_, _, _) = (a, b, c);
+    }
+}