subplotlib: use state for context

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

4 files changed, 79 insertions, 48 deletions

diff --git a/subplotlib/Cargo.toml b/subplotlib/Cargo.toml
index abea4c7..5ca2616 100644
--- a/subplotlib/Cargo.toml
+++ b/subplotlib/Cargo.toml
@@ -10,3 +10,4 @@ fehler = "1"
 subplotlib-derive = { version="0.1", path = "../subplotlib-derive" }
 lazy_static = "1"
 base64 = "0.13"
+state = "0.4"
\ No newline at end of file
diff --git a/subplotlib/src/prelude.rs b/subplotlib/src/prelude.rs
index d093e0f..88dd089 100644
--- a/subplotlib/src/prelude.rs
+++ b/subplotlib/src/prelude.rs
@@ -31,6 +31,8 @@ pub use crate::file::SubplotDataFile;
 // Finally a scenario itself
 #[doc(inline)]
 pub use crate::scenario::Scenario;
+#[doc(inline)]
+pub use crate::scenario::ScenarioContext;
 
 // Utility functions
 #[doc(inline)]
diff --git a/subplotlib/src/scenario.rs b/subplotlib/src/scenario.rs
index 7b2b1ba..d8bdd76 100644
--- a/subplotlib/src/scenario.rs
+++ b/subplotlib/src/scenario.rs
@@ -4,27 +4,82 @@
 //! Instead instances of it are constructed in the generated test functions and
 //! will be run automatically.
 
+use std::sync::Mutex;
+
+use state::Container;
+
 use crate::step::ScenarioStep;
-use crate::types::StepError;
+use crate::types::{StepError, StepResult};
+
+/// A scenario context wrapper for a particular context type
+struct ScenarioContextItem<C>(Mutex<C>);
+/// A container for all scenario contexts
+///
+/// This container allows the running of code within a given scenario context.
+pub struct ScenarioContext {
+    inner: Container,
+}
+
+impl ScenarioContext {
+    fn new() -> Self {
+        Self {
+            inner: Container::new(),
+        }
+    }
+
+    /// With the extracted immutable context, run the function f.
+    pub fn with<C, F>(&self, func: F, permit_poison: bool) -> StepResult
+    where
+        F: Fn(&C) -> StepResult,
+        C: Default + Send + 'static,
+    {
+        self.with_mut(|c: &mut C| func(&*c), permit_poison)
+    }
+
+    /// With the extracted mutable context, run the function f.
+    pub fn with_mut<C, F>(&self, func: F, permit_poison: bool) -> StepResult
+    where
+        F: Fn(&mut C) -> StepResult,
+        C: Default + Send + 'static,
+    {
+        let sci: &ScenarioContextItem<C> = if let Some(sci) = self.inner.try_get() {
+            sci
+        } else {
+            self.inner
+                .set(ScenarioContextItem(Mutex::new(C::default())));
+            self.inner.get()
+        };
+        let mut lock = match sci.0.lock() {
+            Ok(lock) => lock,
+            Err(pe) => {
+                if permit_poison {
+                    pe.into_inner()
+                } else {
+                    return Err("step paniced".into());
+                }
+            }
+        };
+        func(&mut lock)
+    }
+}
 
 /// The embodiment of a subplot scenario
 ///
 /// Scenario objects are built up by the generated test functions and then run
 /// to completion.  In rare cases they may be built up and cached for reuse
 /// for example if a scenario is a subroutine.
-pub struct Scenario<C> {
+pub struct Scenario {
     title: String,
-    steps: Vec<(ScenarioStep<C>, Option<ScenarioStep<C>>)>,
+    contexts: ScenarioContext,
+    steps: Vec<(ScenarioStep, Option<ScenarioStep>)>,
 }
 
-impl<C> Scenario<C>
-where
-    C: Default,
-{
+impl Scenario {
     /// Create a new scenario with the given title
     pub fn new(title: &str) -> Self {
         Self {
             title: title.to_string(),
+            contexts: ScenarioContext::new(),
             steps: Vec::new(),
         }
     }
@@ -35,7 +90,7 @@ where
     }
 
     /// Add a scenario step, with optional cleanup step function.
-    pub fn add_step(&mut self, step: ScenarioStep<C>, cleanup: Option<ScenarioStep<C>>) {
+    pub fn add_step(&mut self, step: ScenarioStep, cleanup: Option<ScenarioStep>) {
         self.steps.push((step, cleanup));
     }
 
@@ -51,10 +106,9 @@ where
     ///
     pub fn run(self) -> Result<(), StepError> {
         let mut highest = None;
-        let mut context = C::default();
         let mut ret = Ok(());
         for (i, step) in self.steps.iter().map(|(step, _)| step).enumerate() {
-            let res = step.call(&mut context);
+            let res = step.call(&self.contexts);
             if res.is_err() {
                 ret = res;
                 break;
@@ -64,7 +118,7 @@ where
         if let Some(n) = highest {
             for stepn in (0..=n).rev() {
                 if let (_, Some(cleanup)) = &self.steps[stepn] {
-                    if let Err(e) = cleanup.call(&mut context) {
+                    if let Err(e) = cleanup.call(&self.contexts) {
                         panic!("Failure during cleanup: {:?}", e);
                     }
                 }
diff --git a/subplotlib/src/step.rs b/subplotlib/src/step.rs
index 6b4c334..3eb6a0c 100644
--- a/subplotlib/src/step.rs
+++ b/subplotlib/src/step.rs
@@ -4,65 +4,39 @@
 //! here to provide wrappers to make the scenarios easier to work with.
 //!
 
+use crate::scenario::ScenarioContext;
 use crate::types::StepResult;
 
-/// Step functions
-///
-/// Each step function might take a mutable or immutable borrow of the context
-/// which is used in the scenario.
-enum StepFunctions<C> {
-    Immutable(Box<dyn Fn(&C) -> StepResult>),
-    Mutable(Box<dyn Fn(&mut C) -> StepResult>),
-}
-
 /// A ScenarioStep is one step in a scenario.
 ///
 /// In essence, a scenario step is a named closure.  Its name can be used when
 /// reporting an error encountered in running a scenario.
-pub struct ScenarioStep<C> {
+pub struct ScenarioStep {
     name: String,
-    func: StepFunctions<C>,
+    func: Box<dyn Fn(&ScenarioContext) -> StepResult>,
 }
 
-impl<C> ScenarioStep<C> {
-    /// Create a new scenario step taking an immutable context
+impl ScenarioStep {
+    /// Create a new scenario step taking the scenario context
     ///
     /// This is used to construct a scenario step from a function which
-    /// takes an immutable borrow of the context.  This will generally be
+    /// takes the scenario context container.  This will generally be
     /// called from the generated build method for the step.
-    pub fn new_immutable<F>(name: &str, func: F) -> Self
+    pub fn new<F>(name: &str, func: F) -> Self
     where
-        F: Fn(&C) -> StepResult + 'static,
+        F: Fn(&ScenarioContext) -> StepResult + 'static,
     {
         Self {
             name: name.to_string(),
-            func: StepFunctions::Immutable(Box::new(func)),
-        }
-    }
-
-    /// Create a new scenario step taking a mutable context
-    ///
-    /// This is used to construct a scenario step from a functino which
-    /// takes a mutable borrow of the context.  This will generally be
-    /// called from the generated build method for the step.
-    pub fn new_mutable<F>(name: &str, func: F) -> Self
-    where
-        F: Fn(&mut C) -> StepResult + 'static,
-    {
-        Self {
-            name: name.to_string(),
-            func: StepFunctions::Mutable(Box::new(func)),
+            func: Box::new(func),
         }
     }
 
     /// Call the step function
     ///
     /// This simply calls the encased step function
-    pub fn call(&self, context: &mut C) -> StepResult {
-        match &self.func {
-            StepFunctions::Immutable(f) => f(context),
-            StepFunctions::Mutable(f) => f(context),
-        }
+    pub fn call(&self, context: &ScenarioContext) -> StepResult {
+        (*self.func)(context)
     }
 
     /// Return the name of this step