subplotlib: Force explicit registration of contexts

In order to move toward scenario contexts being told about
step entry/exit etc, require explicit registration of context
types, in a way which can only be done before the scenario starts.

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

1 files changed, 15 insertions, 1 deletions

diff --git a/subplotlib-derive/src/lib.rs b/subplotlib-derive/src/lib.rs
index 0a0f4f1..1f52de0 100644
--- a/subplotlib-derive/src/lib.rs
+++ b/subplotlib-derive/src/lib.rs
@@ -281,7 +281,8 @@ fn process_step(input: ItemFn) -> proc_macro2::TokenStream {
 
                 pub fn build(self) -> ScenarioStep {
                     ScenarioStep::new(#stepnamestr, move |ctx, _defuse_poison|
-                        #builder_body
+                        #builder_body,
+                        |scenario| register_contexts(scenario)
                     )
                 }
             }
@@ -301,6 +302,14 @@ fn process_step(input: ItemFn) -> proc_macro2::TokenStream {
         }
     };
 
+    let register_fn_body = if ty_is_scenariocontext(&contexttype) {
+        quote! {}
+    } else {
+        quote! {
+            scenario.register_context_type::<#contexttype>();
+        }
+    };
+
     let ret = quote! {
         #vis mod #stepname {
             use super::*;
@@ -316,6 +325,11 @@ fn process_step(input: ItemFn) -> proc_macro2::TokenStream {
             pub fn call(___context___: &ScenarioContext, #(#inputargs),*) -> StepResult {
                 #call_body
             }
+
+            #[allow(unused_variables)]
+            pub fn register_contexts(scenario: &mut Scenario) {
+                #register_fn_body
+            }
     }
     };