Skip to main content

robost_script/
lib.rs

1use std::collections::HashMap;
2
3use rhai::{Dynamic, Engine, EvalAltResult, Scope};
4use thiserror::Error;
5use tracing::instrument;
6
7#[derive(Debug, Error)]
8pub enum ScriptError {
9    #[error("script error: {0}")]
10    Eval(Box<EvalAltResult>),
11    #[error("type error: {0}")]
12    Type(String),
13}
14
15pub type Result<T> = std::result::Result<T, ScriptError>;
16
17/// Wraps a Rhai engine. Reuse one instance per scenario to share the compiled
18/// function cache across steps.
19pub struct ScriptEngine {
20    engine: Engine,
21}
22
23impl Default for ScriptEngine {
24    fn default() -> Self {
25        Self::new()
26    }
27}
28
29impl ScriptEngine {
30    pub fn new() -> Self {
31        let mut engine = Engine::new();
32        engine.set_max_expr_depths(64, 64);
33        engine.set_max_operations(1_000_000);
34        register_date_fns(&mut engine);
35        Self { engine }
36    }
37
38    /// Evaluate `expr` as a Rhai boolean expression using current scenario variables.
39    pub fn eval_bool(&self, expr: &str, vars: &HashMap<String, serde_json::Value>) -> Result<bool> {
40        let (result, _) = self.run(expr, vars)?;
41        result
42            .as_bool()
43            .map_err(|_| ScriptError::Type(format!("condition '{expr}' did not return bool")))
44    }
45
46    /// Run `source` (Rhai script) inside a sandboxed scope seeded from scenario variables.
47    /// Rhai provides no file system or OS access — the engine itself is the sandbox boundary.
48    #[instrument(name = "script_run", skip(self, source, vars))]
49    pub fn run(
50        &self,
51        source: &str,
52        vars: &HashMap<String, serde_json::Value>,
53    ) -> Result<(Dynamic, HashMap<String, Dynamic>)> {
54        let mut scope = Scope::new();
55        for (k, v) in vars {
56            scope.push_dynamic(k.clone(), json_to_dynamic(v));
57        }
58
59        let result = self
60            .engine
61            .eval_with_scope::<Dynamic>(&mut scope, source)
62            .map_err(ScriptError::Eval)?;
63
64        let exports: HashMap<String, Dynamic> = scope
65            .iter_raw()
66            .map(|(name, _, val)| (name.to_owned(), val.clone()))
67            .collect();
68
69        Ok((result, exports))
70    }
71}
72
73// ── Date helper functions registered into the Rhai engine ──────────────────
74
75fn register_date_fns(engine: &mut Engine) {
76    use chrono::{Datelike, Local, NaiveDate};
77
78    engine.register_fn("now_year", || Local::now().year() as i64);
79    engine.register_fn("now_month", || Local::now().month() as i64);
80    engine.register_fn("now_day", || Local::now().day() as i64);
81    engine.register_fn("today", || Local::now().format("%Y-%m-%d").to_string());
82
83    /// Returns `Some((next_year, next_month))` for the first day of the following month,
84    /// or `None` if any value is out of range.
85    fn next_month(year: i64, month: i64) -> Option<(i32, u32)> {
86        if !(1..=12).contains(&month) {
87            return None;
88        }
89        if year < i32::MIN as i64 || year > i32::MAX as i64 {
90            return None;
91        }
92        if month == 12 {
93            let y2 = year.checked_add(1)?;
94            if y2 > i32::MAX as i64 {
95                return None;
96            }
97            Some((y2 as i32, 1))
98        } else {
99            Some((year as i32, (month + 1) as u32))
100        }
101    }
102
103    fn valid_ym(year: i64, month: i64) -> Option<(i32, u32)> {
104        if !(1..=12).contains(&month) {
105            return None;
106        }
107        if year < i32::MIN as i64 || year > i32::MAX as i64 {
108            return None;
109        }
110        Some((year as i32, month as u32))
111    }
112
113    engine.register_fn("end_of_month", |year: i64, month: i64| -> i64 {
114        next_month(year, month)
115            .and_then(|(y2, m2)| NaiveDate::from_ymd_opt(y2, m2, 1))
116            .and_then(|d| d.pred_opt())
117            .map(|d| d.day() as i64)
118            .unwrap_or(0)
119    });
120
121    engine.register_fn(
122        "end_of_month_str",
123        |year: i64, month: i64, fmt: &str| -> String {
124            next_month(year, month)
125                .and_then(|(y2, m2)| NaiveDate::from_ymd_opt(y2, m2, 1))
126                .and_then(|d| d.pred_opt())
127                .map(|d| d.format(fmt).to_string())
128                .unwrap_or_default()
129        },
130    );
131
132    engine.register_fn("start_of_month", |year: i64, month: i64| -> i64 {
133        valid_ym(year, month)
134            .and_then(|(y, m)| NaiveDate::from_ymd_opt(y, m, 1))
135            .map(|d| d.day() as i64)
136            .unwrap_or(0)
137    });
138
139    engine.register_fn(
140        "start_of_month_str",
141        |year: i64, month: i64, fmt: &str| -> String {
142            valid_ym(year, month)
143                .and_then(|(y, m)| NaiveDate::from_ymd_opt(y, m, 1))
144                .map(|d| d.format(fmt).to_string())
145                .unwrap_or_default()
146        },
147    );
148}
149
150fn json_to_dynamic(v: &serde_json::Value) -> Dynamic {
151    match v {
152        serde_json::Value::Null => Dynamic::UNIT,
153        serde_json::Value::Bool(b) => Dynamic::from(*b),
154        serde_json::Value::Number(n) => {
155            if let Some(i) = n.as_i64() {
156                Dynamic::from(i)
157            } else {
158                Dynamic::from(n.as_f64().unwrap_or(0.0))
159            }
160        }
161        serde_json::Value::String(s) => Dynamic::from(s.clone()),
162        serde_json::Value::Array(arr) => {
163            Dynamic::from(arr.iter().map(json_to_dynamic).collect::<Vec<_>>())
164        }
165        serde_json::Value::Object(map) => {
166            let m: rhai::Map = map
167                .iter()
168                .map(|(k, v)| (k.clone().into(), json_to_dynamic(v)))
169                .collect();
170            Dynamic::from(m)
171        }
172    }
173}