subplotlib: Support direct calling of other steps

Signed-off-by: Daniel Silverstone <dsilvers@digital-scurf.org>

1 files changed, 43 insertions, 5 deletions

diff --git a/subplotlib-derive/src/lib.rs b/subplotlib-derive/src/lib.rs
index adc0c01..0e7e267 100644
--- a/subplotlib-derive/src/lib.rs
+++ b/subplotlib-derive/src/lib.rs
@@ -33,6 +33,14 @@ fn ty_is_datafile(ty: &Type) -> bool {
     }
 }
 
+fn ty_is_scenariocontext(ty: &Type) -> bool {
+    if let Type::Path(ty) = ty {
+        ty.path.is_ident("ScenarioContext")
+    } else {
+        false
+    }
+}
+
 #[throws(Error)]
 fn check_step_declaration(step: &ItemFn) {
     // Step functions must be declared very simply as:
@@ -208,6 +216,12 @@ fn process_step(input: ItemFn) -> proc_macro2::TokenStream {
         }
     };
 
+    let withfn = if mutablectx {
+        Ident::new("with_mut", Span::call_site())
+    } else {
+        Ident::new("with", Span::call_site())
+    };
+
     let structimpl = {
         let fieldfns: Vec<_> = fields
             .iter()
@@ -251,10 +265,14 @@ fn process_step(input: ItemFn) -> proc_macro2::TokenStream {
 
         let stepnamestr = format!("{}", stepname);
 
-        let withfn = if mutablectx {
-            Ident::new("with_mut", Span::call_site())
+        let builder_body = if ty_is_scenariocontext(&contexttype) {
+            quote! {
+                #stepname(ctx,#(#buildargs),*)
+            }
         } else {
-            Ident::new("with", Span::call_site())
+            quote! {
+                ctx.#withfn (|ctx| #stepname(ctx, #(#buildargs),*),false)
+            }
         };
 
         quote! {
@@ -263,22 +281,42 @@ fn process_step(input: ItemFn) -> proc_macro2::TokenStream {
 
                 pub fn build(self) -> ScenarioStep {
                     ScenarioStep::new(#stepnamestr, move |ctx|
-                        ctx.#withfn (|ctx| #stepname(ctx, #(#buildargs),*),false))
+                        #builder_body
+                    )
                 }
             }
         }
     };
 
+    let inputargs: Vec<_> = fields.iter().map(|(i, t)| quote!(#i : #t)).collect();
+    let argnames: Vec<_> = fields.iter().map(|(i, _)| i).collect();
+
+    let call_body = if ty_is_scenariocontext(&contexttype) {
+        quote! {
+            #stepname(___context___,#(#argnames),*)
+        }
+    } else {
+        quote! {
+            ___context___.#withfn (move |ctx| #stepname(ctx, #(#argnames),*),false)
+        }
+    };
+
     let ret = quote! {
         #vis mod #stepname {
             use super::*;
+            pub(crate) use super::#contexttype;
+
             #structdef
             #structimpl
 
             #[throws(StepError)]
             #[allow(unused)] // It's okay for step functions to not be used
             #input
-        }
+
+            pub fn call(___context___: &ScenarioContext, #(#inputargs),*) -> StepResult {
+                #call_body
+            }
+    }
     };
 
     ret