Async tasks can now take arguments at spawn again

10 files changed, 141 insertions, 75 deletions

diff --git a/macros/src/codegen/module.rs b/macros/src/codegen/module.rs
index 19cf2417..666bd042 100644
--- a/macros/src/codegen/module.rs
+++ b/macros/src/codegen/module.rs
@@ -146,6 +146,8 @@ pub fn codegen(ctxt: Context, app: &App, analysis: &Analysis) -> TokenStream2 {
         };
 
         let internal_spawn_ident = util::internal_task_ident(name, "spawn");
+        let (input_args, input_tupled, input_untupled, input_ty) =
+            util::regroup_inputs(&spawnee.inputs);
 
         // Spawn caller
         items.push(quote!(
@@ -153,17 +155,18 @@ pub fn codegen(ctxt: Context, app: &App, analysis: &Analysis) -> TokenStream2 {
             /// Spawns the task directly
             #[allow(non_snake_case)]
             #[doc(hidden)]
-            pub fn #internal_spawn_ident() -> Result<(), ()> {
+            pub fn #internal_spawn_ident(#(#input_args,)*) -> Result<(), #input_ty> {
                 if #exec_name.try_reserve() {
+                    // This unsafe is protected by `try_reserve`, see its documentation for details
                     unsafe {
-                        // TODO: Add args here
-                        #exec_name.spawn_unchecked(#name(#name::Context::new()));
+                        #exec_name.spawn_unchecked(#name(#name::Context::new() #(,#input_untupled)*));
                     }
+
                     #pend_interrupt
 
                     Ok(())
                 } else {
-                    Err(())
+                    Err(#input_tupled)
                 }
             }
         ));
diff --git a/macros/src/codegen/software_tasks.rs b/macros/src/codegen/software_tasks.rs
index b9232832..34fc851a 100644
--- a/macros/src/codegen/software_tasks.rs
+++ b/macros/src/codegen/software_tasks.rs
@@ -36,12 +36,13 @@ pub fn codegen(app: &App, analysis: &Analysis) -> TokenStream2 {
             let attrs = &task.attrs;
             let cfgs = &task.cfgs;
             let stmts = &task.stmts;
+            let inputs = &task.inputs;
 
             user_tasks.push(quote!(
                 #(#attrs)*
                 #(#cfgs)*
                 #[allow(non_snake_case)]
-                async fn #name(#context: #name::Context<'static>) {
+                async fn #name<'a>(#context: #name::Context<'a> #(,#inputs)*) {
                     use rtic::Mutex as _;
                     use rtic::mutex::prelude::*;
 
diff --git a/macros/src/codegen/util.rs b/macros/src/codegen/util.rs
index 0558d9d1..e121487c 100644
--- a/macros/src/codegen/util.rs
+++ b/macros/src/codegen/util.rs
@@ -1,9 +1,8 @@
-use core::sync::atomic::{AtomicUsize, Ordering};
-
 use crate::syntax::{ast::App, Context};
+use core::sync::atomic::{AtomicUsize, Ordering};
 use proc_macro2::{Span, TokenStream as TokenStream2};
 use quote::quote;
-use syn::{Attribute, Ident};
+use syn::{Attribute, Ident, PatType};
 
 const RTIC_INTERNAL: &str = "__rtic_internal";
 
@@ -94,6 +93,55 @@ pub fn link_section_uninit() -> TokenStream2 {
     quote!(#[link_section = #section])
 }
 
+/// Regroups the inputs of a task
+///
+/// `inputs` could be &[`input: Foo`] OR &[`mut x: i32`, `ref y: i64`]
+pub fn regroup_inputs(
+    inputs: &[PatType],
+) -> (
+    // args e.g. &[`_0`],  &[`_0: i32`, `_1: i64`]
+    Vec<TokenStream2>,
+    // tupled e.g. `_0`, `(_0, _1)`
+    TokenStream2,
+    // untupled e.g. &[`_0`], &[`_0`, `_1`]
+    Vec<TokenStream2>,
+    // ty e.g. `Foo`, `(i32, i64)`
+    TokenStream2,
+) {
+    if inputs.len() == 1 {
+        let ty = &inputs[0].ty;
+
+        (
+            vec![quote!(_0: #ty)],
+            quote!(_0),
+            vec![quote!(_0)],
+            quote!(#ty),
+        )
+    } else {
+        let mut args = vec![];
+        let mut pats = vec![];
+        let mut tys = vec![];
+
+        for (i, input) in inputs.iter().enumerate() {
+            let i = Ident::new(&format!("_{}", i), Span::call_site());
+            let ty = &input.ty;
+
+            args.push(quote!(#i: #ty));
+
+            pats.push(quote!(#i));
+
+            tys.push(quote!(#ty));
+        }
+
+        let tupled = {
+            let pats = pats.clone();
+            quote!((#(#pats,)*))
+        };
+        let ty = quote!((#(#tys,)*));
+        (args, tupled, pats, ty)
+    }
+}
+
 /// Get the ident for the name of the task
 pub fn get_task_name(ctxt: Context, app: &App) -> Ident {
     let s = match ctxt {
diff --git a/macros/src/syntax/analyze.rs b/macros/src/syntax/analyze.rs
index b70ceb8b..3ed14877 100644
--- a/macros/src/syntax/analyze.rs
+++ b/macros/src/syntax/analyze.rs
@@ -287,6 +287,11 @@ pub(crate) fn app(app: &App) -> Result<Analysis, syn::Error> {
 
         let channel = channels.entry(spawnee_prio).or_default();
         channel.tasks.insert(name.clone());
+
+        // All inputs are send as we do not know from where they may be spawned.
+        spawnee.inputs.iter().for_each(|input| {
+            send_types.insert(input.ty.clone());
+        });
     }
 
     // No channel should ever be empty
diff --git a/macros/src/syntax/ast.rs b/macros/src/syntax/ast.rs
index da6016ad..27e6773f 100644
--- a/macros/src/syntax/ast.rs
+++ b/macros/src/syntax/ast.rs
@@ -1,6 +1,6 @@
 //! Abstract Syntax Tree
 
-use syn::{Attribute, Expr, Ident, Item, ItemUse, Pat, Path, Stmt, Type};
+use syn::{Attribute, Expr, Ident, Item, ItemUse, Pat, PatType, Path, Stmt, Type};
 
 use crate::syntax::Map;
 
@@ -205,6 +205,9 @@ pub struct SoftwareTask {
     /// The context argument
     pub context: Box<Pat>,
 
+    /// The inputs of this software task
+    pub inputs: Vec<PatType>,
+
     /// The statements that make up the task handler
     pub stmts: Vec<Stmt>,
 
diff --git a/macros/src/syntax/parse/hardware_task.rs b/macros/src/syntax/parse/hardware_task.rs
index c13426f4..7f6dfbe4 100644
--- a/macros/src/syntax/parse/hardware_task.rs
+++ b/macros/src/syntax/parse/hardware_task.rs
@@ -15,25 +15,20 @@ impl HardwareTask {
 
         let name = item.sig.ident.to_string();
 
-        if name == "init" || name == "idle" {
-            return Err(parse::Error::new(
-                span,
-                "tasks cannot be named `init` or `idle`",
-            ));
-        }
-
         if valid_signature {
-            if let Some(context) = util::parse_inputs(item.sig.inputs, &name) {
-                let FilterAttrs { cfgs, attrs, .. } = util::filter_attributes(item.attrs);
+            if let Some((context, Ok(rest))) = util::parse_inputs(item.sig.inputs, &name) {
+                if rest.is_empty() {
+                    let FilterAttrs { cfgs, attrs, .. } = util::filter_attributes(item.attrs);
 
-                return Ok(HardwareTask {
-                    args,
-                    cfgs,
-                    attrs,
-                    context,
-                    stmts: item.block.stmts,
-                    is_extern: false,
-                });
+                    return Ok(HardwareTask {
+                        args,
+                        cfgs,
+                        attrs,
+                        context,
+                        stmts: item.block.stmts,
+                        is_extern: false,
+                    });
+                }
             }
         }
 
@@ -56,25 +51,20 @@ impl HardwareTask {
 
         let name = item.sig.ident.to_string();
 
-        if name == "init" || name == "idle" {
-            return Err(parse::Error::new(
-                span,
-                "tasks cannot be named `init` or `idle`",
-            ));
-        }
-
         if valid_signature {
-            if let Some(context) = util::parse_inputs(item.sig.inputs, &name) {
-                let FilterAttrs { cfgs, attrs, .. } = util::filter_attributes(item.attrs);
+            if let Some((context, Ok(rest))) = util::parse_inputs(item.sig.inputs, &name) {
+                if rest.is_empty() {
+                    let FilterAttrs { cfgs, attrs, .. } = util::filter_attributes(item.attrs);
 
-                return Ok(HardwareTask {
-                    args,
-                    cfgs,
-                    attrs,
-                    context,
-                    stmts: Vec::<Stmt>::new(),
-                    is_extern: true,
-                });
+                    return Ok(HardwareTask {
+                        args,
+                        cfgs,
+                        attrs,
+                        context,
+                        stmts: Vec::<Stmt>::new(),
+                        is_extern: true,
+                    });
+                }
             }
         }
 
diff --git a/macros/src/syntax/parse/idle.rs b/macros/src/syntax/parse/idle.rs
index f049cca8..124c1366 100644
--- a/macros/src/syntax/parse/idle.rs
+++ b/macros/src/syntax/parse/idle.rs
@@ -21,14 +21,16 @@ impl Idle {
         let name = item.sig.ident.to_string();
 
         if valid_signature {
-            if let Some(context) = util::parse_inputs(item.sig.inputs, &name) {
-                return Ok(Idle {
-                    args,
-                    attrs: item.attrs,
-                    context,
-                    name: item.sig.ident,
-                    stmts: item.block.stmts,
-                });
+            if let Some((context, Ok(rest))) = util::parse_inputs(item.sig.inputs, &name) {
+                if rest.is_empty() {
+                    return Ok(Idle {
+                        args,
+                        attrs: item.attrs,
+                        context,
+                        name: item.sig.ident,
+                        stmts: item.block.stmts,
+                    });
+                }
             }
         }
 
diff --git a/macros/src/syntax/parse/init.rs b/macros/src/syntax/parse/init.rs
index 23130c85..0aea20bd 100644
--- a/macros/src/syntax/parse/init.rs
+++ b/macros/src/syntax/parse/init.rs
@@ -25,16 +25,18 @@ impl Init {
             if let Ok((user_shared_struct, user_local_struct)) =
                 util::type_is_init_return(&item.sig.output)
             {
-                if let Some(context) = util::parse_inputs(item.sig.inputs, &name) {
-                    return Ok(Init {
-                        args,
-                        attrs: item.attrs,
-                        context,
-                        name: item.sig.ident,
-                        stmts: item.block.stmts,
-                        user_shared_struct,
-                        user_local_struct,
-                    });
+                if let Some((context, Ok(rest))) = util::parse_inputs(item.sig.inputs, &name) {
+                    if rest.is_empty() {
+                        return Ok(Init {
+                            args,
+                            attrs: item.attrs,
+                            context,
+                            name: item.sig.ident,
+                            stmts: item.block.stmts,
+                            user_shared_struct,
+                            user_local_struct,
+                        });
+                    }
                 }
             }
         }
diff --git a/macros/src/syntax/parse/software_task.rs b/macros/src/syntax/parse/software_task.rs
index fb9b37c4..769aa653 100644
--- a/macros/src/syntax/parse/software_task.rs
+++ b/macros/src/syntax/parse/software_task.rs
@@ -17,7 +17,7 @@ impl SoftwareTask {
         let name = item.sig.ident.to_string();
 
         if valid_signature {
-            if let Some(context) = util::parse_inputs(item.sig.inputs, &name) {
+            if let Some((context, Ok(inputs))) = util::parse_inputs(item.sig.inputs, &name) {
                 let FilterAttrs { cfgs, attrs, .. } = util::filter_attributes(item.attrs);
 
                 return Ok(SoftwareTask {
@@ -25,6 +25,7 @@ impl SoftwareTask {
                     attrs,
                     cfgs,
                     context,
+                    inputs,
                     stmts: item.block.stmts,
                     is_extern: false,
                 });
@@ -33,7 +34,7 @@ impl SoftwareTask {
 
         Err(parse::Error::new(
             span,
-            format!("this task handler must have type signature `async fn({name}::Context)`"),
+            format!("this task handler must have type signature `async fn({name}::Context, ..)`"),
         ))
     }
 }
@@ -52,7 +53,7 @@ impl SoftwareTask {
         let name = item.sig.ident.to_string();
 
         if valid_signature {
-            if let Some(context) = util::parse_inputs(item.sig.inputs, &name) {
+            if let Some((context, Ok(inputs))) = util::parse_inputs(item.sig.inputs, &name) {
                 let FilterAttrs { cfgs, attrs, .. } = util::filter_attributes(item.attrs);
 
                 return Ok(SoftwareTask {
@@ -60,6 +61,7 @@ impl SoftwareTask {
                     attrs,
                     cfgs,
                     context,
+                    inputs,
                     stmts: Vec::<Stmt>::new(),
                     is_extern: true,
                 });
@@ -68,7 +70,7 @@ impl SoftwareTask {
 
         Err(parse::Error::new(
             span,
-            format!("this task handler must have type signature `async fn({name}::Context)`"),
+            format!("this task handler must have type signature `async fn({name}::Context, ..)`"),
         ))
     }
 }
diff --git a/macros/src/syntax/parse/util.rs b/macros/src/syntax/parse/util.rs
index 900ef9d6..5a5e0c0e 100644
--- a/macros/src/syntax/parse/util.rs
+++ b/macros/src/syntax/parse/util.rs
@@ -3,8 +3,8 @@ use syn::{
     parse::{self, ParseStream},
     punctuated::Punctuated,
     spanned::Spanned,
-    Abi, AttrStyle, Attribute, Expr, FnArg, ForeignItemFn, Ident, ItemFn, Pat, Path, PathArguments,
-    ReturnType, Token, Type, Visibility,
+    Abi, AttrStyle, Attribute, Expr, FnArg, ForeignItemFn, Ident, ItemFn, Pat, PatType, Path,
+    PathArguments, ReturnType, Token, Type, Visibility,
 };
 
 use crate::syntax::{
@@ -231,19 +231,29 @@ pub fn parse_local_resources(content: ParseStream<'_>) -> parse::Result<LocalRes
     Ok(resources)
 }
 
-pub fn parse_inputs(inputs: Punctuated<FnArg, Token![,]>, name: &str) -> Option<Box<Pat>> {
+type ParseInputResult = Option<(Box<Pat>, Result<Vec<PatType>, FnArg>)>;
+
+pub fn parse_inputs(inputs: Punctuated<FnArg, Token![,]>, name: &str) -> ParseInputResult {
     let mut inputs = inputs.into_iter();
 
-    if let Some(FnArg::Typed(first)) = inputs.next() {
-        if type_is_path(&first.ty, &[name, "Context"]) {
-            // No more inputs
-            if inputs.next().is_none() {
-                return Some(first.pat);
+    match inputs.next() {
+        Some(FnArg::Typed(first)) => {
+            if type_is_path(&first.ty, &[name, "Context"]) {
+                let rest = inputs
+                    .map(|arg| match arg {
+                        FnArg::Typed(arg) => Ok(arg),
+                        _ => Err(arg),
+                    })
+                    .collect::<Result<Vec<_>, _>>();
+
+                Some((first.pat, rest))
+            } else {
+                None
             }
         }
-    }
 
-    None
+        _ => None,
+    }
 }
 
 pub fn type_is_bottom(ty: &ReturnType) -> bool {