subplotlib: Rework contextelement to be able to register other elements

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

5 files changed, 53 insertions, 31 deletions

diff --git a/subplotlib/helpers/subplotlib_context.rs b/subplotlib/helpers/subplotlib_context.rs
index 098d9f3..b986ed2 100644
--- a/subplotlib/helpers/subplotlib_context.rs
+++ b/subplotlib/helpers/subplotlib_context.rs
@@ -21,3 +21,7 @@ impl Context {
         self.files.insert(name.to_string(), content);
     }
 }
+
+impl ContextElement for Context {
+    // An empty implementation is sufficient for now
+}
diff --git a/subplotlib/src/prelude.rs b/subplotlib/src/prelude.rs
index 88dd089..68f3ec4 100644
--- a/subplotlib/src/prelude.rs
+++ b/subplotlib/src/prelude.rs
@@ -30,6 +30,8 @@ pub use crate::file::SubplotDataFile;
 
 // Finally a scenario itself
 #[doc(inline)]
+pub use crate::scenario::ContextElement;
+#[doc(inline)]
 pub use crate::scenario::Scenario;
 #[doc(inline)]
 pub use crate::scenario::ScenarioContext;
diff --git a/subplotlib/src/scenario.rs b/subplotlib/src/scenario.rs
index 36e113e..9b95737 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::{marker::PhantomData, sync::Mutex};
+use std::{cell::RefCell, marker::PhantomData, sync::Mutex};
 
 use state::Container;
 
@@ -18,16 +18,21 @@ use crate::types::{StepError, StepResult};
 /// be aware that they won't always be called on scenario start and they will
 /// not be caught up the first time they are invoked for a step, simply expected
 /// to get on with life from their first use.
-pub trait ContextElement: Send + 'static {
-    /// Create a context element.
+pub trait ContextElement: Default + Send + 'static {
+    /// A new context element was created.
     ///
     /// In order to permit elements which for example work on disk, this
-    /// function will be invoked with the scenario's name to permit the creation
-    /// of suitably named temporary directories, logging, etc.
+    /// function will be invoked with the scenario's context to permit the
+    /// context to register other contexts it might need, or to permit the
+    /// creation of suitably named temporary directories, logging, etc.
     ///
-    /// There will be an impl of this for Default capable types should the name
-    /// be unnecessary.
-    fn create(scenario_title: &str) -> Self;
+    /// The scenario's title is available via [`scenario_context.title()`][title]
+    ///
+    /// [title]: [ScenarioContext::title]
+    #[allow(unused_variables)]
+    fn created(&mut self, scenario: &Scenario) {
+        // Nothing by default
+    }
 
     /// Scenario starts
     ///
@@ -88,14 +93,6 @@ pub trait ContextElement: Send + 'static {
     }
 }
 
-impl<T> ContextElement for T
-where
-    T: Default + Send + 'static,
-{
-    fn create(_scenario_title: &str) -> Self {
-        Self::default()
-    }
-}
 /// A scenario context wrapper for a particular context type
 struct ScenarioContextItem<C>(Mutex<C>);
 
@@ -154,7 +151,7 @@ where
 pub struct ScenarioContext {
     title: String,
     inner: Container,
-    hooks: Vec<Box<dyn ScenarioContextHookKind>>,
+    hooks: RefCell<Vec<Box<dyn ScenarioContextHookKind>>>,
 }
 
 impl ScenarioContext {
@@ -162,7 +159,7 @@ impl ScenarioContext {
         Self {
             title: title.to_string(),
             inner: Container::new(),
-            hooks: Vec::new(),
+            hooks: RefCell::new(Vec::new()),
         }
     }
 
@@ -172,15 +169,20 @@ impl ScenarioContext {
     }
 
     /// Ensure a context is registered
-    pub fn register_context_type<C>(&mut self)
+    pub(crate) fn register_context_type<C>(&self) -> bool
     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()));
+            let ctx = ScenarioContextItem(Mutex::new(C::default()));
+            self.inner.set(ctx);
+            self.hooks
+                .borrow_mut()
+                .push(Box::new(ScenarioContextHook::<C>::new()));
+            true
+        } else {
+            false
         }
     }
 
@@ -251,11 +253,21 @@ impl Scenario {
     }
 
     /// Register a type with the scenario contexts
-    pub fn register_context_type<C>(&mut self)
+    pub fn register_context_type<C>(&self)
     where
         C: ContextElement,
     {
-        self.contexts.register_context_type::<C>();
+        if self.contexts.register_context_type::<C>() {
+            self.contexts
+                .with_mut(
+                    |c: &mut C| {
+                        c.created(self);
+                        Ok(())
+                    },
+                    false,
+                )
+                .unwrap();
+        }
     }
 
     /// Run the scenario to completion.
@@ -272,7 +284,7 @@ impl Scenario {
         // Firstly, we start all the contexts
         let mut ret = Ok(());
         let mut highest_start = None;
-        for (i, hook) in self.contexts.hooks.iter().enumerate() {
+        for (i, hook) in self.contexts.hooks.borrow().iter().enumerate() {
             let res = hook.scenario_starts(&self.contexts);
             if res.is_err() {
                 ret = res;
@@ -285,7 +297,7 @@ impl Scenario {
             let mut highest = None;
             for (i, step) in self.steps.iter().map(|(step, _)| step).enumerate() {
                 let mut highest_prep = None;
-                for (i, prep) in self.contexts.hooks.iter().enumerate() {
+                for (i, prep) in self.contexts.hooks.borrow().iter().enumerate() {
                     let res = prep.step_starts(&self.contexts, step.name());
                     if res.is_err() {
                         ret = res;
@@ -303,7 +315,7 @@ impl Scenario {
                 }
                 if let Some(n) = highest_prep {
                     for hookn in (0..=n).rev() {
-                        let res = self.contexts.hooks[hookn].step_stops(&self.contexts);
+                        let res = self.contexts.hooks.borrow()[hookn].step_stops(&self.contexts);
                         ret = ret.and(res)
                     }
                 }
@@ -320,7 +332,7 @@ impl Scenario {
 
         if let Some(n) = highest_start {
             for hookn in (0..=n).rev() {
-                let res = self.contexts.hooks[hookn].scenario_stops(&self.contexts);
+                let res = self.contexts.hooks.borrow()[hookn].scenario_stops(&self.contexts);
                 ret = ret.and(res);
             }
         }
diff --git a/subplotlib/src/step.rs b/subplotlib/src/step.rs
index c7a3078..6ce3b0a 100644
--- a/subplotlib/src/step.rs
+++ b/subplotlib/src/step.rs
@@ -17,7 +17,7 @@ use crate::types::StepResult;
 pub struct ScenarioStep {
     name: String,
     func: Box<dyn Fn(&ScenarioContext, bool) -> StepResult>,
-    reg: Box<dyn Fn(&mut Scenario)>,
+    reg: Box<dyn Fn(&Scenario)>,
 }
 
 impl ScenarioStep {
@@ -29,7 +29,7 @@ impl ScenarioStep {
     pub fn new<F, R>(name: &str, func: F, reg: R) -> Self
     where
         F: Fn(&ScenarioContext, bool) -> StepResult + 'static,
-        R: Fn(&mut Scenario) + 'static,
+        R: Fn(&Scenario) + 'static,
     {
         Self {
             name: name.to_string(),
@@ -68,7 +68,7 @@ impl ScenarioStep {
     }
 
     /// Register any context types needed by this step
-    pub fn register_contexts(&self, scenario: &mut Scenario) {
+    pub fn register_contexts(&self, scenario: &Scenario) {
         (*self.reg)(scenario);
     }
 }
diff --git a/subplotlib/tests/subplotlib.rs b/subplotlib/tests/subplotlib.rs
index 84f66d6..a8a9695 100644
--- a/subplotlib/tests/subplotlib.rs
+++ b/subplotlib/tests/subplotlib.rs
@@ -27,6 +27,10 @@ impl Context {
     }
 }
 
+impl ContextElement for Context {
+    // An empty implementation is sufficient for now
+}
+
 // --------------------------------
 // This came from helpers/subplotlib_impl.rs