Skip to main content

stormchaser_engine/
stdlib.rs

1use anyhow::Result;
2use base64::{engine::general_purpose, Engine as _};
3use hcl::eval::{Context as HclContext, FuncArgs, FuncDef, ParamType};
4use hcl::Value;
5
6pub fn register_stdlib(ctx: &mut HclContext) {
7    // String functions
8    ctx.declare_func(
9        "upper",
10        FuncDef::builder().param(ParamType::String).build(upper),
11    );
12    ctx.declare_func(
13        "lower",
14        FuncDef::builder().param(ParamType::String).build(lower),
15    );
16    ctx.declare_func(
17        "trim",
18        FuncDef::builder()
19            .param(ParamType::String)
20            .param(ParamType::String)
21            .build(trim),
22    );
23    ctx.declare_func(
24        "trimspace",
25        FuncDef::builder().param(ParamType::String).build(trimspace),
26    );
27    ctx.declare_func(
28        "split",
29        FuncDef::builder()
30            .param(ParamType::String)
31            .param(ParamType::String)
32            .build(split),
33    );
34    ctx.declare_func(
35        "join",
36        FuncDef::builder()
37            .param(ParamType::String)
38            .param(ParamType::Array(Box::new(ParamType::Any)))
39            .build(join),
40    );
41    ctx.declare_func(
42        "replace",
43        FuncDef::builder()
44            .param(ParamType::String)
45            .param(ParamType::String)
46            .param(ParamType::String)
47            .build(replace),
48    );
49
50    // Encoding functions
51    ctx.declare_func(
52        "jsonencode",
53        FuncDef::builder().param(ParamType::Any).build(jsonencode),
54    );
55    ctx.declare_func(
56        "jsondecode",
57        FuncDef::builder()
58            .param(ParamType::String)
59            .build(jsondecode),
60    );
61    ctx.declare_func(
62        "base64encode",
63        FuncDef::builder()
64            .param(ParamType::String)
65            .build(base64encode),
66    );
67    ctx.declare_func(
68        "base64decode",
69        FuncDef::builder()
70            .param(ParamType::String)
71            .build(base64decode),
72    );
73
74    // Collection functions
75    ctx.declare_func(
76        "length",
77        FuncDef::builder().param(ParamType::Any).build(length),
78    );
79    ctx.declare_func(
80        "keys",
81        FuncDef::builder()
82            .param(ParamType::Object(Box::new(ParamType::Any)))
83            .build(keys),
84    );
85    ctx.declare_func(
86        "values",
87        FuncDef::builder()
88            .param(ParamType::Object(Box::new(ParamType::Any)))
89            .build(values),
90    );
91    ctx.declare_func(
92        "contains",
93        FuncDef::builder()
94            .param(ParamType::Array(Box::new(ParamType::Any)))
95            .param(ParamType::Any)
96            .build(contains),
97    );
98
99    // Logic/Type functions
100    ctx.declare_func(
101        "coalesce",
102        FuncDef::builder()
103            .variadic_param(ParamType::Any)
104            .build(coalesce),
105    );
106}
107
108fn upper(args: FuncArgs) -> Result<Value, String> {
109    if let Some(Value::String(s)) = args.first() {
110        Ok(Value::String(s.to_uppercase()))
111    } else {
112        Err("upper() expects a string argument".to_string())
113    }
114}
115
116fn lower(args: FuncArgs) -> Result<Value, String> {
117    if let Some(Value::String(s)) = args.first() {
118        Ok(Value::String(s.to_lowercase()))
119    } else {
120        Err("lower() expects a string argument".to_string())
121    }
122}
123
124fn trim(args: FuncArgs) -> Result<Value, String> {
125    if args.len() != 2 {
126        return Err("trim() expects exactly 2 arguments".to_string());
127    }
128    match (&args[0], &args[1]) {
129        (Value::String(s), Value::String(cutset)) => {
130            let cutset_chars: Vec<char> = cutset.chars().collect();
131            Ok(Value::String(
132                s.trim_matches(|c| cutset_chars.contains(&c)).to_string(),
133            ))
134        }
135        _ => Err("trim() expects string arguments".to_string()),
136    }
137}
138
139fn trimspace(args: FuncArgs) -> Result<Value, String> {
140    if let Some(Value::String(s)) = args.first() {
141        Ok(Value::String(s.trim().to_string()))
142    } else {
143        Err("trimspace() expects a string argument".to_string())
144    }
145}
146
147fn split(args: FuncArgs) -> Result<Value, String> {
148    if args.len() != 2 {
149        return Err("split() expects exactly 2 arguments".to_string());
150    }
151    match (&args[0], &args[1]) {
152        (Value::String(sep), Value::String(s)) => {
153            let parts: Vec<Value> = s
154                .split(sep)
155                .map(|part| Value::String(part.to_string()))
156                .collect();
157            Ok(Value::Array(parts))
158        }
159        _ => Err("split() expects string arguments".to_string()),
160    }
161}
162
163fn join(args: FuncArgs) -> Result<Value, String> {
164    if args.len() != 2 {
165        return Err("join() expects exactly 2 arguments".to_string());
166    }
167    match (&args[0], &args[1]) {
168        (Value::String(sep), Value::Array(arr)) => {
169            let mut strings = Vec::new();
170            for item in arr {
171                match item {
172                    Value::String(s) => strings.push(s.clone()),
173                    _ => strings.push(item.to_string()),
174                }
175            }
176            Ok(Value::String(strings.join(sep)))
177        }
178        _ => Err("join() expects a string and an array".to_string()),
179    }
180}
181
182fn replace(args: FuncArgs) -> Result<Value, String> {
183    if args.len() != 3 {
184        return Err("replace() expects exactly 3 arguments".to_string());
185    }
186    match (&args[0], &args[1], &args[2]) {
187        (Value::String(s), Value::String(old), Value::String(new)) => {
188            Ok(Value::String(s.replace(old, new)))
189        }
190        _ => Err("replace() expects string arguments".to_string()),
191    }
192}
193
194fn jsonencode(args: FuncArgs) -> Result<Value, String> {
195    if args.len() != 1 {
196        return Err("jsonencode() expects exactly 1 argument".to_string());
197    }
198    match serde_json::to_string(&args[0]) {
199        Ok(s) => Ok(Value::String(s)),
200        Err(e) => Err(format!("jsonencode() failed: {}", e)),
201    }
202}
203
204fn jsondecode(args: FuncArgs) -> Result<Value, String> {
205    if let Some(Value::String(s)) = args.first() {
206        match serde_json::from_str::<Value>(s) {
207            Ok(v) => Ok(v),
208            Err(e) => Err(format!("jsondecode() failed: {}", e)),
209        }
210    } else {
211        Err("jsondecode() expects a string argument".to_string())
212    }
213}
214
215fn base64encode(args: FuncArgs) -> Result<Value, String> {
216    if let Some(Value::String(s)) = args.first() {
217        Ok(Value::String(general_purpose::STANDARD.encode(s)))
218    } else {
219        Err("base64encode() expects a string argument".to_string())
220    }
221}
222
223fn base64decode(args: FuncArgs) -> Result<Value, String> {
224    if let Some(Value::String(s)) = args.first() {
225        match general_purpose::STANDARD.decode(s) {
226            Ok(bytes) => match String::from_utf8(bytes) {
227                Ok(decoded) => Ok(Value::String(decoded)),
228                Err(_) => Err("base64decode() resulted in invalid UTF-8".to_string()),
229            },
230            Err(e) => Err(format!("base64decode() failed: {}", e)),
231        }
232    } else {
233        Err("base64decode() expects a string argument".to_string())
234    }
235}
236
237fn length(args: FuncArgs) -> Result<Value, String> {
238    if args.len() != 1 {
239        return Err("length() expects exactly 1 argument".to_string());
240    }
241    match &args[0] {
242        Value::String(s) => Ok(Value::Number(hcl::Number::from(s.len() as u64))),
243        Value::Array(a) => Ok(Value::Number(hcl::Number::from(a.len() as u64))),
244        Value::Object(o) => Ok(Value::Number(hcl::Number::from(o.len() as u64))),
245        _ => Err("length() expects a string, array, or object".to_string()),
246    }
247}
248
249fn keys(args: FuncArgs) -> Result<Value, String> {
250    if let Some(Value::Object(o)) = args.first() {
251        let keys_arr: Vec<Value> = o.keys().map(|k| Value::String(k.to_string())).collect();
252        Ok(Value::Array(keys_arr))
253    } else {
254        Err("keys() expects an object argument".to_string())
255    }
256}
257
258fn values(args: FuncArgs) -> Result<Value, String> {
259    if let Some(Value::Object(o)) = args.first() {
260        let vals_arr: Vec<Value> = o.values().cloned().collect();
261        Ok(Value::Array(vals_arr))
262    } else {
263        Err("values() expects an object argument".to_string())
264    }
265}
266
267fn contains(args: FuncArgs) -> Result<Value, String> {
268    if args.len() != 2 {
269        return Err("contains() expects exactly 2 arguments".to_string());
270    }
271    if let Value::Array(arr) = &args[0] {
272        let target = &args[1];
273        Ok(Value::Bool(arr.contains(target)))
274    } else {
275        Err("contains() expects an array as its first argument".to_string())
276    }
277}
278
279/// Returns the first argument that is neither `null` nor an empty string `""`.
280///
281/// This matches Terraform's `coalesce` semantics: for non-string types (numbers,
282/// booleans, arrays, objects) only `null` is skipped — `0`, `false`, and empty
283/// collections are all considered valid non-empty values.
284fn coalesce(args: FuncArgs) -> Result<Value, String> {
285    if args.is_empty() {
286        return Err("coalesce() expects at least 1 argument".to_string());
287    }
288    for arg in args.iter() {
289        if arg.is_null() {
290            continue;
291        }
292        if let Value::String(s) = arg {
293            if s.is_empty() {
294                continue;
295            }
296        }
297        return Ok(arg.clone());
298    }
299    Err("coalesce(): no non-null, non-empty-string arguments provided".to_string())
300}
301
302#[cfg(test)]
303mod tests {
304    use super::*;
305    use crate::hcl_eval::evaluate_raw_expr;
306    use hcl::eval::Context;
307    use serde_json::json;
308
309    fn eval_with_stdlib(expr: &str) -> serde_json::Value {
310        let mut ctx = Context::new();
311        register_stdlib(&mut ctx);
312        evaluate_raw_expr(expr, &ctx).expect("evaluation failed")
313    }
314
315    #[test]
316    fn test_string_functions() {
317        assert_eq!(eval_with_stdlib("upper(\"hello\")"), json!("HELLO"));
318        assert_eq!(eval_with_stdlib("lower(\"HELLO\")"), json!("hello"));
319        assert_eq!(
320            eval_with_stdlib("trim(\"?!hello?!\", \"?!\")"),
321            json!("hello")
322        );
323        assert_eq!(
324            eval_with_stdlib("trimspace(\"  hello  \\n\")"),
325            json!("hello")
326        );
327
328        let split_res = eval_with_stdlib("split(\",\", \"a,b,c\")");
329        assert_eq!(split_res, json!(["a", "b", "c"]));
330
331        assert_eq!(
332            eval_with_stdlib("join(\"-\", [\"a\", \"b\", \"c\"])"),
333            json!("a-b-c")
334        );
335        assert_eq!(
336            eval_with_stdlib("replace(\"hello world\", \"world\", \"there\")"),
337            json!("hello there")
338        );
339    }
340
341    #[test]
342    fn test_encoding_functions() {
343        let encode_res = eval_with_stdlib("jsonencode({\"a\" = 1})");
344        assert_eq!(encode_res, json!("{\"a\":1}"));
345
346        let decode_res = eval_with_stdlib("jsondecode(\"{\\\"a\\\": 1}\")");
347        assert_eq!(decode_res, json!({"a": 1}));
348
349        assert_eq!(
350            eval_with_stdlib("base64encode(\"hello\")"),
351            json!("aGVsbG8=")
352        );
353        assert_eq!(
354            eval_with_stdlib("base64decode(\"aGVsbG8=\")"),
355            json!("hello")
356        );
357    }
358
359    #[test]
360    fn test_collection_functions() {
361        assert_eq!(eval_with_stdlib("length(\"hello\")"), json!(5));
362        assert_eq!(eval_with_stdlib("length([1, 2, 3])"), json!(3));
363
364        assert_eq!(
365            eval_with_stdlib("contains([\"a\", \"b\"], \"b\")"),
366            json!(true)
367        );
368        assert_eq!(
369            eval_with_stdlib("contains([\"a\", \"b\"], \"c\")"),
370            json!(false)
371        );
372    }
373
374    #[test]
375    fn test_coalesce_skips_null_and_empty_string() {
376        // Skips null and empty strings, returns first valid string.
377        assert_eq!(
378            eval_with_stdlib("coalesce(\"\", null, \"first\", \"second\")"),
379            json!("first")
380        );
381    }
382
383    #[test]
384    fn test_coalesce_returns_zero_number() {
385        // Terraform: 0 is a valid non-empty value — must not be skipped.
386        assert_eq!(eval_with_stdlib("coalesce(0, 1)"), json!(0));
387    }
388
389    #[test]
390    fn test_coalesce_returns_false_bool() {
391        // Terraform: false is a valid non-empty value — must not be skipped.
392        assert_eq!(eval_with_stdlib("coalesce(false, true)"), json!(false));
393    }
394
395    #[test]
396    fn test_coalesce_returns_empty_array() {
397        // Terraform: empty array is a valid non-empty value — must not be skipped.
398        assert_eq!(eval_with_stdlib("coalesce([], [1, 2])"), json!([]));
399    }
400}