Skip to main content

voirs_cli/workflow/
variables.rs

1//! Variable Resolution and Substitution
2//!
3//! Handles variable resolution, substitution, and scoping for workflows.
4
5use serde::{Deserialize, Serialize};
6use std::collections::HashMap;
7
8/// Variable scope
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub enum VariableScope {
11    /// Global workflow variables
12    Global,
13    /// Step-level variables
14    Step(String),
15    /// Loop iteration variables
16    Loop { step: String, iteration: usize },
17}
18
19/// Variable resolver
20pub struct VariableResolver {
21    /// Global variables
22    global: HashMap<String, serde_json::Value>,
23    /// Step-scoped variables
24    step_scoped: HashMap<String, HashMap<String, serde_json::Value>>,
25}
26
27impl VariableResolver {
28    /// Create new resolver
29    pub fn new(global: HashMap<String, serde_json::Value>) -> Self {
30        Self {
31            global,
32            step_scoped: HashMap::new(),
33        }
34    }
35
36    /// Set a global variable
37    pub fn set_global(&mut self, name: String, value: serde_json::Value) {
38        self.global.insert(name, value);
39    }
40
41    /// Set a step-scoped variable
42    pub fn set_step_variable(&mut self, step: String, name: String, value: serde_json::Value) {
43        self.step_scoped
44            .entry(step)
45            .or_default()
46            .insert(name, value);
47    }
48
49    /// Get a variable (checks step scope first, then global)
50    pub fn get(&self, name: &str, current_step: Option<&str>) -> Option<serde_json::Value> {
51        // Check step scope first
52        if let Some(step_name) = current_step {
53            if let Some(step_vars) = self.step_scoped.get(step_name) {
54                if let Some(value) = step_vars.get(name) {
55                    return Some(value.clone());
56                }
57            }
58        }
59
60        // Check global scope
61        self.global.get(name).cloned()
62    }
63
64    /// Resolve a string with variable substitution
65    pub fn resolve_string(&self, input: &str, current_step: Option<&str>) -> String {
66        let mut result = input.to_string();
67
68        // Find all ${var} patterns
69        while let Some(start) = result.find("${") {
70            if let Some(end) = result[start..].find('}') {
71                let var_name = &result[start + 2..start + end];
72                let replacement = self
73                    .get(var_name, current_step)
74                    .and_then(|v| match v {
75                        serde_json::Value::String(s) => Some(s),
76                        serde_json::Value::Number(n) => Some(n.to_string()),
77                        serde_json::Value::Bool(b) => Some(b.to_string()),
78                        _ => None,
79                    })
80                    .unwrap_or_else(|| format!("${{{}}}", var_name));
81
82                result.replace_range(start..start + end + 1, &replacement);
83            } else {
84                break;
85            }
86        }
87
88        result
89    }
90
91    /// Resolve all variables in a parameters map
92    pub fn resolve_parameters(
93        &self,
94        params: &HashMap<String, serde_json::Value>,
95        current_step: Option<&str>,
96    ) -> HashMap<String, serde_json::Value> {
97        let mut resolved = HashMap::new();
98
99        for (key, value) in params {
100            let resolved_value = self.resolve_value(value, current_step);
101            resolved.insert(key.clone(), resolved_value);
102        }
103
104        resolved
105    }
106
107    /// Resolve a single value
108    fn resolve_value(
109        &self,
110        value: &serde_json::Value,
111        current_step: Option<&str>,
112    ) -> serde_json::Value {
113        match value {
114            serde_json::Value::String(s) => {
115                // Check if it's a variable reference
116                if let Some(var_name) = s.strip_prefix("${").and_then(|s| s.strip_suffix('}')) {
117                    self.get(var_name, current_step)
118                        .unwrap_or(serde_json::Value::Null)
119                } else {
120                    // Resolve inline substitutions
121                    serde_json::Value::String(self.resolve_string(s, current_step))
122                }
123            }
124            serde_json::Value::Array(arr) => serde_json::Value::Array(
125                arr.iter()
126                    .map(|v| self.resolve_value(v, current_step))
127                    .collect(),
128            ),
129            serde_json::Value::Object(obj) => serde_json::Value::Object(
130                obj.iter()
131                    .map(|(k, v)| (k.clone(), self.resolve_value(v, current_step)))
132                    .collect(),
133            ),
134            _ => value.clone(),
135        }
136    }
137
138    /// Get all variables in current scope
139    pub fn get_all(&self, current_step: Option<&str>) -> HashMap<String, serde_json::Value> {
140        let mut all_vars = self.global.clone();
141
142        // Merge step-scoped variables
143        if let Some(step_name) = current_step {
144            if let Some(step_vars) = self.step_scoped.get(step_name) {
145                all_vars.extend(step_vars.clone());
146            }
147        }
148
149        all_vars
150    }
151
152    /// Clear step-scoped variables for a step
153    pub fn clear_step_scope(&mut self, step: &str) {
154        self.step_scoped.remove(step);
155    }
156
157    /// Check if a variable exists
158    pub fn has_variable(&self, name: &str, current_step: Option<&str>) -> bool {
159        self.get(name, current_step).is_some()
160    }
161
162    /// List all variable names
163    pub fn list_variable_names(&self, current_step: Option<&str>) -> Vec<String> {
164        let mut names: Vec<String> = self.global.keys().cloned().collect();
165
166        if let Some(step_name) = current_step {
167            if let Some(step_vars) = self.step_scoped.get(step_name) {
168                names.extend(step_vars.keys().cloned());
169            }
170        }
171
172        names.sort();
173        names.dedup();
174        names
175    }
176}
177
178impl Default for VariableResolver {
179    fn default() -> Self {
180        Self::new(HashMap::new())
181    }
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187
188    #[test]
189    fn test_variable_resolver_creation() {
190        let resolver = VariableResolver::new(HashMap::new());
191        assert_eq!(resolver.global.len(), 0);
192    }
193
194    #[test]
195    fn test_set_and_get_global() {
196        let mut resolver = VariableResolver::new(HashMap::new());
197        resolver.set_global("key1".to_string(), serde_json::json!("value1"));
198
199        let value = resolver.get("key1", None);
200        assert!(value.is_some());
201        assert_eq!(value.unwrap().as_str().unwrap(), "value1");
202    }
203
204    #[test]
205    fn test_set_and_get_step_variable() {
206        let mut resolver = VariableResolver::new(HashMap::new());
207        resolver.set_step_variable(
208            "step1".to_string(),
209            "key1".to_string(),
210            serde_json::json!("value1"),
211        );
212
213        let value = resolver.get("key1", Some("step1"));
214        assert!(value.is_some());
215        assert_eq!(value.unwrap().as_str().unwrap(), "value1");
216
217        let value2 = resolver.get("key1", Some("step2"));
218        assert!(value2.is_none());
219    }
220
221    #[test]
222    fn test_step_scope_overrides_global() {
223        let mut global = HashMap::new();
224        global.insert("key1".to_string(), serde_json::json!("global_value"));
225
226        let mut resolver = VariableResolver::new(global);
227        resolver.set_step_variable(
228            "step1".to_string(),
229            "key1".to_string(),
230            serde_json::json!("step_value"),
231        );
232
233        // Step scope should override global
234        let value = resolver.get("key1", Some("step1"));
235        assert_eq!(value.unwrap().as_str().unwrap(), "step_value");
236
237        // Without step context, should get global
238        let value2 = resolver.get("key1", None);
239        assert_eq!(value2.unwrap().as_str().unwrap(), "global_value");
240    }
241
242    #[test]
243    fn test_resolve_string_simple() {
244        let mut resolver = VariableResolver::new(HashMap::new());
245        resolver.set_global("name".to_string(), serde_json::json!("World"));
246
247        let result = resolver.resolve_string("Hello, ${name}!", None);
248        assert_eq!(result, "Hello, World!");
249    }
250
251    #[test]
252    fn test_resolve_string_multiple_variables() {
253        let mut resolver = VariableResolver::new(HashMap::new());
254        resolver.set_global("first".to_string(), serde_json::json!("John"));
255        resolver.set_global("last".to_string(), serde_json::json!("Doe"));
256
257        let result = resolver.resolve_string("Name: ${first} ${last}", None);
258        assert_eq!(result, "Name: John Doe");
259    }
260
261    #[test]
262    fn test_resolve_string_missing_variable() {
263        let resolver = VariableResolver::new(HashMap::new());
264        let result = resolver.resolve_string("Hello, ${missing}!", None);
265        assert_eq!(result, "Hello, ${missing}!");
266    }
267
268    #[test]
269    fn test_resolve_parameters() {
270        let mut resolver = VariableResolver::new(HashMap::new());
271        resolver.set_global("voice".to_string(), serde_json::json!("en-US-neural"));
272
273        let mut params = HashMap::new();
274        params.insert("voice".to_string(), serde_json::json!("${voice}"));
275        params.insert("text".to_string(), serde_json::json!("Hello"));
276
277        let resolved = resolver.resolve_parameters(&params, None);
278
279        assert_eq!(
280            resolved.get("voice").unwrap().as_str().unwrap(),
281            "en-US-neural"
282        );
283        assert_eq!(resolved.get("text").unwrap().as_str().unwrap(), "Hello");
284    }
285
286    #[test]
287    fn test_has_variable() {
288        let mut resolver = VariableResolver::new(HashMap::new());
289        resolver.set_global("exists".to_string(), serde_json::json!("yes"));
290
291        assert!(resolver.has_variable("exists", None));
292        assert!(!resolver.has_variable("missing", None));
293    }
294
295    #[test]
296    fn test_list_variable_names() {
297        let mut resolver = VariableResolver::new(HashMap::new());
298        resolver.set_global("var1".to_string(), serde_json::json!("value1"));
299        resolver.set_global("var2".to_string(), serde_json::json!("value2"));
300        resolver.set_step_variable(
301            "step1".to_string(),
302            "var3".to_string(),
303            serde_json::json!("value3"),
304        );
305
306        let names = resolver.list_variable_names(Some("step1"));
307        assert_eq!(names.len(), 3);
308        assert!(names.contains(&"var1".to_string()));
309        assert!(names.contains(&"var2".to_string()));
310        assert!(names.contains(&"var3".to_string()));
311    }
312
313    #[test]
314    fn test_clear_step_scope() {
315        let mut resolver = VariableResolver::new(HashMap::new());
316        resolver.set_step_variable(
317            "step1".to_string(),
318            "key1".to_string(),
319            serde_json::json!("value1"),
320        );
321
322        assert!(resolver.get("key1", Some("step1")).is_some());
323
324        resolver.clear_step_scope("step1");
325
326        assert!(resolver.get("key1", Some("step1")).is_none());
327    }
328}