Skip to main content

leviath_scripting/
functions.rs

1//! Leviath functions exposed to Rhai scripts.
2
3use rhai::Engine;
4
5/// Register Leviath functions in the Rhai engine.
6pub fn register_functions(engine: &mut Engine) {
7    // String operations
8    engine.register_fn("contains", |text: &str, pattern: &str| -> bool {
9        text.contains(pattern)
10    });
11
12    engine.register_fn("starts_with", |text: &str, pattern: &str| -> bool {
13        text.starts_with(pattern)
14    });
15
16    engine.register_fn("ends_with", |text: &str, pattern: &str| -> bool {
17        text.ends_with(pattern)
18    });
19
20    engine.register_fn("trim", |text: &str| -> String { text.trim().to_string() });
21
22    engine.register_fn("join", |arr: rhai::Array, separator: &str| -> String {
23        arr.iter()
24            .map(|v| v.to_string())
25            .collect::<Vec<_>>()
26            .join(separator)
27    });
28
29    engine.register_fn("split", |text: &str, separator: &str| -> rhai::Array {
30        text.split(separator)
31            .map(|s| rhai::Dynamic::from(s.to_string()))
32            .collect()
33    });
34
35    // Token counting (approximate)
36    engine.register_fn("count_tokens", |text: &str| -> i64 {
37        leviath_core::estimate_tokens(text) as i64
38    });
39
40    // Content validation
41    engine.register_fn("is_json", |text: &str| -> bool {
42        serde_json::from_str::<serde_json::Value>(text).is_ok()
43    });
44
45    engine.register_fn("is_mermaid", |text: &str| -> bool {
46        text.contains("graph")
47            || text.contains("sequenceDiagram")
48            || text.contains("classDiagram")
49            || text.contains("stateDiagram")
50            || text.contains("erDiagram")
51            || text.contains("flowchart")
52    });
53
54    engine.register_fn("is_markdown", |text: &str| -> bool {
55        // Very permissive - just check for common markdown markers
56        text.contains("##") || text.contains("**") || text.contains("```") || !text.is_empty()
57    });
58
59    engine.register_fn("is_empty", |text: &str| -> bool { text.trim().is_empty() });
60}
61
62#[cfg(test)]
63mod tests {
64    use super::*;
65    use rhai::Engine;
66
67    fn engine() -> Engine {
68        let mut e = Engine::new();
69        register_functions(&mut e);
70        e
71    }
72
73    // --- contains ---
74
75    #[test]
76    fn contains_returns_true_when_pattern_present() {
77        let e = engine();
78        let result: bool = e.eval(r#"contains("hello world", "world")"#).unwrap();
79        assert!(result);
80    }
81
82    #[test]
83    fn contains_returns_false_when_pattern_absent() {
84        let e = engine();
85        let result: bool = e.eval(r#"contains("hello", "xyz")"#).unwrap();
86        assert!(!result);
87    }
88
89    #[test]
90    fn contains_empty_pattern_always_matches() {
91        let e = engine();
92        let result: bool = e.eval(r#"contains("hello", "")"#).unwrap();
93        assert!(result);
94    }
95
96    // --- starts_with ---
97
98    #[test]
99    fn starts_with_true() {
100        let e = engine();
101        let result: bool = e.eval(r#"starts_with("hello", "he")"#).unwrap();
102        assert!(result);
103    }
104
105    #[test]
106    fn starts_with_false() {
107        let e = engine();
108        let result: bool = e.eval(r#"starts_with("hello", "lo")"#).unwrap();
109        assert!(!result);
110    }
111
112    // --- ends_with ---
113
114    #[test]
115    fn ends_with_true() {
116        let e = engine();
117        let result: bool = e.eval(r#"ends_with("hello", "lo")"#).unwrap();
118        assert!(result);
119    }
120
121    #[test]
122    fn ends_with_false() {
123        let e = engine();
124        let result: bool = e.eval(r#"ends_with("hello", "he")"#).unwrap();
125        assert!(!result);
126    }
127
128    // --- trim ---
129
130    #[test]
131    fn trim_removes_whitespace() {
132        let e = engine();
133        let result: String = e.eval(r#"trim("  hi  ")"#).unwrap();
134        assert_eq!(result, "hi");
135    }
136
137    #[test]
138    fn trim_no_op_on_clean_string() {
139        let e = engine();
140        let result: String = e.eval(r#"trim("hi")"#).unwrap();
141        assert_eq!(result, "hi");
142    }
143
144    // --- join ---
145
146    #[test]
147    fn join_with_comma() {
148        let e = engine();
149        let result: String = e.eval(r#"join(["a", "b", "c"], ",")"#).unwrap();
150        assert_eq!(result, "a,b,c");
151    }
152
153    #[test]
154    fn join_empty_array() {
155        let e = engine();
156        let result: String = e.eval(r#"join([], ",")"#).unwrap();
157        assert_eq!(result, "");
158    }
159
160    #[test]
161    fn join_single_element() {
162        let e = engine();
163        let result: String = e.eval(r#"join(["only"], "-")"#).unwrap();
164        assert_eq!(result, "only");
165    }
166
167    // --- split ---
168
169    #[test]
170    fn split_by_comma() {
171        let e = engine();
172        let result: rhai::Array = e.eval(r#"split("a,b,c", ",")"#).unwrap();
173        assert_eq!(result.len(), 3);
174        assert_eq!(result[0].clone_cast::<String>(), "a");
175        assert_eq!(result[1].clone_cast::<String>(), "b");
176        assert_eq!(result[2].clone_cast::<String>(), "c");
177    }
178
179    #[test]
180    fn split_no_separator_found() {
181        let e = engine();
182        let result: rhai::Array = e.eval(r#"split("abc", ",")"#).unwrap();
183        assert_eq!(result.len(), 1);
184        assert_eq!(result[0].clone_cast::<String>(), "abc");
185    }
186
187    // --- count_tokens ---
188
189    #[test]
190    fn count_tokens_approximate() {
191        let e = engine();
192        let result: i64 = e.eval(r#"count_tokens("hello world")"#).unwrap();
193        // "hello world" is 11 chars, ceil(11/4) = 3
194        assert_eq!(result, 3);
195    }
196
197    #[test]
198    fn count_tokens_empty() {
199        let e = engine();
200        let result: i64 = e.eval(r#"count_tokens("")"#).unwrap();
201        assert_eq!(result, 0);
202    }
203
204    // --- is_json ---
205
206    #[test]
207    fn is_json_valid_object() {
208        let e = engine();
209        let result: bool = e.eval(r#"is_json("{}")"#).unwrap();
210        assert!(result);
211    }
212
213    #[test]
214    fn is_json_valid_array() {
215        let e = engine();
216        let result: bool = e.eval(r#"is_json("[1,2,3]")"#).unwrap();
217        assert!(result);
218    }
219
220    #[test]
221    fn is_json_invalid() {
222        let e = engine();
223        let result: bool = e.eval(r#"is_json("not json")"#).unwrap();
224        assert!(!result);
225    }
226
227    // --- is_mermaid ---
228
229    #[test]
230    fn is_mermaid_with_graph_keyword() {
231        let e = engine();
232        let result: bool = e.eval(r#"is_mermaid("graph TD; A-->B")"#).unwrap();
233        assert!(result);
234    }
235
236    #[test]
237    fn is_mermaid_with_sequence_diagram() {
238        let e = engine();
239        let result: bool = e
240            .eval(r#"is_mermaid("sequenceDiagram\nA->>B: Hi")"#)
241            .unwrap();
242        assert!(result);
243    }
244
245    #[test]
246    fn is_mermaid_with_flowchart() {
247        let e = engine();
248        let result: bool = e.eval(r#"is_mermaid("flowchart LR")"#).unwrap();
249        assert!(result);
250    }
251
252    #[test]
253    fn is_mermaid_false_for_plain_text() {
254        let e = engine();
255        let result: bool = e.eval(r#"is_mermaid("just some text")"#).unwrap();
256        assert!(!result);
257    }
258
259    // --- is_markdown ---
260
261    #[test]
262    fn is_markdown_with_heading() {
263        let e = engine();
264        let script = "is_markdown(\"## Heading\")";
265        let result: bool = e.eval(script).unwrap();
266        assert!(result);
267    }
268
269    #[test]
270    fn is_markdown_with_bold() {
271        let e = engine();
272        let result: bool = e.eval(r#"is_markdown("some **bold** text")"#).unwrap();
273        assert!(result);
274    }
275
276    #[test]
277    fn is_markdown_with_code_fence() {
278        let e = engine();
279        let result: bool = e.eval(r#"is_markdown("```code```")"#).unwrap();
280        assert!(result);
281    }
282
283    #[test]
284    fn is_markdown_plain_nonempty_text_returns_true_via_nonempty_fallback() {
285        let e = engine();
286        // Text with none of the markdown markers - falls through to !text.is_empty()
287        let result: bool = e.eval(r#"is_markdown("just plain text")"#).unwrap();
288        assert!(result);
289    }
290
291    // --- is_empty ---
292
293    #[test]
294    fn is_empty_true_for_empty() {
295        let e = engine();
296        let result: bool = e.eval(r#"is_empty("")"#).unwrap();
297        assert!(result);
298    }
299
300    #[test]
301    fn is_empty_true_for_whitespace_only() {
302        let e = engine();
303        let result: bool = e.eval(r#"is_empty("   ")"#).unwrap();
304        assert!(result);
305    }
306
307    #[test]
308    fn is_empty_false_for_content() {
309        let e = engine();
310        let result: bool = e.eval(r#"is_empty("hi")"#).unwrap();
311        assert!(!result);
312    }
313}