jpx_core/extensions/
computing.rs

1//! Computing and bitwise functions.
2
3use std::collections::HashSet;
4
5use serde_json::{Number, Value};
6
7use crate::functions::{Function, number_value};
8use crate::interpreter::SearchResult;
9use crate::registry::register_if_enabled;
10use crate::{Context, Runtime, arg, defn};
11
12// Decimal units (SI): KB, MB, GB, TB, PB
13const DECIMAL_UNITS: &[(&str, f64)] = &[
14    ("PB", 1e15),
15    ("TB", 1e12),
16    ("GB", 1e9),
17    ("MB", 1e6),
18    ("KB", 1e3),
19    ("B", 1.0),
20];
21
22// Binary units (IEC): KiB, MiB, GiB, TiB, PiB
23const BINARY_UNITS: &[(&str, f64)] = &[
24    ("PiB", 1125899906842624.0),
25    ("TiB", 1099511627776.0),
26    ("GiB", 1073741824.0),
27    ("MiB", 1048576.0),
28    ("KiB", 1024.0),
29    ("B", 1.0),
30];
31
32defn!(ParseBytesFn, vec![arg!(string)], None);
33
34impl Function for ParseBytesFn {
35    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
36        self.signature.validate(args, ctx)?;
37
38        let s = args[0]
39            .as_str()
40            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected string"))?;
41
42        match parse_bytes_str(s) {
43            Some(bytes) => Ok(number_value(bytes)),
44            None => Ok(Value::Null),
45        }
46    }
47}
48
49defn!(FormatBytesFn, vec![arg!(number)], None);
50
51impl Function for FormatBytesFn {
52    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
53        self.signature.validate(args, ctx)?;
54
55        let bytes = args[0]
56            .as_f64()
57            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected number"))?;
58
59        let formatted = format_bytes_with_units(bytes, DECIMAL_UNITS);
60        Ok(Value::String(formatted))
61    }
62}
63
64defn!(FormatBytesBinaryFn, vec![arg!(number)], None);
65
66impl Function for FormatBytesBinaryFn {
67    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
68        self.signature.validate(args, ctx)?;
69
70        let bytes = args[0]
71            .as_f64()
72            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected number"))?;
73
74        let formatted = format_bytes_with_units(bytes, BINARY_UNITS);
75        Ok(Value::String(formatted))
76    }
77}
78
79defn!(BitAndFn, vec![arg!(number), arg!(number)], None);
80
81impl Function for BitAndFn {
82    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
83        self.signature.validate(args, ctx)?;
84
85        let a = args[0]
86            .as_f64()
87            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected integer"))?
88            as i64;
89
90        let b = args[1]
91            .as_f64()
92            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected integer"))?
93            as i64;
94
95        Ok(Value::Number(Number::from(a & b)))
96    }
97}
98
99defn!(BitOrFn, vec![arg!(number), arg!(number)], None);
100
101impl Function for BitOrFn {
102    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
103        self.signature.validate(args, ctx)?;
104
105        let a = args[0]
106            .as_f64()
107            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected integer"))?
108            as i64;
109
110        let b = args[1]
111            .as_f64()
112            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected integer"))?
113            as i64;
114
115        Ok(Value::Number(Number::from(a | b)))
116    }
117}
118
119defn!(BitXorFn, vec![arg!(number), arg!(number)], None);
120
121impl Function for BitXorFn {
122    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
123        self.signature.validate(args, ctx)?;
124
125        let a = args[0]
126            .as_f64()
127            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected integer"))?
128            as i64;
129
130        let b = args[1]
131            .as_f64()
132            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected integer"))?
133            as i64;
134
135        Ok(Value::Number(Number::from(a ^ b)))
136    }
137}
138
139defn!(BitNotFn, vec![arg!(number)], None);
140
141impl Function for BitNotFn {
142    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
143        self.signature.validate(args, ctx)?;
144
145        let a = args[0]
146            .as_f64()
147            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected integer"))?
148            as i64;
149
150        Ok(Value::Number(Number::from(!a)))
151    }
152}
153
154defn!(BitShiftLeftFn, vec![arg!(number), arg!(number)], None);
155
156impl Function for BitShiftLeftFn {
157    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
158        self.signature.validate(args, ctx)?;
159
160        let a = args[0]
161            .as_f64()
162            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected integer"))?
163            as i64;
164
165        let n = args[1]
166            .as_f64()
167            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected non-negative integer"))?
168            as u32;
169
170        Ok(Value::Number(Number::from(a << n)))
171    }
172}
173
174defn!(BitShiftRightFn, vec![arg!(number), arg!(number)], None);
175
176impl Function for BitShiftRightFn {
177    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
178        self.signature.validate(args, ctx)?;
179
180        let a = args[0]
181            .as_f64()
182            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected integer"))?
183            as i64;
184
185        let n = args[1]
186            .as_f64()
187            .ok_or_else(|| crate::functions::custom_error(ctx, "Expected non-negative integer"))?
188            as u32;
189
190        Ok(Value::Number(Number::from(a >> n)))
191    }
192}
193
194// Helper functions
195
196/// Parse a byte string like "1.5 GB" or "100 MiB" into bytes.
197fn parse_bytes_str(s: &str) -> Option<f64> {
198    let s = s.trim();
199    if s.is_empty() {
200        return None;
201    }
202
203    let mut num_end = 0;
204    let chars: Vec<char> = s.chars().collect();
205
206    for (i, c) in chars.iter().enumerate() {
207        if c.is_ascii_digit() || *c == '.' || *c == '-' || *c == '+' {
208            num_end = i + 1;
209        } else if !c.is_whitespace() {
210            break;
211        }
212    }
213
214    if num_end == 0 {
215        return None;
216    }
217
218    let num_str: String = chars[..num_end].iter().collect();
219    let num: f64 = num_str.trim().parse().ok()?;
220
221    let unit: String = chars[num_end..].iter().collect();
222    let unit = unit.trim().to_uppercase();
223
224    if unit.is_empty() || unit == "B" || unit == "BYTES" || unit == "BYTE" {
225        return Some(num);
226    }
227
228    let multiplier = match unit.as_str() {
229        "PIB" | "PEBIBYTE" | "PEBIBYTES" => 1125899906842624.0,
230        "TIB" | "TEBIBYTE" | "TEBIBYTES" => 1099511627776.0,
231        "GIB" | "GIBIBYTE" | "GIBIBYTES" => 1073741824.0,
232        "MIB" | "MEBIBYTE" | "MEBIBYTES" => 1048576.0,
233        "KIB" | "KIBIBYTE" | "KIBIBYTES" => 1024.0,
234        "PB" | "PETABYTE" | "PETABYTES" => 1e15,
235        "TB" | "TERABYTE" | "TERABYTES" => 1e12,
236        "GB" | "GIGABYTE" | "GIGABYTES" => 1e9,
237        "MB" | "MEGABYTE" | "MEGABYTES" => 1e6,
238        "KB" | "KILOBYTE" | "KILOBYTES" => 1e3,
239        _ => return None,
240    };
241
242    Some(num * multiplier)
243}
244
245/// Format bytes with the given unit system.
246fn format_bytes_with_units(bytes: f64, units: &[(&str, f64)]) -> String {
247    if bytes == 0.0 {
248        return "0 B".to_string();
249    }
250
251    let abs_bytes = bytes.abs();
252
253    for (unit, threshold) in units {
254        if abs_bytes >= *threshold {
255            let value = bytes / threshold;
256            let formatted = format!("{:.2}", value);
257            let formatted = formatted.trim_end_matches('0').trim_end_matches('.');
258            return format!("{} {}", formatted, unit);
259        }
260    }
261
262    format!("{} B", bytes)
263}
264
265/// Register computing functions that are in the enabled set.
266pub fn register_filtered(runtime: &mut Runtime, enabled: &HashSet<&str>) {
267    register_if_enabled(
268        runtime,
269        "parse_bytes",
270        enabled,
271        Box::new(ParseBytesFn::new()),
272    );
273    register_if_enabled(
274        runtime,
275        "format_bytes",
276        enabled,
277        Box::new(FormatBytesFn::new()),
278    );
279    register_if_enabled(
280        runtime,
281        "format_bytes_binary",
282        enabled,
283        Box::new(FormatBytesBinaryFn::new()),
284    );
285    register_if_enabled(runtime, "bit_and", enabled, Box::new(BitAndFn::new()));
286    register_if_enabled(runtime, "bit_or", enabled, Box::new(BitOrFn::new()));
287    register_if_enabled(runtime, "bit_xor", enabled, Box::new(BitXorFn::new()));
288    register_if_enabled(runtime, "bit_not", enabled, Box::new(BitNotFn::new()));
289    register_if_enabled(
290        runtime,
291        "bit_shift_left",
292        enabled,
293        Box::new(BitShiftLeftFn::new()),
294    );
295    register_if_enabled(
296        runtime,
297        "bit_shift_right",
298        enabled,
299        Box::new(BitShiftRightFn::new()),
300    );
301}
302
303#[cfg(test)]
304mod tests {
305    use crate::Runtime;
306    use serde_json::json;
307
308    fn setup_runtime() -> Runtime {
309        Runtime::builder()
310            .with_standard()
311            .with_all_extensions()
312            .build()
313    }
314
315    #[test]
316    fn test_parse_bytes() {
317        use super::{BINARY_UNITS, DECIMAL_UNITS, format_bytes_with_units, parse_bytes_str};
318
319        assert_eq!(parse_bytes_str("100"), Some(100.0));
320        assert_eq!(parse_bytes_str("100 B"), Some(100.0));
321        assert_eq!(parse_bytes_str("1 KB"), Some(1000.0));
322        assert_eq!(parse_bytes_str("1.5 KB"), Some(1500.0));
323        assert_eq!(parse_bytes_str("1 MB"), Some(1_000_000.0));
324        assert_eq!(parse_bytes_str("1.5 GB"), Some(1_500_000_000.0));
325        assert_eq!(parse_bytes_str("1 TB"), Some(1_000_000_000_000.0));
326        assert_eq!(parse_bytes_str("1 KiB"), Some(1024.0));
327        assert_eq!(parse_bytes_str("1 MiB"), Some(1_048_576.0));
328        assert_eq!(parse_bytes_str("1 GiB"), Some(1_073_741_824.0));
329        assert_eq!(parse_bytes_str("1 gb"), Some(1_000_000_000.0));
330        assert_eq!(parse_bytes_str(""), None);
331        assert_eq!(parse_bytes_str("invalid"), None);
332
333        // Also verify format_bytes_with_units helper
334        assert_eq!(format_bytes_with_units(0.0, DECIMAL_UNITS), "0 B");
335        assert_eq!(format_bytes_with_units(500.0, DECIMAL_UNITS), "500 B");
336        assert_eq!(format_bytes_with_units(1000.0, DECIMAL_UNITS), "1 KB");
337        assert_eq!(format_bytes_with_units(1500.0, DECIMAL_UNITS), "1.5 KB");
338        assert_eq!(
339            format_bytes_with_units(1_500_000_000.0, DECIMAL_UNITS),
340            "1.5 GB"
341        );
342        assert_eq!(format_bytes_with_units(1024.0, BINARY_UNITS), "1 KiB");
343        assert_eq!(format_bytes_with_units(1536.0, BINARY_UNITS), "1.5 KiB");
344        assert_eq!(
345            format_bytes_with_units(1_073_741_824.0, BINARY_UNITS),
346            "1 GiB"
347        );
348    }
349
350    #[test]
351    fn test_format_bytes_via_runtime() {
352        let runtime = setup_runtime();
353
354        let expr = runtime.compile("format_bytes(`1500`)").unwrap();
355        let result = expr.search(&json!(null)).unwrap();
356        assert_eq!(result.as_str().unwrap(), "1.5 KB");
357
358        let expr = runtime.compile("format_bytes_binary(`1024`)").unwrap();
359        let result = expr.search(&json!(null)).unwrap();
360        assert_eq!(result.as_str().unwrap(), "1 KiB");
361    }
362
363    #[test]
364    fn test_bitwise_ops() {
365        let runtime = setup_runtime();
366
367        let expr = runtime.compile("bit_and(`12`, `10`)").unwrap();
368        let result = expr.search(&json!(null)).unwrap();
369        assert_eq!(result, json!(8)); // 1100 & 1010 = 1000
370
371        let expr = runtime.compile("bit_or(`12`, `10`)").unwrap();
372        let result = expr.search(&json!(null)).unwrap();
373        assert_eq!(result, json!(14)); // 1100 | 1010 = 1110
374
375        let expr = runtime.compile("bit_xor(`12`, `10`)").unwrap();
376        let result = expr.search(&json!(null)).unwrap();
377        assert_eq!(result, json!(6)); // 1100 ^ 1010 = 0110
378
379        let expr = runtime.compile("bit_shift_left(`1`, `4`)").unwrap();
380        let result = expr.search(&json!(null)).unwrap();
381        assert_eq!(result, json!(16)); // 1 << 4 = 16
382
383        let expr = runtime.compile("bit_shift_right(`16`, `2`)").unwrap();
384        let result = expr.search(&json!(null)).unwrap();
385        assert_eq!(result, json!(4)); // 16 >> 2 = 4
386    }
387}
jpx_core/extensions/computing.rs

jpx_core/extensions/
computing.rs
