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>

2 files changed, 78 insertions, 10 deletions

diff --git a/subplotlib/src/scenario.rs b/subplotlib/src/scenario.rs
index f3b8dce..0028af9 100644
--- a/subplotlib/src/scenario.rs
+++ b/subplotlib/src/scenario.rs
@@ -4,7 +4,7 @@
 //! Instead instances of it are constructed in the generated test functions and
 //! will be run automatically.
 
-use std::sync::Mutex;
+use std::{marker::PhantomData, sync::Mutex};
 
 use state::Container;
 
@@ -40,12 +40,49 @@ where
 }
 /// A scenario context wrapper for a particular context type
 struct ScenarioContextItem<C>(Mutex<C>);
+
+/// A type hook used purely in order to be able to look up contexts in the
+/// container in order to be able to iterate them during scenario execution
+struct ScenarioContextHook<C>(PhantomData<C>);
+
+impl<C> ScenarioContextHook<C>
+where
+    C: ContextElement,
+{
+    fn new() -> Self {
+        Self(PhantomData::default())
+    }
+}
+
+/// A trait used to permit the holding of multiple hooks in one vector
+trait ScenarioContextHookKind {
+    /// Enter a step
+    fn enter_step(&self);
+
+    /// Leave a step
+    fn exit_step(&self);
+}
+
+impl<C> ScenarioContextHookKind for ScenarioContextHook<C>
+where
+    C: ContextElement,
+{
+    fn enter_step(&self) {
+        //
+    }
+
+    fn exit_step(&self) {
+        //
+    }
+}
+
 /// A container for all scenario contexts
 ///
 /// This container allows the running of code within a given scenario context.
 pub struct ScenarioContext {
     title: String,
     inner: Container,
+    hooks: Vec<Box<dyn ScenarioContextHookKind>>,
 }
 
 impl ScenarioContext {
@@ -53,6 +90,7 @@ impl ScenarioContext {
         Self {
             title: title.to_string(),
             inner: Container::new(),
+            hooks: Vec::new(),
         }
     }
 
@@ -61,6 +99,19 @@ impl ScenarioContext {
         &self.title
     }
 
+    /// Ensure a context is registered
+    pub fn register_context_type<C>(&mut self)
+    where
+        C: ContextElement,
+    {
+        let sci: Option<&ScenarioContextItem<C>> = self.inner.try_get();
+        if sci.is_none() {
+            self.inner
+                .set(ScenarioContextItem(Mutex::new(C::create(&self.title))));
+            self.hooks.push(Box::new(ScenarioContextHook::<C>::new()));
+        }
+    }
+
     /// With the extracted immutable context, run the function f.
     pub fn with<C, F>(&self, func: F, defuse_poison: bool) -> StepResult
     where
@@ -76,13 +127,10 @@ impl ScenarioContext {
         F: FnOnce(&mut C) -> StepResult,
         C: ContextElement,
     {
-        let sci: &ScenarioContextItem<C> = if let Some(sci) = self.inner.try_get() {
-            sci
-        } else {
-            self.inner
-                .set(ScenarioContextItem(Mutex::new(C::create(&self.title))));
-            self.inner.get()
-        };
+        let sci: &ScenarioContextItem<C> = self
+            .inner
+            .try_get()
+            .ok_or("required context type not registered with scenario")?;
         let mut lock = match sci.0.lock() {
             Ok(lock) => lock,
             Err(pe) => {
@@ -123,9 +171,21 @@ impl Scenario {
 
     /// Add a scenario step, with optional cleanup step function.
     pub fn add_step(&mut self, step: ScenarioStep, cleanup: Option<ScenarioStep>) {
+        step.register_contexts(self);
+        if let Some(s) = cleanup.as_ref() {
+            s.register_contexts(self)
+        }
         self.steps.push((step, cleanup));
     }
 
+    /// Register a type with the scenario contexts
+    pub fn register_context_type<C>(&mut self)
+    where
+        C: ContextElement,
+    {
+        self.contexts.register_context_type::<C>();
+    }
+
     /// Run the scenario to completion.
     ///
     /// Running the scenario to completion requires running each step in turn.
diff --git a/subplotlib/src/step.rs b/subplotlib/src/step.rs
index ca40550..c7a3078 100644
--- a/subplotlib/src/step.rs
+++ b/subplotlib/src/step.rs
@@ -7,7 +7,7 @@
 use std::any::Any;
 use std::panic::{catch_unwind, AssertUnwindSafe};
 
-use crate::scenario::ScenarioContext;
+use crate::scenario::{Scenario, ScenarioContext};
 use crate::types::StepResult;
 
 /// A ScenarioStep is one step in a scenario.
@@ -17,6 +17,7 @@ use crate::types::StepResult;
 pub struct ScenarioStep {
     name: String,
     func: Box<dyn Fn(&ScenarioContext, bool) -> StepResult>,
+    reg: Box<dyn Fn(&mut Scenario)>,
 }
 
 impl ScenarioStep {
@@ -25,13 +26,15 @@ impl ScenarioStep {
     /// This is used to construct a scenario step from a function which
     /// takes the scenario context container.  This will generally be
     /// called from the generated build method for the step.
-    pub fn new<F>(name: &str, func: F) -> Self
+    pub fn new<F, R>(name: &str, func: F, reg: R) -> Self
     where
         F: Fn(&ScenarioContext, bool) -> StepResult + 'static,
+        R: Fn(&mut Scenario) + 'static,
     {
         Self {
             name: name.to_string(),
             func: Box::new(func),
+            reg: Box::new(reg),
         }
     }
 
@@ -63,4 +66,9 @@ impl ScenarioStep {
     pub fn name(&self) -> &str {
         &self.name
     }
+
+    /// Register any context types needed by this step
+    pub fn register_contexts(&self, scenario: &mut Scenario) {
+        (*self.reg)(scenario);
+    }
 }