Skip to main content

jpx_core/extensions/
utility.rs

1//! Utility functions.
2
3use std::collections::HashSet;
4
5use serde_json::{Number, Value};
6
7use crate::functions::Function;
8use crate::interpreter::SearchResult;
9use crate::registry::register_if_enabled;
10use crate::{Context, Runtime, arg, defn};
11
12// =============================================================================
13// now(fallback?) -> number (Unix timestamp in seconds)
14// =============================================================================
15
16defn!(NowFn, vec![], Some(arg!(number)));
17
18impl Function for NowFn {
19    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
20        self.signature.validate(args, ctx)?;
21
22        if let Some(fallback) = args.first()
23            && let Some(n) = fallback.as_f64()
24        {
25            return Ok(Value::Number(
26                Number::from_f64(n).unwrap_or_else(|| Number::from(0)),
27            ));
28        }
29
30        let timestamp = std::time::SystemTime::now()
31            .duration_since(std::time::UNIX_EPOCH)
32            .map(|d| d.as_secs())
33            .unwrap_or(0);
34
35        Ok(Value::Number(Number::from(timestamp)))
36    }
37}
38
39// =============================================================================
40// now_ms(fallback?) -> number (Unix timestamp in milliseconds)
41// =============================================================================
42
43defn!(NowMsFn, vec![], Some(arg!(number)));
44
45impl Function for NowMsFn {
46    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
47        self.signature.validate(args, ctx)?;
48
49        if let Some(fallback) = args.first()
50            && let Some(n) = fallback.as_f64()
51        {
52            return Ok(Value::Number(
53                Number::from_f64(n).unwrap_or_else(|| Number::from(0)),
54            ));
55        }
56
57        let timestamp = std::time::SystemTime::now()
58            .duration_since(std::time::UNIX_EPOCH)
59            .map(|d| d.as_millis() as u64)
60            .unwrap_or(0);
61
62        Ok(Value::Number(Number::from(timestamp)))
63    }
64}
65
66// =============================================================================
67// default(value, default_value) -> value if not null, else default
68// =============================================================================
69
70defn!(DefaultFn, vec![arg!(any), arg!(any)], None);
71
72impl Function for DefaultFn {
73    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
74        self.signature.validate(args, ctx)?;
75
76        if args[0].is_null() {
77            Ok(args[1].clone())
78        } else {
79            Ok(args[0].clone())
80        }
81    }
82}
83
84// =============================================================================
85// if(condition, then_value, else_value) -> any
86// =============================================================================
87
88defn!(IfFn, vec![arg!(any), arg!(any), arg!(any)], None);
89
90impl Function for IfFn {
91    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
92        self.signature.validate(args, ctx)?;
93
94        let condition = &args[0];
95        let then_value = &args[1];
96        let else_value = &args[2];
97
98        let is_truthy = match condition {
99            Value::Bool(b) => *b,
100            Value::Null => false,
101            _ => true,
102        };
103
104        if is_truthy {
105            Ok(then_value.clone())
106        } else {
107            Ok(else_value.clone())
108        }
109    }
110}
111
112// =============================================================================
113// coalesce(...) -> any (first non-null value)
114// =============================================================================
115
116defn!(CoalesceFn, vec![], Some(arg!(any)));
117
118impl Function for CoalesceFn {
119    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
120        self.signature.validate(args, ctx)?;
121        for arg in args {
122            if !arg.is_null() {
123                return Ok(arg.clone());
124            }
125        }
126        Ok(Value::Null)
127    }
128}
129
130// =============================================================================
131// json_encode(any) -> string
132// =============================================================================
133
134defn!(JsonEncodeFn, vec![arg!(any)], None);
135
136impl Function for JsonEncodeFn {
137    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
138        self.signature.validate(args, ctx)?;
139
140        let json_str = serde_json::to_string(&args[0])
141            .map_err(|_| crate::functions::custom_error(ctx, "Failed to encode as JSON"))?;
142
143        Ok(Value::String(json_str))
144    }
145}
146
147// =============================================================================
148// pretty(any, indent?) -> string
149// =============================================================================
150
151defn!(PrettyFn, vec![arg!(any)], Some(arg!(number)));
152
153impl Function for PrettyFn {
154    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
155        self.signature.validate(args, ctx)?;
156
157        let indent = if args.len() > 1 {
158            args[1].as_f64().unwrap_or(2.0) as usize
159        } else {
160            2
161        };
162
163        // For default indent of 2, use built-in to_string_pretty
164        if indent == 2 {
165            let pretty_str = serde_json::to_string_pretty(&args[0])
166                .map_err(|_| crate::functions::custom_error(ctx, "Failed to serialize as JSON"))?;
167            return Ok(Value::String(pretty_str));
168        }
169
170        // For custom indent, manually format
171        let json_str = serde_json::to_string(&args[0])
172            .map_err(|_| crate::functions::custom_error(ctx, "Failed to serialize as JSON"))?;
173
174        let pretty_str = pretty_print_json(&json_str, indent);
175        Ok(Value::String(pretty_str))
176    }
177}
178
179/// Pretty print JSON with custom indentation.
180fn pretty_print_json(json: &str, indent_size: usize) -> String {
181    let mut result = String::new();
182    let mut depth = 0;
183    let mut in_string = false;
184    let mut escape_next = false;
185    let indent = " ".repeat(indent_size);
186
187    for ch in json.chars() {
188        if escape_next {
189            result.push(ch);
190            escape_next = false;
191            continue;
192        }
193
194        if ch == '\\' && in_string {
195            result.push(ch);
196            escape_next = true;
197            continue;
198        }
199
200        if ch == '"' {
201            in_string = !in_string;
202            result.push(ch);
203            continue;
204        }
205
206        if in_string {
207            result.push(ch);
208            continue;
209        }
210
211        match ch {
212            '{' | '[' => {
213                result.push(ch);
214                depth += 1;
215                result.push('\n');
216                for _ in 0..depth {
217                    result.push_str(&indent);
218                }
219            }
220            '}' | ']' => {
221                depth -= 1;
222                result.push('\n');
223                for _ in 0..depth {
224                    result.push_str(&indent);
225                }
226                result.push(ch);
227            }
228            ',' => {
229                result.push(ch);
230                result.push('\n');
231                for _ in 0..depth {
232                    result.push_str(&indent);
233                }
234            }
235            ':' => {
236                result.push_str(": ");
237            }
238            ' ' | '\n' | '\t' | '\r' => {
239                // Skip whitespace in compact JSON
240            }
241            _ => {
242                result.push(ch);
243            }
244        }
245    }
246
247    result
248}
249
250// =============================================================================
251// json_decode(string) -> any
252// =============================================================================
253
254defn!(JsonDecodeFn, vec![arg!(string)], None);
255
256impl Function for JsonDecodeFn {
257    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
258        self.signature.validate(args, ctx)?;
259
260        let s = args[0]
261            .as_str()
262            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected string argument"))?;
263
264        // Return null for invalid JSON instead of erroring
265        match serde_json::from_str::<Value>(s) {
266            Ok(val) => Ok(val),
267            Err(_) => Ok(Value::Null),
268        }
269    }
270}
271
272// =============================================================================
273// json_pointer(any, string) -> any (RFC 6901 JSON Pointer)
274// =============================================================================
275
276defn!(JsonPointerFn, vec![arg!(any), arg!(string)], None);
277
278impl Function for JsonPointerFn {
279    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
280        self.signature.validate(args, ctx)?;
281
282        let pointer = args[1].as_str().ok_or_else(|| {
283            crate::functions::custom_error(ctx, "Expected string pointer argument")
284        })?;
285
286        // Use serde_json's built-in pointer method directly on the Value
287        match args[0].pointer(pointer) {
288            Some(result) => Ok(result.clone()),
289            None => Ok(Value::Null),
290        }
291    }
292}
293
294// =============================================================================
295// env() -> object (all environment variables)
296// =============================================================================
297
298defn!(EnvFn, vec![], None);
299
300impl Function for EnvFn {
301    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
302        self.signature.validate(args, ctx)?;
303
304        let mut map = serde_json::Map::new();
305        for (key, value) in std::env::vars() {
306            map.insert(key, Value::String(value));
307        }
308
309        Ok(Value::Object(map))
310    }
311}
312
313// =============================================================================
314// get_env(name) -> string | null (single environment variable)
315// =============================================================================
316
317defn!(GetEnvFn, vec![arg!(string)], None);
318
319impl Function for GetEnvFn {
320    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
321        self.signature.validate(args, ctx)?;
322
323        let name = args[0]
324            .as_str()
325            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected string argument"))?;
326
327        match std::env::var(name) {
328            Ok(value) => Ok(Value::String(value)),
329            Err(_) => Ok(Value::Null),
330        }
331    }
332}
333
334/// Register utility functions that are in the enabled set.
335pub fn register_filtered(runtime: &mut Runtime, enabled: &HashSet<&str>) {
336    register_if_enabled(runtime, "now", enabled, Box::new(NowFn::new()));
337    register_if_enabled(runtime, "now_ms", enabled, Box::new(NowMsFn::new()));
338    register_if_enabled(runtime, "default", enabled, Box::new(DefaultFn::new()));
339    register_if_enabled(runtime, "if", enabled, Box::new(IfFn::new()));
340    register_if_enabled(runtime, "coalesce", enabled, Box::new(CoalesceFn::new()));
341    register_if_enabled(
342        runtime,
343        "json_encode",
344        enabled,
345        Box::new(JsonEncodeFn::new()),
346    );
347    register_if_enabled(runtime, "to_json", enabled, Box::new(JsonEncodeFn::new()));
348    register_if_enabled(
349        runtime,
350        "json_decode",
351        enabled,
352        Box::new(JsonDecodeFn::new()),
353    );
354    register_if_enabled(
355        runtime,
356        "json_pointer",
357        enabled,
358        Box::new(JsonPointerFn::new()),
359    );
360    register_if_enabled(runtime, "pretty", enabled, Box::new(PrettyFn::new()));
361    register_if_enabled(runtime, "env", enabled, Box::new(EnvFn::new()));
362    register_if_enabled(runtime, "get_env", enabled, Box::new(GetEnvFn::new()));
363}