Skip to main content

jpx_core/extensions/
expression.rs

1//! Expression functions.
2//!
3//! Higher-order functions that use expression references (exprefs) to apply
4//! transformations, filters, and reductions over arrays and objects.
5
6use std::collections::HashSet;
7
8use serde_json::{Map, Number, Value};
9
10use crate::ast::Ast;
11use crate::functions::{Function, custom_error, number_value};
12use crate::interpreter::{SearchResult, interpret};
13use crate::registry::register_if_enabled;
14use crate::value_ext::ValueExt;
15use crate::{Context, Runtime, arg, defn, get_expref_id};
16
17/// Helper to extract an expref AST from a function argument.
18fn get_expref_ast<'a>(value: &Value, ctx: &'a Context<'_>) -> Option<&'a Ast> {
19    get_expref_id(value).and_then(|id| ctx.get_expref(id))
20}
21
22/// Convert a Value to a string key for grouping/deduplication.
23fn value_to_string(value: &Value) -> String {
24    match value {
25        Value::String(s) => s.clone(),
26        Value::Number(n) => n.to_string(),
27        Value::Bool(b) => b.to_string(),
28        Value::Null => "null".to_string(),
29        _ => serde_json::to_string(value).unwrap_or_default(),
30    }
31}
32
33/// Compare two values for sorting purposes.
34fn compare_values(a: &Value, b: &Value) -> std::cmp::Ordering {
35    use std::cmp::Ordering;
36    match (a, b) {
37        (Value::Number(an), Value::Number(bn)) => {
38            let a_f = an.as_f64().unwrap_or(0.0);
39            let b_f = bn.as_f64().unwrap_or(0.0);
40            a_f.partial_cmp(&b_f).unwrap_or(Ordering::Equal)
41        }
42        (Value::String(a_s), Value::String(b_s)) => a_s.cmp(b_s),
43        (Value::Null, Value::Null) => Ordering::Equal,
44        (Value::Null, _) => Ordering::Less,
45        (_, Value::Null) => Ordering::Greater,
46        _ => Ordering::Equal,
47    }
48}
49
50/// Register only the expression functions that are in the enabled set.
51pub fn register_filtered(runtime: &mut Runtime, enabled: &HashSet<&str>) {
52    register_if_enabled(runtime, "map_expr", enabled, Box::new(MapExprFn::new()));
53    register_if_enabled(
54        runtime,
55        "filter_expr",
56        enabled,
57        Box::new(FilterExprFn::new()),
58    );
59    register_if_enabled(runtime, "any_expr", enabled, Box::new(AnyExprFn::new()));
60    register_if_enabled(runtime, "all_expr", enabled, Box::new(AllExprFn::new()));
61    register_if_enabled(runtime, "find_expr", enabled, Box::new(FindExprFn::new()));
62    register_if_enabled(
63        runtime,
64        "find_index_expr",
65        enabled,
66        Box::new(FindIndexExprFn::new()),
67    );
68    register_if_enabled(runtime, "count_expr", enabled, Box::new(CountExprFn::new()));
69    register_if_enabled(
70        runtime,
71        "sort_by_expr",
72        enabled,
73        Box::new(SortByExprFn::new()),
74    );
75    register_if_enabled(
76        runtime,
77        "group_by_expr",
78        enabled,
79        Box::new(GroupByExprFn::new()),
80    );
81    register_if_enabled(
82        runtime,
83        "partition_expr",
84        enabled,
85        Box::new(PartitionExprFn::new()),
86    );
87    register_if_enabled(
88        runtime,
89        "min_by_expr",
90        enabled,
91        Box::new(MinByExprFn::new()),
92    );
93    register_if_enabled(
94        runtime,
95        "max_by_expr",
96        enabled,
97        Box::new(MaxByExprFn::new()),
98    );
99    register_if_enabled(
100        runtime,
101        "unique_by_expr",
102        enabled,
103        Box::new(UniqueByExprFn::new()),
104    );
105    register_if_enabled(
106        runtime,
107        "flat_map_expr",
108        enabled,
109        Box::new(FlatMapExprFn::new()),
110    );
111
112    // Clojure-style alias for flat_map_expr
113    register_if_enabled(runtime, "mapcat", enabled, Box::new(FlatMapExprFn::new()));
114
115    // Lodash-style aliases
116    register_if_enabled(runtime, "some", enabled, Box::new(AnyExprFn::new()));
117    register_if_enabled(runtime, "every", enabled, Box::new(AllExprFn::new()));
118    register_if_enabled(runtime, "reject", enabled, Box::new(RejectFn::new()));
119    register_if_enabled(runtime, "map_keys", enabled, Box::new(MapKeysFn::new()));
120    register_if_enabled(runtime, "map_values", enabled, Box::new(MapValuesFn::new()));
121    register_if_enabled(runtime, "order_by", enabled, Box::new(OrderByFn::new()));
122    register_if_enabled(
123        runtime,
124        "reduce_expr",
125        enabled,
126        Box::new(ReduceExprFn::new()),
127    );
128    register_if_enabled(runtime, "scan_expr", enabled, Box::new(ScanExprFn::new()));
129    // Alias for reduce_expr (lodash-style)
130    register_if_enabled(runtime, "fold", enabled, Box::new(ReduceExprFn::new()));
131    // Clojure-style alias for scan_expr
132    register_if_enabled(runtime, "reductions", enabled, Box::new(ScanExprFn::new()));
133    // none - opposite of any_expr/some
134    register_if_enabled(runtime, "none", enabled, Box::new(NoneFn::new()));
135    register_if_enabled(runtime, "count_by", enabled, Box::new(CountByFn::new()));
136
137    // Partial application functions
138    register_if_enabled(runtime, "partial", enabled, Box::new(PartialFn::new()));
139    register_if_enabled(runtime, "apply", enabled, Box::new(ApplyFn::new()));
140
141    // Functional array operations
142    register_if_enabled(runtime, "take_while", enabled, Box::new(TakeWhileFn::new()));
143    register_if_enabled(runtime, "drop_while", enabled, Box::new(DropWhileFn::new()));
144    register_if_enabled(runtime, "zip_with", enabled, Box::new(ZipWithFn::new()));
145
146    // Recursive transformation
147    register_if_enabled(runtime, "walk", enabled, Box::new(WalkFn::new()));
148
149    // Recursive descent (jq parity)
150    register_if_enabled(runtime, "recurse", enabled, Box::new(RecurseFn::new()));
151    register_if_enabled(
152        runtime,
153        "recurse_with",
154        enabled,
155        Box::new(RecurseWithFn::new()),
156    );
157
158    // Loop functions (jq parity)
159    register_if_enabled(runtime, "while_expr", enabled, Box::new(WhileExprFn::new()));
160    register_if_enabled(runtime, "until_expr", enabled, Box::new(UntilExprFn::new()));
161}
162
163// =============================================================================
164// map_expr(expr, array) -> array
165// =============================================================================
166
167defn!(MapExprFn, vec![arg!(expref), arg!(array)], None);
168
169impl Function for MapExprFn {
170    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
171        self.signature.validate(args, ctx)?;
172
173        let ast = get_expref_ast(&args[0], ctx)
174            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
175            .clone();
176        let arr = args[1].as_array().unwrap();
177
178        let mut results = Vec::with_capacity(arr.len());
179        for item in arr {
180            results.push(interpret(item, &ast, ctx)?);
181        }
182
183        Ok(Value::Array(results))
184    }
185}
186
187// =============================================================================
188// filter_expr(expr, array) -> array
189// =============================================================================
190
191defn!(FilterExprFn, vec![arg!(expref), arg!(array)], None);
192
193impl Function for FilterExprFn {
194    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
195        self.signature.validate(args, ctx)?;
196
197        let ast = get_expref_ast(&args[0], ctx)
198            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
199            .clone();
200        let arr = args[1].as_array().unwrap();
201
202        let mut results = Vec::new();
203        for item in arr {
204            let result = interpret(item, &ast, ctx)?;
205            if result.is_truthy() {
206                results.push(item.clone());
207            }
208        }
209
210        Ok(Value::Array(results))
211    }
212}
213
214// =============================================================================
215// any_expr(expr, array) -> bool
216// =============================================================================
217
218defn!(AnyExprFn, vec![arg!(expref), arg!(array)], None);
219
220impl Function for AnyExprFn {
221    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
222        self.signature.validate(args, ctx)?;
223
224        let ast = get_expref_ast(&args[0], ctx)
225            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
226            .clone();
227        let arr = args[1].as_array().unwrap();
228
229        for item in arr {
230            let result = interpret(item, &ast, ctx)?;
231            if result.is_truthy() {
232                return Ok(Value::Bool(true));
233            }
234        }
235
236        Ok(Value::Bool(false))
237    }
238}
239
240// =============================================================================
241// none(expr, array) -> bool (opposite of any_expr/some)
242// =============================================================================
243
244defn!(NoneFn, vec![arg!(expref), arg!(array)], None);
245
246impl Function for NoneFn {
247    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
248        self.signature.validate(args, ctx)?;
249
250        let ast = get_expref_ast(&args[0], ctx)
251            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
252            .clone();
253        let arr = args[1].as_array().unwrap();
254
255        // Empty array returns true (vacuously, no elements satisfy the predicate)
256        if arr.is_empty() {
257            return Ok(Value::Bool(true));
258        }
259
260        for item in arr {
261            let result = interpret(item, &ast, ctx)?;
262            if result.is_truthy() {
263                return Ok(Value::Bool(false));
264            }
265        }
266
267        Ok(Value::Bool(true))
268    }
269}
270
271// =============================================================================
272// all_expr(expr, array) -> bool
273// =============================================================================
274
275defn!(AllExprFn, vec![arg!(expref), arg!(array)], None);
276
277impl Function for AllExprFn {
278    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
279        self.signature.validate(args, ctx)?;
280
281        let ast = get_expref_ast(&args[0], ctx)
282            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
283            .clone();
284        let arr = args[1].as_array().unwrap();
285
286        // Empty array returns true (vacuous truth)
287        if arr.is_empty() {
288            return Ok(Value::Bool(true));
289        }
290
291        for item in arr {
292            let result = interpret(item, &ast, ctx)?;
293            if !result.is_truthy() {
294                return Ok(Value::Bool(false));
295            }
296        }
297
298        Ok(Value::Bool(true))
299    }
300}
301
302// =============================================================================
303// find_expr(expr, array) -> element | null
304// =============================================================================
305
306defn!(FindExprFn, vec![arg!(expref), arg!(array)], None);
307
308impl Function for FindExprFn {
309    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
310        self.signature.validate(args, ctx)?;
311
312        let ast = get_expref_ast(&args[0], ctx)
313            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
314            .clone();
315        let arr = args[1].as_array().unwrap();
316
317        for item in arr {
318            let result = interpret(item, &ast, ctx)?;
319            if result.is_truthy() {
320                return Ok(item.clone());
321            }
322        }
323
324        Ok(Value::Null)
325    }
326}
327
328// =============================================================================
329// find_index_expr(expr, array) -> number | null
330// =============================================================================
331
332defn!(FindIndexExprFn, vec![arg!(expref), arg!(array)], None);
333
334impl Function for FindIndexExprFn {
335    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
336        self.signature.validate(args, ctx)?;
337
338        let ast = get_expref_ast(&args[0], ctx)
339            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
340            .clone();
341        let arr = args[1].as_array().unwrap();
342
343        for (i, item) in arr.iter().enumerate() {
344            let result = interpret(item, &ast, ctx)?;
345            if result.is_truthy() {
346                return Ok(number_value(i as f64));
347            }
348        }
349
350        Ok(number_value(-1.0))
351    }
352}
353
354// =============================================================================
355// count_expr(expr, array) -> number
356// =============================================================================
357
358defn!(CountExprFn, vec![arg!(expref), arg!(array)], None);
359
360impl Function for CountExprFn {
361    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
362        self.signature.validate(args, ctx)?;
363
364        let ast = get_expref_ast(&args[0], ctx)
365            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
366            .clone();
367        let arr = args[1].as_array().unwrap();
368
369        let mut count = 0i64;
370        for item in arr {
371            let result = interpret(item, &ast, ctx)?;
372            if result.is_truthy() {
373                count += 1;
374            }
375        }
376
377        Ok(Value::Number(Number::from(count)))
378    }
379}
380
381// =============================================================================
382// sort_by_expr(expr, array) -> array
383// =============================================================================
384
385defn!(SortByExprFn, vec![arg!(expref), arg!(array)], None);
386
387impl Function for SortByExprFn {
388    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
389        self.signature.validate(args, ctx)?;
390
391        let ast = get_expref_ast(&args[0], ctx)
392            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
393            .clone();
394        let arr = args[1].as_array().unwrap();
395
396        if arr.is_empty() {
397            return Ok(Value::Array(vec![]));
398        }
399
400        // Compute sort keys for each element
401        let mut keyed: Vec<(Value, Value)> = Vec::with_capacity(arr.len());
402        for item in arr {
403            let key = interpret(item, &ast, ctx)?;
404            keyed.push((item.clone(), key));
405        }
406
407        // Sort by key
408        keyed.sort_by(|a, b| compare_values(&a.1, &b.1));
409
410        let results: Vec<Value> = keyed.into_iter().map(|(item, _)| item).collect();
411        Ok(Value::Array(results))
412    }
413}
414
415// =============================================================================
416// group_by_expr(expr, array) -> object
417// =============================================================================
418
419defn!(GroupByExprFn, vec![arg!(expref), arg!(array)], None);
420
421impl Function for GroupByExprFn {
422    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
423        self.signature.validate(args, ctx)?;
424
425        let ast = get_expref_ast(&args[0], ctx)
426            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
427            .clone();
428        let arr = args[1].as_array().unwrap();
429
430        // Use an index-based approach to preserve insertion order
431        let mut group_keys: Vec<String> = Vec::new();
432        let mut group_map: std::collections::HashMap<String, Vec<Value>> =
433            std::collections::HashMap::new();
434
435        for item in arr {
436            let key_val = interpret(item, &ast, ctx)?;
437            let key = value_to_string(&key_val);
438            if !group_map.contains_key(&key) {
439                group_keys.push(key.clone());
440            }
441            group_map.entry(key).or_default().push(item.clone());
442        }
443
444        let mut result = Map::new();
445        for key in group_keys {
446            if let Some(items) = group_map.remove(&key) {
447                result.insert(key, Value::Array(items));
448            }
449        }
450
451        Ok(Value::Object(result))
452    }
453}
454
455// =============================================================================
456// count_by(expr, array) -> object
457// =============================================================================
458
459defn!(CountByFn, vec![arg!(expref), arg!(array)], None);
460
461impl Function for CountByFn {
462    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
463        self.signature.validate(args, ctx)?;
464
465        let ast = get_expref_ast(&args[0], ctx)
466            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
467            .clone();
468        let arr = args[1].as_array().unwrap();
469
470        let mut count_keys: Vec<String> = Vec::new();
471        let mut counts: std::collections::HashMap<String, i64> = std::collections::HashMap::new();
472
473        for item in arr {
474            let key_val = interpret(item, &ast, ctx)?;
475            let key = value_to_string(&key_val);
476            if !counts.contains_key(&key) {
477                count_keys.push(key.clone());
478            }
479            *counts.entry(key).or_insert(0) += 1;
480        }
481
482        let mut result = Map::new();
483        for key in count_keys {
484            if let Some(&count) = counts.get(&key) {
485                result.insert(key, Value::Number(Number::from(count)));
486            }
487        }
488
489        Ok(Value::Object(result))
490    }
491}
492
493// =============================================================================
494// partition_expr(expr, array) -> [matches, non_matches]
495// =============================================================================
496
497defn!(PartitionExprFn, vec![arg!(expref), arg!(array)], None);
498
499impl Function for PartitionExprFn {
500    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
501        self.signature.validate(args, ctx)?;
502
503        let ast = get_expref_ast(&args[0], ctx)
504            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
505            .clone();
506        let arr = args[1].as_array().unwrap();
507
508        let mut matches = Vec::new();
509        let mut non_matches = Vec::new();
510
511        for item in arr {
512            let result = interpret(item, &ast, ctx)?;
513            if result.is_truthy() {
514                matches.push(item.clone());
515            } else {
516                non_matches.push(item.clone());
517            }
518        }
519
520        Ok(Value::Array(vec![
521            Value::Array(matches),
522            Value::Array(non_matches),
523        ]))
524    }
525}
526
527// =============================================================================
528// min_by_expr(expr, array) -> element | null
529// =============================================================================
530
531defn!(MinByExprFn, vec![arg!(expref), arg!(array)], None);
532
533impl Function for MinByExprFn {
534    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
535        self.signature.validate(args, ctx)?;
536
537        let ast = get_expref_ast(&args[0], ctx)
538            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
539            .clone();
540        let arr = args[1].as_array().unwrap();
541
542        if arr.is_empty() {
543            return Ok(Value::Null);
544        }
545
546        let mut min_item = arr[0].clone();
547        let mut min_key = interpret(&arr[0], &ast, ctx)?;
548
549        for item in arr.iter().skip(1) {
550            let key = interpret(item, &ast, ctx)?;
551            if compare_values(&key, &min_key) == std::cmp::Ordering::Less {
552                min_item = item.clone();
553                min_key = key;
554            }
555        }
556
557        Ok(min_item)
558    }
559}
560
561// =============================================================================
562// max_by_expr(expr, array) -> element | null
563// =============================================================================
564
565defn!(MaxByExprFn, vec![arg!(expref), arg!(array)], None);
566
567impl Function for MaxByExprFn {
568    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
569        self.signature.validate(args, ctx)?;
570
571        let ast = get_expref_ast(&args[0], ctx)
572            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
573            .clone();
574        let arr = args[1].as_array().unwrap();
575
576        if arr.is_empty() {
577            return Ok(Value::Null);
578        }
579
580        let mut max_item = arr[0].clone();
581        let mut max_key = interpret(&arr[0], &ast, ctx)?;
582
583        for item in arr.iter().skip(1) {
584            let key = interpret(item, &ast, ctx)?;
585            if compare_values(&key, &max_key) == std::cmp::Ordering::Greater {
586                max_item = item.clone();
587                max_key = key;
588            }
589        }
590
591        Ok(max_item)
592    }
593}
594
595// =============================================================================
596// unique_by_expr(expr, array) -> array
597// =============================================================================
598
599defn!(UniqueByExprFn, vec![arg!(expref), arg!(array)], None);
600
601impl Function for UniqueByExprFn {
602    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
603        self.signature.validate(args, ctx)?;
604
605        let ast = get_expref_ast(&args[0], ctx)
606            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
607            .clone();
608        let arr = args[1].as_array().unwrap();
609
610        let mut seen: HashSet<String> = HashSet::new();
611        let mut results = Vec::new();
612
613        for item in arr {
614            let key_val = interpret(item, &ast, ctx)?;
615            let key = value_to_string(&key_val);
616            if seen.insert(key) {
617                results.push(item.clone());
618            }
619        }
620
621        Ok(Value::Array(results))
622    }
623}
624
625// =============================================================================
626// flat_map_expr(expr, array) -> array
627// =============================================================================
628
629defn!(FlatMapExprFn, vec![arg!(expref), arg!(array)], None);
630
631impl Function for FlatMapExprFn {
632    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
633        self.signature.validate(args, ctx)?;
634
635        let ast = get_expref_ast(&args[0], ctx)
636            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
637            .clone();
638        let arr = args[1].as_array().unwrap();
639
640        let mut results = Vec::new();
641        for item in arr {
642            let result = interpret(item, &ast, ctx)?;
643            match result {
644                Value::Array(inner) => {
645                    results.extend(inner);
646                }
647                Value::Null => {
648                    // Skip nulls
649                }
650                _ => {
651                    results.push(result);
652                }
653            }
654        }
655
656        Ok(Value::Array(results))
657    }
658}
659
660// =============================================================================
661// reject(expr, array) -> array (inverse of filter_expr)
662// =============================================================================
663
664defn!(RejectFn, vec![arg!(expref), arg!(array)], None);
665
666impl Function for RejectFn {
667    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
668        self.signature.validate(args, ctx)?;
669
670        let ast = get_expref_ast(&args[0], ctx)
671            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
672            .clone();
673        let arr = args[1].as_array().unwrap();
674
675        let mut results = Vec::new();
676        for item in arr {
677            let result = interpret(item, &ast, ctx)?;
678            // Keep items where expression is falsy (inverse of filter)
679            if !result.is_truthy() {
680                results.push(item.clone());
681            }
682        }
683
684        Ok(Value::Array(results))
685    }
686}
687
688// =============================================================================
689// map_keys(expr, object) -> object
690// =============================================================================
691
692defn!(MapKeysFn, vec![arg!(expref), arg!(object)], None);
693
694impl Function for MapKeysFn {
695    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
696        self.signature.validate(args, ctx)?;
697
698        let ast = get_expref_ast(&args[0], ctx)
699            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
700            .clone();
701        let obj = args[1].as_object().unwrap();
702
703        let mut result = Map::new();
704        for (key, value) in obj.iter() {
705            // Apply expression to the key
706            let key_val = Value::String(key.clone());
707            let new_key_val = interpret(&key_val, &ast, ctx)?;
708
709            let new_key_str = match &new_key_val {
710                Value::String(s) => s.clone(),
711                Value::Number(n) => n.to_string(),
712                _ => key.clone(), // Keep original if result is not a string/number
713            };
714
715            result.insert(new_key_str, value.clone());
716        }
717
718        Ok(Value::Object(result))
719    }
720}
721
722// =============================================================================
723// map_values(expr, object) -> object
724// =============================================================================
725
726defn!(MapValuesFn, vec![arg!(expref), arg!(object)], None);
727
728impl Function for MapValuesFn {
729    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
730        self.signature.validate(args, ctx)?;
731
732        let ast = get_expref_ast(&args[0], ctx)
733            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
734            .clone();
735        let obj = args[1].as_object().unwrap();
736
737        let mut result = Map::new();
738        for (key, value) in obj.iter() {
739            let new_value = interpret(value, &ast, ctx)?;
740            result.insert(key.clone(), new_value);
741        }
742
743        Ok(Value::Object(result))
744    }
745}
746
747// =============================================================================
748// order_by(array, criteria) -> array
749// =============================================================================
750
751defn!(OrderByFn, vec![arg!(array), arg!(array)], None);
752
753impl Function for OrderByFn {
754    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
755        self.signature.validate(args, ctx)?;
756
757        let arr = args[0].as_array().unwrap();
758        let criteria = args[1].as_array().unwrap();
759
760        if arr.is_empty() {
761            return Ok(Value::Array(vec![]));
762        }
763
764        // Parse criteria: each element should be [field, direction]
765        let mut sort_specs: Vec<(String, bool)> = Vec::new(); // (field, ascending)
766        for criterion in criteria {
767            let crit_arr = criterion.as_array().ok_or_else(|| {
768                custom_error(ctx, "Each criterion must be an array [field, direction]")
769            })?;
770
771            if crit_arr.len() < 2 {
772                return Err(custom_error(
773                    ctx,
774                    "Each criterion must have [field, direction]",
775                ));
776            }
777
778            let field = crit_arr[0]
779                .as_str()
780                .ok_or_else(|| custom_error(ctx, "Field name must be a string"))?;
781
782            let direction = crit_arr[1]
783                .as_str()
784                .ok_or_else(|| custom_error(ctx, "Direction must be 'asc' or 'desc'"))?;
785
786            let ascending = match direction.to_lowercase().as_str() {
787                "asc" | "ascending" => true,
788                "desc" | "descending" => false,
789                _ => {
790                    return Err(custom_error(ctx, "Direction must be 'asc' or 'desc'"));
791                }
792            };
793
794            sort_specs.push((field.to_string(), ascending));
795        }
796
797        // Clone and sort the array
798        let mut result: Vec<Value> = arr.clone();
799        result.sort_by(|a, b| {
800            for (field, ascending) in &sort_specs {
801                let a_val = a
802                    .as_object()
803                    .and_then(|o| o.get(field.as_str()))
804                    .unwrap_or(&Value::Null);
805                let b_val = b
806                    .as_object()
807                    .and_then(|o| o.get(field.as_str()))
808                    .unwrap_or(&Value::Null);
809
810                let cmp = compare_values(a_val, b_val);
811                if cmp != std::cmp::Ordering::Equal {
812                    return if *ascending { cmp } else { cmp.reverse() };
813                }
814            }
815            std::cmp::Ordering::Equal
816        });
817
818        Ok(Value::Array(result))
819    }
820}
821
822// =============================================================================
823// reduce_expr(expr, array, initial) -> any
824// =============================================================================
825
826defn!(
827    ReduceExprFn,
828    vec![arg!(expref), arg!(array), arg!(any)],
829    None
830);
831
832impl Function for ReduceExprFn {
833    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
834        self.signature.validate(args, ctx)?;
835
836        let ast = get_expref_ast(&args[0], ctx)
837            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
838            .clone();
839        let arr = args[1].as_array().unwrap();
840        let initial = args[2].clone();
841
842        if arr.is_empty() {
843            return Ok(initial);
844        }
845
846        let mut accumulator = initial;
847
848        for (idx, item) in arr.iter().enumerate() {
849            // Create context object with accumulator, current, and index
850            let mut context_map = Map::new();
851            context_map.insert("accumulator".to_string(), accumulator.clone());
852            context_map.insert("current".to_string(), item.clone());
853            context_map.insert("index".to_string(), Value::Number(Number::from(idx as i64)));
854            let context_val = Value::Object(context_map);
855
856            accumulator = interpret(&context_val, &ast, ctx)?;
857        }
858
859        Ok(accumulator)
860    }
861}
862
863// =============================================================================
864// scan_expr(expr, array, initial) -> array
865// =============================================================================
866
867defn!(ScanExprFn, vec![arg!(expref), arg!(array), arg!(any)], None);
868
869impl Function for ScanExprFn {
870    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
871        self.signature.validate(args, ctx)?;
872
873        let ast = get_expref_ast(&args[0], ctx)
874            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
875            .clone();
876        let arr = args[1].as_array().unwrap();
877        let initial = args[2].clone();
878
879        if arr.is_empty() {
880            return Ok(Value::Array(vec![]));
881        }
882
883        let mut accumulator = initial;
884        let mut results: Vec<Value> = Vec::with_capacity(arr.len());
885
886        for (idx, item) in arr.iter().enumerate() {
887            // Create context object with accumulator, current, and index
888            let mut context_map = Map::new();
889            context_map.insert("accumulator".to_string(), accumulator.clone());
890            context_map.insert("current".to_string(), item.clone());
891            context_map.insert("index".to_string(), Value::Number(Number::from(idx as i64)));
892            let context_val = Value::Object(context_map);
893
894            accumulator = interpret(&context_val, &ast, ctx)?;
895            results.push(accumulator.clone());
896        }
897
898        Ok(Value::Array(results))
899    }
900}
901
902// =============================================================================
903// partial(fn_name, ...args) -> partial object
904// =============================================================================
905
906defn!(PartialFn, vec![arg!(string)], Some(arg!(any)));
907
908impl Function for PartialFn {
909    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
910        self.signature.validate(args, ctx)?;
911
912        let fn_name = args[0].as_str().ok_or_else(|| {
913            custom_error(
914                ctx,
915                "partial() first argument must be a function name string",
916            )
917        })?;
918
919        // Collect the pre-filled arguments
920        let prefilled_args: Vec<Value> = args[1..].to_vec();
921
922        // Create the partial object
923        let mut partial_obj = Map::new();
924        partial_obj.insert("__partial__".to_string(), Value::Bool(true));
925        partial_obj.insert("fn".to_string(), Value::String(fn_name.to_string()));
926        partial_obj.insert("args".to_string(), Value::Array(prefilled_args));
927
928        Ok(Value::Object(partial_obj))
929    }
930}
931
932// =============================================================================
933// apply(partial_or_fn, ...args) -> result
934// =============================================================================
935
936defn!(ApplyFn, vec![arg!(any)], Some(arg!(any)));
937
938impl Function for ApplyFn {
939    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
940        self.signature.validate(args, ctx)?;
941
942        let first_arg = &args[0];
943        let additional_args = &args[1..];
944
945        // Check if it's a partial object
946        if let Some(obj) = first_arg.as_object()
947            && obj.get("__partial__").and_then(|v| v.as_bool()) == Some(true)
948        {
949            // It's a partial - extract fn name and pre-filled args
950            let fn_name = obj
951                .get("fn")
952                .and_then(|v| v.as_str())
953                .ok_or_else(|| custom_error(ctx, "Invalid partial object: missing 'fn' field"))?;
954
955            let prefilled = obj
956                .get("args")
957                .and_then(|v| v.as_array())
958                .ok_or_else(|| custom_error(ctx, "Invalid partial object: missing 'args' field"))?;
959
960            return invoke_function(fn_name, prefilled, additional_args, ctx);
961        }
962
963        // If it's a string, treat as function name
964        if let Some(fn_name) = first_arg.as_str() {
965            return invoke_function(fn_name, &[], additional_args, ctx);
966        }
967
968        Err(custom_error(
969            ctx,
970            "apply() first argument must be a partial object or function name string",
971        ))
972    }
973}
974
975/// Helper to invoke a function by name with pre-filled and additional arguments.
976fn invoke_function(
977    fn_name: &str,
978    prefilled: &[Value],
979    additional: &[Value],
980    ctx: &mut Context<'_>,
981) -> SearchResult {
982    // Build the argument list for the expression
983    let mut all_args_json: Vec<String> = Vec::new();
984
985    // Add pre-filled args as literals
986    for a in prefilled {
987        all_args_json.push(format!("`{}`", serde_json::to_string(a).unwrap()));
988    }
989
990    // Add additional args as literals
991    for a in additional {
992        all_args_json.push(format!("`{}`", serde_json::to_string(a).unwrap()));
993    }
994
995    // Build and execute the expression
996    let expr_str = format!("{}({})", fn_name, all_args_json.join(", "));
997
998    let compiled = ctx.runtime.compile(&expr_str).map_err(|_| {
999        custom_error(
1000            ctx,
1001            &format!("Failed to compile function call '{}'", expr_str),
1002        )
1003    })?;
1004
1005    compiled
1006        .search(&Value::Null)
1007        .map_err(|_| custom_error(ctx, &format!("Failed to execute '{}'", fn_name)))
1008}
1009
1010// =============================================================================
1011// take_while(expr, array) -> array
1012// =============================================================================
1013
1014defn!(TakeWhileFn, vec![arg!(expref), arg!(array)], None);
1015
1016impl Function for TakeWhileFn {
1017    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
1018        self.signature.validate(args, ctx)?;
1019
1020        let ast = get_expref_ast(&args[0], ctx)
1021            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
1022            .clone();
1023        let arr = args[1].as_array().unwrap();
1024
1025        let mut results = Vec::new();
1026        for item in arr {
1027            let result = interpret(item, &ast, ctx)?;
1028            if result.is_truthy() {
1029                results.push(item.clone());
1030            } else {
1031                break;
1032            }
1033        }
1034
1035        Ok(Value::Array(results))
1036    }
1037}
1038
1039// =============================================================================
1040// drop_while(expr, array) -> array
1041// =============================================================================
1042
1043defn!(DropWhileFn, vec![arg!(expref), arg!(array)], None);
1044
1045impl Function for DropWhileFn {
1046    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
1047        self.signature.validate(args, ctx)?;
1048
1049        let ast = get_expref_ast(&args[0], ctx)
1050            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
1051            .clone();
1052        let arr = args[1].as_array().unwrap();
1053
1054        let mut dropping = true;
1055        let mut results = Vec::new();
1056        for item in arr {
1057            if dropping {
1058                let result = interpret(item, &ast, ctx)?;
1059                if !result.is_truthy() {
1060                    dropping = false;
1061                    results.push(item.clone());
1062                }
1063            } else {
1064                results.push(item.clone());
1065            }
1066        }
1067
1068        Ok(Value::Array(results))
1069    }
1070}
1071
1072// =============================================================================
1073// zip_with(expr, array1, array2) -> array
1074// =============================================================================
1075
1076defn!(
1077    ZipWithFn,
1078    vec![arg!(expref), arg!(array), arg!(array)],
1079    None
1080);
1081
1082impl Function for ZipWithFn {
1083    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
1084        self.signature.validate(args, ctx)?;
1085
1086        let ast = get_expref_ast(&args[0], ctx)
1087            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
1088            .clone();
1089        let arr1 = args[1].as_array().unwrap();
1090        let arr2 = args[2].as_array().unwrap();
1091
1092        let min_len = arr1.len().min(arr2.len());
1093        let mut results = Vec::with_capacity(min_len);
1094
1095        for i in 0..min_len {
1096            // Create a pair array [element1, element2] as input to the expression
1097            let pair = Value::Array(vec![arr1[i].clone(), arr2[i].clone()]);
1098            let result = interpret(&pair, &ast, ctx)?;
1099            results.push(result);
1100        }
1101
1102        Ok(Value::Array(results))
1103    }
1104}
1105
1106// =============================================================================
1107// walk(expr, value) -> value (recursive transformation)
1108// =============================================================================
1109
1110defn!(WalkFn, vec![arg!(expref), arg!(any)], None);
1111
1112/// Recursively walk a value, applying the expression bottom-up.
1113fn walk_value(value: &Value, ast: &Ast, ctx: &mut Context<'_>) -> SearchResult {
1114    match value {
1115        Value::Array(arr) => {
1116            // First, recursively walk all elements
1117            let walked_elements: Result<Vec<Value>, _> =
1118                arr.iter().map(|elem| walk_value(elem, ast, ctx)).collect();
1119            let new_array = Value::Array(walked_elements?);
1120            // Then apply the expression to the array itself
1121            interpret(&new_array, ast, ctx)
1122        }
1123        Value::Object(obj) => {
1124            // First, recursively walk all values
1125            let mut walked_obj = Map::new();
1126            for (k, v) in obj.iter() {
1127                walked_obj.insert(k.clone(), walk_value(v, ast, ctx)?);
1128            }
1129            let new_object = Value::Object(walked_obj);
1130            // Then apply the expression to the object itself
1131            interpret(&new_object, ast, ctx)
1132        }
1133        // For scalars (string, number, bool, null), just apply the expression
1134        _ => interpret(value, ast, ctx),
1135    }
1136}
1137
1138impl Function for WalkFn {
1139    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
1140        self.signature.validate(args, ctx)?;
1141
1142        let ast = get_expref_ast(&args[0], ctx)
1143            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
1144            .clone();
1145
1146        walk_value(&args[1], &ast, ctx)
1147    }
1148}
1149
1150// =============================================================================
1151// recurse(value) -> array (collect all nested values, jq parity)
1152// =============================================================================
1153
1154defn!(RecurseFn, vec![arg!(any)], None);
1155
1156/// Helper to collect all values recursively.
1157fn collect_recursive(value: &Value, results: &mut Vec<Value>) {
1158    results.push(value.clone());
1159    match value {
1160        Value::Array(arr) => {
1161            for elem in arr {
1162                collect_recursive(elem, results);
1163            }
1164        }
1165        Value::Object(obj) => {
1166            for (_, v) in obj.iter() {
1167                collect_recursive(v, results);
1168            }
1169        }
1170        _ => {}
1171    }
1172}
1173
1174impl Function for RecurseFn {
1175    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
1176        self.signature.validate(args, ctx)?;
1177
1178        let mut results = Vec::new();
1179        collect_recursive(&args[0], &mut results);
1180
1181        Ok(Value::Array(results))
1182    }
1183}
1184
1185// =============================================================================
1186// recurse_with(value, expr) -> array (recursive descent with filter, jq parity)
1187// =============================================================================
1188
1189defn!(RecurseWithFn, vec![arg!(any), arg!(expref)], None);
1190
1191impl Function for RecurseWithFn {
1192    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
1193        self.signature.validate(args, ctx)?;
1194
1195        let ast = get_expref_ast(&args[1], ctx)
1196            .ok_or_else(|| custom_error(ctx, "Expected expref"))?
1197            .clone();
1198
1199        let mut results = Vec::new();
1200        let mut queue = vec![args[0].clone()];
1201        let max_iterations = 10000; // Safety limit
1202        let mut iterations = 0;
1203
1204        while let Some(current) = queue.pop() {
1205            if iterations >= max_iterations {
1206                return Err(custom_error(
1207                    ctx,
1208                    "recurse_with exceeded maximum iterations",
1209                ));
1210            }
1211            iterations += 1;
1212
1213            // Skip null values
1214            if current.is_null() {
1215                continue;
1216            }
1217
1218            results.push(current.clone());
1219
1220            // Apply expression to get next values
1221            let next = interpret(&current, &ast, ctx)?;
1222
1223            match next {
1224                Value::Null => {}
1225                Value::Array(arr) => {
1226                    // Add non-null elements to queue (in reverse to maintain order)
1227                    for elem in arr.into_iter().rev() {
1228                        if !elem.is_null() {
1229                            queue.push(elem);
1230                        }
1231                    }
1232                }
1233                _ => {
1234                    queue.push(next);
1235                }
1236            }
1237        }
1238
1239        Ok(Value::Array(results))
1240    }
1241}
1242
1243// =============================================================================
1244// while_expr(init, condition, update) -> value
1245// =============================================================================
1246
1247defn!(
1248    WhileExprFn,
1249    vec![arg!(any), arg!(expref), arg!(expref)],
1250    None
1251);
1252
1253impl Function for WhileExprFn {
1254    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
1255        self.signature.validate(args, ctx)?;
1256
1257        let cond_ast = get_expref_ast(&args[1], ctx)
1258            .ok_or_else(|| custom_error(ctx, "Expected expref for condition"))?
1259            .clone();
1260        let update_ast = get_expref_ast(&args[2], ctx)
1261            .ok_or_else(|| custom_error(ctx, "Expected expref for update"))?
1262            .clone();
1263
1264        let mut current = args[0].clone();
1265        let max_iterations = 100000; // Safety limit
1266        let mut iterations = 0;
1267
1268        loop {
1269            if iterations >= max_iterations {
1270                return Err(custom_error(ctx, "while_expr exceeded maximum iterations"));
1271            }
1272            iterations += 1;
1273
1274            // Check condition
1275            let cond_result = interpret(&current, &cond_ast, ctx)?;
1276            if !cond_result.is_truthy() {
1277                break;
1278            }
1279
1280            // Apply update
1281            current = interpret(&current, &update_ast, ctx)?;
1282        }
1283
1284        Ok(current)
1285    }
1286}
1287
1288// =============================================================================
1289// until_expr(init, condition, update) -> value
1290// =============================================================================
1291
1292defn!(
1293    UntilExprFn,
1294    vec![arg!(any), arg!(expref), arg!(expref)],
1295    None
1296);
1297
1298impl Function for UntilExprFn {
1299    fn evaluate(&self, args: &[Value], ctx: &mut Context<'_>) -> SearchResult {
1300        self.signature.validate(args, ctx)?;
1301
1302        let cond_ast = get_expref_ast(&args[1], ctx)
1303            .ok_or_else(|| custom_error(ctx, "Expected expref for condition"))?
1304            .clone();
1305        let update_ast = get_expref_ast(&args[2], ctx)
1306            .ok_or_else(|| custom_error(ctx, "Expected expref for update"))?
1307            .clone();
1308
1309        let mut current = args[0].clone();
1310        let max_iterations = 100000; // Safety limit
1311        let mut iterations = 0;
1312
1313        loop {
1314            if iterations >= max_iterations {
1315                return Err(custom_error(ctx, "until_expr exceeded maximum iterations"));
1316            }
1317            iterations += 1;
1318
1319            // Check condition (stop when true, opposite of while)
1320            let cond_result = interpret(&current, &cond_ast, ctx)?;
1321            if cond_result.is_truthy() {
1322                break;
1323            }
1324
1325            // Apply update
1326            current = interpret(&current, &update_ast, ctx)?;
1327        }
1328
1329        Ok(current)
1330    }
1331}
1332
1333#[cfg(test)]
1334mod tests {
1335    use crate::Runtime;
1336    use serde_json::json;
1337
1338    fn setup_runtime() -> Runtime {
1339        Runtime::builder()
1340            .with_standard()
1341            .with_all_extensions()
1342            .build()
1343    }
1344
1345    #[test]
1346    fn test_map_expr_field() {
1347        let runtime = setup_runtime();
1348        let data = json!([{"name": "Alice"}, {"name": "Bob"}]);
1349        let expr = runtime.compile("map_expr(&name, @)").unwrap();
1350        let result = expr.search(&data).unwrap();
1351        let arr = result.as_array().unwrap();
1352        assert_eq!(arr.len(), 2);
1353        assert_eq!(arr[0].as_str().unwrap(), "Alice");
1354        assert_eq!(arr[1].as_str().unwrap(), "Bob");
1355    }
1356
1357    #[test]
1358    fn test_map_expr_transform() {
1359        let runtime = setup_runtime();
1360        let data = json!(["hello", "world"]);
1361        let expr = runtime.compile("map_expr(&length(@), @)").unwrap();
1362        let result = expr.search(&data).unwrap();
1363        let arr = result.as_array().unwrap();
1364        assert_eq!(arr[0].as_f64().unwrap(), 5.0);
1365        assert_eq!(arr[1].as_f64().unwrap(), 5.0);
1366    }
1367
1368    #[test]
1369    fn test_filter_expr() {
1370        let runtime = setup_runtime();
1371        let data = json!([{"age": 25}, {"age": 17}, {"age": 30}]);
1372        let expr = runtime.compile("filter_expr(&(age >= `18`), @)").unwrap();
1373        let result = expr.search(&data).unwrap();
1374        let arr = result.as_array().unwrap();
1375        assert_eq!(arr.len(), 2);
1376    }
1377
1378    #[test]
1379    fn test_filter_expr_empty() {
1380        let runtime = setup_runtime();
1381        let data = json!([1, 2, 3]);
1382        let expr = runtime.compile("filter_expr(&(@ > `10`), @)").unwrap();
1383        let result = expr.search(&data).unwrap();
1384        let arr = result.as_array().unwrap();
1385        assert_eq!(arr.len(), 0);
1386    }
1387
1388    #[test]
1389    fn test_any_expr_true() {
1390        let runtime = setup_runtime();
1391        let data = json!([{"active": false}, {"active": true}]);
1392        let expr = runtime.compile("any_expr(&active, @)").unwrap();
1393        let result = expr.search(&data).unwrap();
1394        assert!(result.as_bool().unwrap());
1395    }
1396
1397    #[test]
1398    fn test_any_expr_false() {
1399        let runtime = setup_runtime();
1400        let data = json!([{"active": false}, {"active": false}]);
1401        let expr = runtime.compile("any_expr(&active, @)").unwrap();
1402        let result = expr.search(&data).unwrap();
1403        assert!(!result.as_bool().unwrap());
1404    }
1405
1406    #[test]
1407    fn test_all_expr_true() {
1408        let runtime = setup_runtime();
1409        let data = json!([{"active": true}, {"active": true}]);
1410        let expr = runtime.compile("all_expr(&active, @)").unwrap();
1411        let result = expr.search(&data).unwrap();
1412        assert!(result.as_bool().unwrap());
1413    }
1414
1415    #[test]
1416    fn test_all_expr_false() {
1417        let runtime = setup_runtime();
1418        let data = json!([{"active": true}, {"active": false}]);
1419        let expr = runtime.compile("all_expr(&active, @)").unwrap();
1420        let result = expr.search(&data).unwrap();
1421        assert!(!result.as_bool().unwrap());
1422    }
1423
1424    #[test]
1425    fn test_all_expr_empty() {
1426        let runtime = setup_runtime();
1427        let data = json!([]);
1428        let expr = runtime.compile("all_expr(&active, @)").unwrap();
1429        let result = expr.search(&data).unwrap();
1430        assert!(result.as_bool().unwrap()); // vacuous truth
1431    }
1432
1433    #[test]
1434    fn test_find_expr_found() {
1435        let runtime = setup_runtime();
1436        let data = json!([{"id": 1, "name": "Alice"}, {"id": 2, "name": "Bob"}]);
1437        let expr = runtime.compile("find_expr(&(id == `2`), @)").unwrap();
1438        let result = expr.search(&data).unwrap();
1439        let obj = result.as_object().unwrap();
1440        assert_eq!(obj.get("name").unwrap().as_str().unwrap(), "Bob");
1441    }
1442
1443    #[test]
1444    fn test_find_expr_not_found() {
1445        let runtime = setup_runtime();
1446        let data = json!([{"id": 1}, {"id": 2}]);
1447        let expr = runtime.compile("find_expr(&(id == `99`), @)").unwrap();
1448        let result = expr.search(&data).unwrap();
1449        assert!(result.is_null());
1450    }
1451
1452    #[test]
1453    fn test_sort_by_expr_numbers() {
1454        let runtime = setup_runtime();
1455        let data = json!([{"val": 3}, {"val": 1}, {"val": 2}]);
1456        let expr = runtime.compile("sort_by_expr(&val, @)").unwrap();
1457        let result = expr.search(&data).unwrap();
1458        let arr = result.as_array().unwrap();
1459        assert_eq!(
1460            arr[0]
1461                .as_object()
1462                .unwrap()
1463                .get("val")
1464                .unwrap()
1465                .as_f64()
1466                .unwrap(),
1467            1.0
1468        );
1469        assert_eq!(
1470            arr[1]
1471                .as_object()
1472                .unwrap()
1473                .get("val")
1474                .unwrap()
1475                .as_f64()
1476                .unwrap(),
1477            2.0
1478        );
1479        assert_eq!(
1480            arr[2]
1481                .as_object()
1482                .unwrap()
1483                .get("val")
1484                .unwrap()
1485                .as_f64()
1486                .unwrap(),
1487            3.0
1488        );
1489    }
1490
1491    #[test]
1492    fn test_sort_by_expr_strings() {
1493        let runtime = setup_runtime();
1494        let data = json!([{"name": "Charlie"}, {"name": "Alice"}, {"name": "Bob"}]);
1495        let expr = runtime.compile("sort_by_expr(&name, @)").unwrap();
1496        let result = expr.search(&data).unwrap();
1497        let arr = result.as_array().unwrap();
1498        assert_eq!(
1499            arr[0]
1500                .as_object()
1501                .unwrap()
1502                .get("name")
1503                .unwrap()
1504                .as_str()
1505                .unwrap(),
1506            "Alice"
1507        );
1508        assert_eq!(
1509            arr[1]
1510                .as_object()
1511                .unwrap()
1512                .get("name")
1513                .unwrap()
1514                .as_str()
1515                .unwrap(),
1516            "Bob"
1517        );
1518        assert_eq!(
1519            arr[2]
1520                .as_object()
1521                .unwrap()
1522                .get("name")
1523                .unwrap()
1524                .as_str()
1525                .unwrap(),
1526            "Charlie"
1527        );
1528    }
1529
1530    #[test]
1531    fn test_find_index_expr_found() {
1532        let runtime = setup_runtime();
1533        let data = json!([{"id": 1}, {"id": 2}, {"id": 3}]);
1534        let expr = runtime.compile("find_index_expr(&(id == `2`), @)").unwrap();
1535        let result = expr.search(&data).unwrap();
1536        assert_eq!(result.as_f64().unwrap(), 1.0);
1537    }
1538
1539    #[test]
1540    fn test_find_index_expr_not_found() {
1541        let runtime = setup_runtime();
1542        let data = json!([{"id": 1}, {"id": 2}]);
1543        let expr = runtime
1544            .compile("find_index_expr(&(id == `99`), @)")
1545            .unwrap();
1546        let result = expr.search(&data).unwrap();
1547        assert_eq!(result.as_f64().unwrap(), -1.0);
1548    }
1549
1550    #[test]
1551    fn test_count_expr() {
1552        let runtime = setup_runtime();
1553        let data = json!([{"active": true}, {"active": false}, {"active": true}]);
1554        let expr = runtime.compile("count_expr(&active, @)").unwrap();
1555        let result = expr.search(&data).unwrap();
1556        assert_eq!(result.as_f64().unwrap(), 2.0);
1557    }
1558
1559    #[test]
1560    fn test_count_expr_none() {
1561        let runtime = setup_runtime();
1562        let data = json!([1, 2, 3]);
1563        let expr = runtime.compile("count_expr(&(@ > `10`), @)").unwrap();
1564        let result = expr.search(&data).unwrap();
1565        assert_eq!(result.as_f64().unwrap(), 0.0);
1566    }
1567
1568    #[test]
1569    fn test_group_by_expr() {
1570        let runtime = setup_runtime();
1571        let data = json!([
1572            {"type": "a", "val": 1},
1573            {"type": "b", "val": 2},
1574            {"type": "a", "val": 3}
1575        ]);
1576        let expr = runtime.compile("group_by_expr(&type, @)").unwrap();
1577        let result = expr.search(&data).unwrap();
1578        let obj = result.as_object().unwrap();
1579        assert_eq!(obj.get("a").unwrap().as_array().unwrap().len(), 2);
1580        assert_eq!(obj.get("b").unwrap().as_array().unwrap().len(), 1);
1581    }
1582
1583    #[test]
1584    fn test_partition_expr() {
1585        let runtime = setup_runtime();
1586        let data = json!([1, 2, 3, 4, 5]);
1587        let expr = runtime.compile("partition_expr(&(@ > `3`), @)").unwrap();
1588        let result = expr.search(&data).unwrap();
1589        let arr = result.as_array().unwrap();
1590        let matches = arr[0].as_array().unwrap();
1591        let non_matches = arr[1].as_array().unwrap();
1592        assert_eq!(matches.len(), 2); // 4, 5
1593        assert_eq!(non_matches.len(), 3); // 1, 2, 3
1594    }
1595
1596    #[test]
1597    fn test_min_by_expr() {
1598        let runtime = setup_runtime();
1599        let data = json!([{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}]);
1600        let expr = runtime.compile("min_by_expr(&age, @)").unwrap();
1601        let result = expr.search(&data).unwrap();
1602        let obj = result.as_object().unwrap();
1603        assert_eq!(obj.get("name").unwrap().as_str().unwrap(), "Bob");
1604    }
1605
1606    #[test]
1607    fn test_min_by_expr_empty() {
1608        let runtime = setup_runtime();
1609        let data = json!([]);
1610        let expr = runtime.compile("min_by_expr(&age, @)").unwrap();
1611        let result = expr.search(&data).unwrap();
1612        assert!(result.is_null());
1613    }
1614
1615    #[test]
1616    fn test_max_by_expr() {
1617        let runtime = setup_runtime();
1618        let data = json!([{"name": "Alice", "age": 30}, {"name": "Bob", "age": 25}]);
1619        let expr = runtime.compile("max_by_expr(&age, @)").unwrap();
1620        let result = expr.search(&data).unwrap();
1621        let obj = result.as_object().unwrap();
1622        assert_eq!(obj.get("name").unwrap().as_str().unwrap(), "Alice");
1623    }
1624
1625    #[test]
1626    fn test_unique_by_expr() {
1627        let runtime = setup_runtime();
1628        let data = json!([
1629            {"type": "a", "val": 1},
1630            {"type": "b", "val": 2},
1631            {"type": "a", "val": 3}
1632        ]);
1633        let expr = runtime.compile("unique_by_expr(&type, @)").unwrap();
1634        let result = expr.search(&data).unwrap();
1635        let arr = result.as_array().unwrap();
1636        assert_eq!(arr.len(), 2); // First "a" and first "b"
1637        assert_eq!(
1638            arr[0]
1639                .as_object()
1640                .unwrap()
1641                .get("val")
1642                .unwrap()
1643                .as_f64()
1644                .unwrap(),
1645            1.0
1646        );
1647    }
1648
1649    #[test]
1650    fn test_flat_map_expr() {
1651        let runtime = setup_runtime();
1652        let data = json!([
1653            {"tags": ["a", "b"]},
1654            {"tags": ["c"]},
1655            {"tags": ["d", "e"]}
1656        ]);
1657        let expr = runtime.compile("flat_map_expr(&tags, @)").unwrap();
1658        let result = expr.search(&data).unwrap();
1659        let arr = result.as_array().unwrap();
1660        assert_eq!(arr.len(), 5);
1661        assert_eq!(arr[0].as_str().unwrap(), "a");
1662        assert_eq!(arr[4].as_str().unwrap(), "e");
1663    }
1664
1665    #[test]
1666    fn test_flat_map_expr_non_array() {
1667        let runtime = setup_runtime();
1668        let data = json!([{"name": "Alice"}, {"name": "Bob"}]);
1669        let expr = runtime.compile("flat_map_expr(&name, @)").unwrap();
1670        let result = expr.search(&data).unwrap();
1671        let arr = result.as_array().unwrap();
1672        assert_eq!(arr.len(), 2);
1673        assert_eq!(arr[0].as_str().unwrap(), "Alice");
1674    }
1675
1676    #[test]
1677    fn test_some_alias() {
1678        let runtime = setup_runtime();
1679        let data = json!([1, 2, 3, 4, 5]);
1680        let expr = runtime.compile("some(&(@ > `3`), @)").unwrap();
1681        let result = expr.search(&data).unwrap();
1682        assert!(result.as_bool().unwrap());
1683    }
1684
1685    #[test]
1686    fn test_every_alias() {
1687        let runtime = setup_runtime();
1688        let data = json!([2, 4, 6]);
1689        let expr = runtime.compile("every(&(@ > `0`), @)").unwrap();
1690        let result = expr.search(&data).unwrap();
1691        assert!(result.as_bool().unwrap());
1692    }
1693
1694    #[test]
1695    fn test_none_true() {
1696        let runtime = setup_runtime();
1697        let data = json!([1, 2, 3]);
1698        let expr = runtime.compile("none(&(@ > `5`), @)").unwrap();
1699        let result = expr.search(&data).unwrap();
1700        assert!(result.as_bool().unwrap()); // No element > 5
1701    }
1702
1703    #[test]
1704    fn test_none_false() {
1705        let runtime = setup_runtime();
1706        let data = json!([1, 2, 10]);
1707        let expr = runtime.compile("none(&(@ > `5`), @)").unwrap();
1708        let result = expr.search(&data).unwrap();
1709        assert!(!result.as_bool().unwrap()); // 10 > 5
1710    }
1711
1712    #[test]
1713    fn test_none_empty() {
1714        let runtime = setup_runtime();
1715        let data = json!([]);
1716        let expr = runtime.compile("none(&(@ > `0`), @)").unwrap();
1717        let result = expr.search(&data).unwrap();
1718        assert!(result.as_bool().unwrap()); // Vacuously true
1719    }
1720
1721    #[test]
1722    fn test_none_objects() {
1723        let runtime = setup_runtime();
1724        let data = json!([{"active": false}, {"active": false}]);
1725        let expr = runtime.compile("none(&active, @)").unwrap();
1726        let result = expr.search(&data).unwrap();
1727        assert!(result.as_bool().unwrap()); // No active elements
1728    }
1729
1730    #[test]
1731    fn test_mapcat_alias() {
1732        let runtime = setup_runtime();
1733        let data = json!([
1734            {"tags": ["a", "b"]},
1735            {"tags": ["c"]},
1736            {"tags": ["d", "e"]}
1737        ]);
1738        let expr = runtime.compile("mapcat(&tags, @)").unwrap();
1739        let result = expr.search(&data).unwrap();
1740        let arr = result.as_array().unwrap();
1741        assert_eq!(arr.len(), 5);
1742        assert_eq!(arr[0].as_str().unwrap(), "a");
1743        assert_eq!(arr[4].as_str().unwrap(), "e");
1744    }
1745
1746    #[test]
1747    fn test_reductions_alias() {
1748        let runtime = setup_runtime();
1749        let data = json!([1, 2, 3, 4]);
1750        let expr = runtime
1751            .compile("reductions(&sum([accumulator, current]), @, `0`)")
1752            .unwrap();
1753        let result = expr.search(&data).unwrap();
1754        let arr = result.as_array().unwrap();
1755        // Running sum: [1, 3, 6, 10]
1756        assert_eq!(arr.len(), 4);
1757        assert_eq!(arr[0].as_f64().unwrap(), 1.0);
1758        assert_eq!(arr[1].as_f64().unwrap(), 3.0);
1759        assert_eq!(arr[2].as_f64().unwrap(), 6.0);
1760        assert_eq!(arr[3].as_f64().unwrap(), 10.0);
1761    }
1762
1763    #[test]
1764    fn test_reject() {
1765        let runtime = setup_runtime();
1766        let data = json!([1, 2, 3, 4, 5]);
1767        let expr = runtime.compile("reject(&(@ > `2`), @)").unwrap();
1768        let result = expr.search(&data).unwrap();
1769        let arr = result.as_array().unwrap();
1770        assert_eq!(arr.len(), 2); // 1, 2
1771        assert_eq!(arr[0].as_f64().unwrap(), 1.0);
1772        assert_eq!(arr[1].as_f64().unwrap(), 2.0);
1773    }
1774
1775    #[test]
1776    fn test_reject_objects() {
1777        let runtime = setup_runtime();
1778        let data = json!([{"active": true}, {"active": false}, {"active": true}]);
1779        let expr = runtime.compile("reject(&active, @)").unwrap();
1780        let result = expr.search(&data).unwrap();
1781        let arr = result.as_array().unwrap();
1782        assert_eq!(arr.len(), 1); // Only the inactive one
1783    }
1784
1785    #[test]
1786    fn test_map_keys() {
1787        let runtime = setup_runtime();
1788        // Use length to transform key to its length (as string)
1789        let data = json!({"abc": 1, "de": 2});
1790        let expr = runtime.compile("map_keys(&length(@), @)").unwrap();
1791        let result = expr.search(&data).unwrap();
1792        let obj = result.as_object().unwrap();
1793        // "abc" -> 3, "de" -> 2 (converted to string keys)
1794        assert!(obj.contains_key("3") || obj.contains_key("2"));
1795    }
1796
1797    #[test]
1798    fn test_map_values_add() {
1799        let runtime = setup_runtime();
1800        // Use sum to get sum of a literal array
1801        let data = json!({"a": 1, "b": 2, "c": 3});
1802        let expr = runtime.compile("map_values(&sum(`[1]`), @)").unwrap();
1803        let result = expr.search(&data).unwrap();
1804        let obj = result.as_object().unwrap();
1805        // Each value becomes 1 (sum of [1])
1806        assert_eq!(obj.get("a").unwrap().as_f64().unwrap(), 1.0);
1807    }
1808
1809    #[test]
1810    fn test_map_values_length() {
1811        let runtime = setup_runtime();
1812        let data = json!({"name": "alice", "city": "boston"});
1813        let expr = runtime.compile("map_values(&length(@), @)").unwrap();
1814        let result = expr.search(&data).unwrap();
1815        let obj = result.as_object().unwrap();
1816        assert_eq!(obj.get("name").unwrap().as_f64().unwrap(), 5.0); // "alice" = 5 chars
1817        assert_eq!(obj.get("city").unwrap().as_f64().unwrap(), 6.0); // "boston" = 6 chars
1818    }
1819
1820    #[test]
1821    fn test_map_values_with_string_fns() {
1822        let runtime = setup_runtime();
1823        let data = json!({"name": "alice", "city": "boston"});
1824        let expr = runtime.compile("map_values(&upper(@), @)").unwrap();
1825        let result = expr.search(&data).unwrap();
1826        let obj = result.as_object().unwrap();
1827        assert_eq!(obj.get("name").unwrap().as_str().unwrap(), "ALICE");
1828        assert_eq!(obj.get("city").unwrap().as_str().unwrap(), "BOSTON");
1829    }
1830
1831    #[test]
1832    fn test_map_keys_with_string_fns() {
1833        let runtime = setup_runtime();
1834        let data = json!({"hello": 1, "world": 2});
1835        let expr = runtime.compile("map_keys(&upper(@), @)").unwrap();
1836        let result = expr.search(&data).unwrap();
1837        let obj = result.as_object().unwrap();
1838        assert!(obj.contains_key("HELLO"));
1839        assert!(obj.contains_key("WORLD"));
1840    }
1841
1842    #[test]
1843    fn test_order_by_single_field_asc() {
1844        let runtime = setup_runtime();
1845        let data = json!([
1846            {"name": "Charlie", "age": 30},
1847            {"name": "Alice", "age": 25},
1848            {"name": "Bob", "age": 35}
1849        ]);
1850        let expr = runtime
1851            .compile(r#"order_by(@, `[["name", "asc"]]`)"#)
1852            .unwrap();
1853        let result = expr.search(&data).unwrap();
1854        let arr = result.as_array().unwrap();
1855        assert_eq!(
1856            arr[0]
1857                .as_object()
1858                .unwrap()
1859                .get("name")
1860                .unwrap()
1861                .as_str()
1862                .unwrap(),
1863            "Alice"
1864        );
1865        assert_eq!(
1866            arr[1]
1867                .as_object()
1868                .unwrap()
1869                .get("name")
1870                .unwrap()
1871                .as_str()
1872                .unwrap(),
1873            "Bob"
1874        );
1875        assert_eq!(
1876            arr[2]
1877                .as_object()
1878                .unwrap()
1879                .get("name")
1880                .unwrap()
1881                .as_str()
1882                .unwrap(),
1883            "Charlie"
1884        );
1885    }
1886
1887    #[test]
1888    fn test_order_by_single_field_desc() {
1889        let runtime = setup_runtime();
1890        let data = json!([
1891            {"name": "Alice", "age": 25},
1892            {"name": "Bob", "age": 35},
1893            {"name": "Charlie", "age": 30}
1894        ]);
1895        let expr = runtime
1896            .compile(r#"order_by(@, `[["age", "desc"]]`)"#)
1897            .unwrap();
1898        let result = expr.search(&data).unwrap();
1899        let arr = result.as_array().unwrap();
1900        assert_eq!(
1901            arr[0]
1902                .as_object()
1903                .unwrap()
1904                .get("age")
1905                .unwrap()
1906                .as_f64()
1907                .unwrap(),
1908            35.0
1909        );
1910        assert_eq!(
1911            arr[1]
1912                .as_object()
1913                .unwrap()
1914                .get("age")
1915                .unwrap()
1916                .as_f64()
1917                .unwrap(),
1918            30.0
1919        );
1920        assert_eq!(
1921            arr[2]
1922                .as_object()
1923                .unwrap()
1924                .get("age")
1925                .unwrap()
1926                .as_f64()
1927                .unwrap(),
1928            25.0
1929        );
1930    }
1931
1932    #[test]
1933    fn test_order_by_multiple_fields() {
1934        let runtime = setup_runtime();
1935        let data = json!([
1936            {"dept": "sales", "name": "Bob"},
1937            {"dept": "eng", "name": "Alice"},
1938            {"dept": "sales", "name": "Alice"}
1939        ]);
1940        let expr = runtime
1941            .compile(r#"order_by(@, `[["dept", "asc"], ["name", "asc"]]`)"#)
1942            .unwrap();
1943        let result = expr.search(&data).unwrap();
1944        let arr = result.as_array().unwrap();
1945        // eng comes first, then sales (sorted by dept)
1946        assert_eq!(
1947            arr[0]
1948                .as_object()
1949                .unwrap()
1950                .get("dept")
1951                .unwrap()
1952                .as_str()
1953                .unwrap(),
1954            "eng"
1955        );
1956        // Within sales, Alice comes before Bob
1957        assert_eq!(
1958            arr[1]
1959                .as_object()
1960                .unwrap()
1961                .get("name")
1962                .unwrap()
1963                .as_str()
1964                .unwrap(),
1965            "Alice"
1966        );
1967        assert_eq!(
1968            arr[2]
1969                .as_object()
1970                .unwrap()
1971                .get("name")
1972                .unwrap()
1973                .as_str()
1974                .unwrap(),
1975            "Bob"
1976        );
1977    }
1978
1979    #[test]
1980    fn test_reduce_expr_sum() {
1981        let runtime = setup_runtime();
1982        let data = json!([1, 2, 3, 4, 5]);
1983        let expr = runtime
1984            .compile("reduce_expr(&sum([accumulator, current]), @, `0`)")
1985            .unwrap();
1986        let result = expr.search(&data).unwrap();
1987        assert_eq!(result.as_f64().unwrap(), 15.0);
1988    }
1989
1990    #[test]
1991    fn test_reduce_expr_max() {
1992        let runtime = setup_runtime();
1993        let data = json!([3, 1, 4, 1, 5, 9, 2, 6]);
1994        let expr = runtime
1995            .compile("reduce_expr(&max([accumulator, current]), @, `0`)")
1996            .unwrap();
1997        let result = expr.search(&data).unwrap();
1998        assert_eq!(result.as_f64().unwrap(), 9.0);
1999    }
2000
2001    #[test]
2002    fn test_reduce_expr_empty() {
2003        let runtime = setup_runtime();
2004        let data = json!([]);
2005        let expr = runtime
2006            .compile("reduce_expr(&sum([accumulator, current]), @, `42`)")
2007            .unwrap();
2008        let result = expr.search(&data).unwrap();
2009        assert_eq!(result.as_f64().unwrap(), 42.0); // Returns initial value
2010    }
2011
2012    #[test]
2013    fn test_fold_alias() {
2014        let runtime = setup_runtime();
2015        let data = json!([1, 2, 3]);
2016        let expr = runtime
2017            .compile("fold(&sum([accumulator, current]), @, `0`)")
2018            .unwrap();
2019        let result = expr.search(&data).unwrap();
2020        assert_eq!(result.as_f64().unwrap(), 6.0);
2021    }
2022
2023    #[test]
2024    fn test_scan_expr_running_sum() {
2025        let runtime = setup_runtime();
2026        let data = json!([1, 2, 3, 4]);
2027        let expr = runtime
2028            .compile("scan_expr(&sum([accumulator, current]), @, `0`)")
2029            .unwrap();
2030        let result = expr.search(&data).unwrap();
2031        let arr = result.as_array().unwrap();
2032        // Running sum: [1, 3, 6, 10]
2033        assert_eq!(arr.len(), 4);
2034        assert_eq!(arr[0].as_f64().unwrap(), 1.0);
2035        assert_eq!(arr[1].as_f64().unwrap(), 3.0);
2036        assert_eq!(arr[2].as_f64().unwrap(), 6.0);
2037        assert_eq!(arr[3].as_f64().unwrap(), 10.0);
2038    }
2039
2040    #[test]
2041    fn test_scan_expr_empty() {
2042        let runtime = setup_runtime();
2043        let data = json!([]);
2044        let expr = runtime
2045            .compile("scan_expr(&sum([accumulator, current]), @, `0`)")
2046            .unwrap();
2047        let result = expr.search(&data).unwrap();
2048        let arr = result.as_array().unwrap();
2049        assert_eq!(arr.len(), 0);
2050    }
2051
2052    #[test]
2053    fn test_reduce_expr_with_index() {
2054        let runtime = setup_runtime();
2055        // Access the index in the reduce expression
2056        let data = json!([10, 20, 30]);
2057        let expr = runtime
2058            .compile("reduce_expr(&sum([accumulator, index]), @, `0`)")
2059            .unwrap();
2060        let result = expr.search(&data).unwrap();
2061        // 0 + 1 + 2 = 3
2062        assert_eq!(result.as_f64().unwrap(), 3.0);
2063    }
2064
2065    #[test]
2066    fn test_count_by_objects() {
2067        let runtime = setup_runtime();
2068        let data = json!([
2069            {"type": "a"},
2070            {"type": "b"},
2071            {"type": "a"},
2072            {"type": "a"}
2073        ]);
2074        let expr = runtime.compile("count_by(&type, @)").unwrap();
2075        let result = expr.search(&data).unwrap();
2076        let obj = result.as_object().unwrap();
2077        assert_eq!(obj.get("a").unwrap().as_f64().unwrap(), 3.0);
2078        assert_eq!(obj.get("b").unwrap().as_f64().unwrap(), 1.0);
2079    }
2080
2081    #[test]
2082    fn test_count_by_strings() {
2083        let runtime = setup_runtime();
2084        let data = json!(["a", "b", "a", "c", "a"]);
2085        let expr = runtime.compile("count_by(&@, @)").unwrap();
2086        let result = expr.search(&data).unwrap();
2087        let obj = result.as_object().unwrap();
2088        assert_eq!(obj.get("a").unwrap().as_f64().unwrap(), 3.0);
2089        assert_eq!(obj.get("b").unwrap().as_f64().unwrap(), 1.0);
2090        assert_eq!(obj.get("c").unwrap().as_f64().unwrap(), 1.0);
2091    }
2092
2093    #[test]
2094    fn test_count_by_empty() {
2095        let runtime = setup_runtime();
2096        let data = json!([]);
2097        let expr = runtime.compile("count_by(&type, @)").unwrap();
2098        let result = expr.search(&data).unwrap();
2099        let obj = result.as_object().unwrap();
2100        assert!(obj.is_empty());
2101    }
2102
2103    #[test]
2104    fn test_count_by_numbers() {
2105        let runtime = setup_runtime();
2106        let data = json!([1, 2, 1, 3, 1, 2]);
2107        let expr = runtime.compile("count_by(&@, @)").unwrap();
2108        let result = expr.search(&data).unwrap();
2109        let obj = result.as_object().unwrap();
2110        assert_eq!(obj.get("1").unwrap().as_f64().unwrap(), 3.0);
2111        assert_eq!(obj.get("2").unwrap().as_f64().unwrap(), 2.0);
2112        assert_eq!(obj.get("3").unwrap().as_f64().unwrap(), 1.0);
2113    }
2114
2115    // =============================================================================
2116    // Partial application tests
2117    // =============================================================================
2118
2119    #[test]
2120    fn test_partial_creates_object() {
2121        let runtime = setup_runtime();
2122        let data = json!(null);
2123        let expr = runtime.compile("partial('length')").unwrap();
2124        let result = expr.search(&data).unwrap();
2125        let obj = result.as_object().unwrap();
2126        assert!(obj.get("__partial__").unwrap().as_bool().unwrap());
2127        assert_eq!(obj.get("fn").unwrap().as_str().unwrap(), "length");
2128        assert!(obj.get("args").unwrap().as_array().unwrap().is_empty());
2129    }
2130
2131    #[test]
2132    fn test_partial_with_args() {
2133        let runtime = setup_runtime();
2134        let data = json!(null);
2135        let expr = runtime
2136            .compile("partial('contains', `\"hello world\"`)")
2137            .unwrap();
2138        let result = expr.search(&data).unwrap();
2139        let obj = result.as_object().unwrap();
2140        assert!(obj.get("__partial__").unwrap().as_bool().unwrap());
2141        assert_eq!(obj.get("fn").unwrap().as_str().unwrap(), "contains");
2142        let args = obj.get("args").unwrap().as_array().unwrap();
2143        assert_eq!(args.len(), 1);
2144        assert_eq!(args[0].as_str().unwrap(), "hello world");
2145    }
2146
2147    #[test]
2148    fn test_apply_with_fn_name() {
2149        let runtime = setup_runtime();
2150        let data = json!(null);
2151        let expr = runtime.compile("apply('length', `\"hello\"`)").unwrap();
2152        let result = expr.search(&data).unwrap();
2153        assert_eq!(result.as_f64().unwrap(), 5.0);
2154    }
2155
2156    #[test]
2157    fn test_apply_with_partial() {
2158        let runtime = setup_runtime();
2159        let data = json!(null);
2160        // Create partial with first arg, then apply with second arg
2161        let expr = runtime
2162            .compile("apply(partial('contains', `\"hello world\"`), `\"world\"`)")
2163            .unwrap();
2164        let result = expr.search(&data).unwrap();
2165        assert!(result.as_bool().unwrap());
2166    }
2167
2168    #[test]
2169    fn test_apply_partial_not_found() {
2170        let runtime = setup_runtime();
2171        let data = json!(null);
2172        let expr = runtime
2173            .compile("apply(partial('contains', `\"hello world\"`), `\"xyz\"`)")
2174            .unwrap();
2175        let result = expr.search(&data).unwrap();
2176        assert!(!result.as_bool().unwrap());
2177    }
2178
2179    #[test]
2180    fn test_partial_with_multiple_prefilled_args() {
2181        let runtime = setup_runtime();
2182        let data = json!(null);
2183        // partial with 2 args pre-filled
2184        let expr = runtime.compile("partial('join', `\"-\"`)").unwrap();
2185        let result = expr.search(&data).unwrap();
2186        let obj = result.as_object().unwrap();
2187        let args = obj.get("args").unwrap().as_array().unwrap();
2188        assert_eq!(args.len(), 1);
2189        assert_eq!(args[0].as_str().unwrap(), "-");
2190    }
2191
2192    #[test]
2193    fn test_apply_partial_join() {
2194        let runtime = setup_runtime();
2195        let data = json!(null);
2196        // Create a join with "-" separator, then apply to array
2197        let expr = runtime
2198            .compile("apply(partial('join', `\"-\"`), `[\"a\", \"b\", \"c\"]`)")
2199            .unwrap();
2200        let result = expr.search(&data).unwrap();
2201        assert_eq!(result.as_str().unwrap(), "a-b-c");
2202    }
2203
2204    // =========================================================================
2205    // Pipeline pattern tests
2206    // =========================================================================
2207
2208    #[test]
2209    fn test_pipeline_filter_sort_products() {
2210        let runtime = setup_runtime();
2211        let data = json!({
2212            "products": [
2213                {"name": "A", "price": 30, "in_stock": true},
2214                {"name": "B", "price": 10, "in_stock": true},
2215                {"name": "C", "price": 20, "in_stock": false},
2216                {"name": "D", "price": 5, "in_stock": true}
2217            ]
2218        });
2219        let expr = runtime
2220            .compile("products | filter_expr(&in_stock, @) | sort_by_expr(&price, @)")
2221            .unwrap();
2222        let result = expr.search(&data).unwrap();
2223        let arr = result.as_array().unwrap();
2224        assert_eq!(arr.len(), 3);
2225        assert_eq!(
2226            arr[0]
2227                .as_object()
2228                .unwrap()
2229                .get("name")
2230                .unwrap()
2231                .as_str()
2232                .unwrap(),
2233            "D"
2234        ); // $5
2235        assert_eq!(
2236            arr[1]
2237                .as_object()
2238                .unwrap()
2239                .get("name")
2240                .unwrap()
2241                .as_str()
2242                .unwrap(),
2243            "B"
2244        ); // $10
2245    }
2246
2247    #[test]
2248    fn test_pipeline_funnel_errors() {
2249        let runtime = setup_runtime();
2250        let data = json!({
2251            "events": [
2252                {"level": "error", "timestamp": 1704067300, "message": "Disk full"},
2253                {"level": "info", "timestamp": 1704067200, "message": "Started"},
2254                {"level": "error", "timestamp": 1704067400, "message": "Connection lost"},
2255                {"level": "warn", "timestamp": 1704067350, "message": "High memory"}
2256            ]
2257        });
2258        let expr = runtime
2259            .compile(
2260                r#"events | filter_expr(&(level == `"error"`), @) | sort_by_expr(&timestamp, @)"#,
2261            )
2262            .unwrap();
2263        let result = expr.search(&data).unwrap();
2264        let arr = result.as_array().unwrap();
2265        assert_eq!(arr.len(), 2);
2266        // Sorted by timestamp ascending
2267        assert_eq!(
2268            arr[0]
2269                .as_object()
2270                .unwrap()
2271                .get("message")
2272                .unwrap()
2273                .as_str()
2274                .unwrap(),
2275            "Disk full"
2276        );
2277    }
2278
2279    #[test]
2280    fn test_pipeline_transactions_completed() {
2281        let runtime = setup_runtime();
2282        let data = json!({
2283            "transactions": [
2284                {"amount": 100, "status": "completed"},
2285                {"amount": 50, "status": "completed"},
2286                {"amount": 75, "status": "pending"},
2287                {"amount": 200, "status": "completed"}
2288            ]
2289        });
2290        let expr = runtime
2291            .compile(
2292                r#"transactions | filter_expr(&(status == `"completed"`), @) | map_expr(&amount, @)"#,
2293            )
2294            .unwrap();
2295        let result = expr.search(&data).unwrap();
2296        let arr = result.as_array().unwrap();
2297        assert_eq!(arr.len(), 3);
2298        assert_eq!(arr[0].as_f64().unwrap(), 100.0);
2299        assert_eq!(arr[1].as_f64().unwrap(), 50.0);
2300        assert_eq!(arr[2].as_f64().unwrap(), 200.0);
2301    }
2302
2303    #[test]
2304    fn test_pipeline_fork_join() {
2305        let runtime = setup_runtime();
2306        let data = json!({
2307            "items": [
2308                {"name": "A", "price": 150},
2309                {"name": "B", "price": 50},
2310                {"name": "C", "price": 200},
2311                {"name": "D", "price": 25}
2312            ]
2313        });
2314        let expr = runtime
2315            .compile(
2316                r#"@.{
2317                    expensive: items | filter_expr(&(price > `100`), @),
2318                    cheap: items | filter_expr(&(price <= `100`), @)
2319                }"#,
2320            )
2321            .unwrap();
2322        let result = expr.search(&data).unwrap();
2323        let obj = result.as_object().unwrap();
2324        assert_eq!(obj.get("expensive").unwrap().as_array().unwrap().len(), 2);
2325        assert_eq!(obj.get("cheap").unwrap().as_array().unwrap().len(), 2);
2326    }
2327
2328    #[test]
2329    fn test_pipeline_nested_users() {
2330        let runtime = setup_runtime();
2331        let data = json!({
2332            "users": [
2333                {"name": "Alice", "orders": [{"total": 100}, {"total": 50}]},
2334                {"name": "Bob", "orders": [{"total": 200}]},
2335                {"name": "Carol", "orders": []}
2336            ]
2337        });
2338        // Filter users with orders, then map to get names
2339        let expr = runtime
2340            .compile("users | filter_expr(&(length(orders) > `0`), @) | map_expr(&name, @)")
2341            .unwrap();
2342        let result = expr.search(&data).unwrap();
2343        let arr = result.as_array().unwrap();
2344        assert_eq!(arr.len(), 2);
2345        assert_eq!(arr[0].as_str().unwrap(), "Alice");
2346        assert_eq!(arr[1].as_str().unwrap(), "Bob");
2347    }
2348
2349    #[test]
2350    fn test_pipeline_rag_chunks() {
2351        let runtime = setup_runtime();
2352        let data = json!({
2353            "chunks": [
2354                {"content": "Redis is fast", "score": 0.9},
2355                {"content": "Redis is in-memory", "score": 0.85},
2356                {"content": "Unrelated content", "score": 0.5},
2357                {"content": "Redis supports modules", "score": 0.75}
2358            ]
2359        });
2360        let expr = runtime
2361            .compile("chunks | filter_expr(&(score > `0.7`), @) | sort_by_expr(&score, @)")
2362            .unwrap();
2363        let result = expr.search(&data).unwrap();
2364        let arr = result.as_array().unwrap();
2365        assert_eq!(arr.len(), 3);
2366        // Sorted ascending by score
2367        assert_eq!(
2368            arr[0]
2369                .as_object()
2370                .unwrap()
2371                .get("score")
2372                .unwrap()
2373                .as_f64()
2374                .unwrap(),
2375            0.75
2376        );
2377    }
2378
2379    // =========================================================================
2380    // Additional reduce_expr/scan_expr tests
2381    // =========================================================================
2382
2383    #[test]
2384    fn test_reduce_expr_product() {
2385        let runtime = setup_runtime();
2386        // Test reduce with min (similar to existing max test but finds minimum)
2387        let data = json!([5, 3, 8, 1, 9]);
2388        let expr = runtime
2389            .compile("reduce_expr(&min([accumulator, current]), @, `100`)")
2390            .unwrap();
2391        let result = expr.search(&data).unwrap();
2392        assert_eq!(result.as_f64().unwrap(), 1.0);
2393    }
2394
2395    #[test]
2396    fn test_scan_expr_running_balance() {
2397        let runtime = setup_runtime();
2398        // Test scan with running max - shows progressive maximum
2399        let data = json!([3, 1, 4, 1, 5, 9]);
2400        let expr = runtime
2401            .compile("scan_expr(&max([accumulator, current]), @, `0`)")
2402            .unwrap();
2403        let result = expr.search(&data).unwrap();
2404        let arr = result.as_array().unwrap();
2405        // Running max: 3, 3, 4, 4, 5, 9
2406        assert_eq!(arr[0].as_f64().unwrap(), 3.0);
2407        assert_eq!(arr[1].as_f64().unwrap(), 3.0);
2408        assert_eq!(arr[2].as_f64().unwrap(), 4.0);
2409        assert_eq!(arr[3].as_f64().unwrap(), 4.0);
2410        assert_eq!(arr[4].as_f64().unwrap(), 5.0);
2411        assert_eq!(arr[5].as_f64().unwrap(), 9.0);
2412    }
2413
2414    // =========================================================================
2415    // Additional order_by tests
2416    // =========================================================================
2417
2418    #[test]
2419    fn test_order_by_three_fields() {
2420        let runtime = setup_runtime();
2421        let data = json!([
2422            {"dept": "Engineering", "level": "senior", "name": "Charlie"},
2423            {"dept": "Engineering", "level": "junior", "name": "Alice"},
2424            {"dept": "Engineering", "level": "senior", "name": "Bob"},
2425            {"dept": "Sales", "level": "senior", "name": "David"}
2426        ]);
2427        let expr = runtime
2428            .compile(r#"order_by(@, `[["dept", "asc"], ["level", "desc"], ["name", "asc"]]`)"#)
2429            .unwrap();
2430        let result = expr.search(&data).unwrap();
2431        let arr = result.as_array().unwrap();
2432        // Engineering seniors first (alphabetical), then Engineering juniors, then Sales
2433        assert_eq!(
2434            arr[0]
2435                .as_object()
2436                .unwrap()
2437                .get("name")
2438                .unwrap()
2439                .as_str()
2440                .unwrap(),
2441            "Bob"
2442        );
2443        assert_eq!(
2444            arr[1]
2445                .as_object()
2446                .unwrap()
2447                .get("name")
2448                .unwrap()
2449                .as_str()
2450                .unwrap(),
2451            "Charlie"
2452        );
2453    }
2454
2455    #[test]
2456    fn test_order_by_empty() {
2457        let runtime = setup_runtime();
2458        let data = json!([]);
2459        let expr = runtime
2460            .compile(r#"order_by(@, `[["name", "asc"]]`)"#)
2461            .unwrap();
2462        let result = expr.search(&data).unwrap();
2463        let arr = result.as_array().unwrap();
2464        assert!(arr.is_empty());
2465    }
2466
2467    // =========================================================================
2468    // Additional partition_expr tests
2469    // =========================================================================
2470
2471    #[test]
2472    fn test_partition_expr_scores() {
2473        let runtime = setup_runtime();
2474        let data = json!([85, 42, 91, 67, 55, 78, 33, 99]);
2475        let expr = runtime.compile("partition_expr(&(@ >= `60`), @)").unwrap();
2476        let result = expr.search(&data).unwrap();
2477        let arr = result.as_array().unwrap();
2478        let passing = arr[0].as_array().unwrap();
2479        let failing = arr[1].as_array().unwrap();
2480        assert_eq!(passing.len(), 5); // 85, 91, 67, 78, 99
2481        assert_eq!(failing.len(), 3); // 42, 55, 33
2482    }
2483
2484    #[test]
2485    fn test_partition_expr_active() {
2486        let runtime = setup_runtime();
2487        let data = json!([{"active": true}, {"active": false}, {"active": true}]);
2488        let expr = runtime.compile("partition_expr(&active, @)").unwrap();
2489        let result = expr.search(&data).unwrap();
2490        let arr = result.as_array().unwrap();
2491        assert_eq!(arr[0].as_array().unwrap().len(), 2);
2492        assert_eq!(arr[1].as_array().unwrap().len(), 1);
2493    }
2494
2495    // =========================================================================
2496    // Additional map_values/map_keys tests
2497    // =========================================================================
2498
2499    #[test]
2500    fn test_map_values_discount() {
2501        let runtime = setup_runtime();
2502        // Test with string transformation since nested expressions don't have extension math functions
2503        let data = json!({"apple": "FRUIT", "banana": "ITEM"});
2504        let expr = runtime.compile("map_values(&length(@), @)").unwrap();
2505        let result = expr.search(&data).unwrap();
2506        let obj = result.as_object().unwrap();
2507        assert_eq!(obj.get("apple").unwrap().as_f64().unwrap(), 5.0);
2508        assert_eq!(obj.get("banana").unwrap().as_f64().unwrap(), 4.0);
2509    }
2510
2511    // =========================================================================
2512    // Additional group_by_expr tests
2513    // =========================================================================
2514
2515    #[test]
2516    fn test_group_by_expr_type() {
2517        let runtime = setup_runtime();
2518        let data = json!([
2519            {"type": "fruit", "name": "apple"},
2520            {"type": "vegetable", "name": "carrot"},
2521            {"type": "fruit", "name": "banana"}
2522        ]);
2523        let expr = runtime.compile("group_by_expr(&type, @)").unwrap();
2524        let result = expr.search(&data).unwrap();
2525        let obj = result.as_object().unwrap();
2526        assert_eq!(obj.get("fruit").unwrap().as_array().unwrap().len(), 2);
2527        assert_eq!(obj.get("vegetable").unwrap().as_array().unwrap().len(), 1);
2528    }
2529
2530    #[test]
2531    fn test_group_by_expr_computed() {
2532        let runtime = setup_runtime();
2533        // Group strings by their length using built-in length function
2534        let data = json!(["a", "bb", "ccc", "dd", "eee", "f"]);
2535        let expr = runtime
2536            .compile("group_by_expr(&to_string(length(@)), @)")
2537            .unwrap();
2538        let result = expr.search(&data).unwrap();
2539        let obj = result.as_object().unwrap();
2540        assert!(obj.contains_key("1")); // "a", "f"
2541        assert!(obj.contains_key("2")); // "bb", "dd"
2542        assert!(obj.contains_key("3")); // "ccc", "eee"
2543        assert_eq!(obj.get("1").unwrap().as_array().unwrap().len(), 2);
2544        assert_eq!(obj.get("2").unwrap().as_array().unwrap().len(), 2);
2545        assert_eq!(obj.get("3").unwrap().as_array().unwrap().len(), 2);
2546    }
2547
2548    // =========================================================================
2549    // Additional unique_by_expr tests
2550    // =========================================================================
2551
2552    #[test]
2553    fn test_unique_by_expr_id() {
2554        let runtime = setup_runtime();
2555        let data = json!([
2556            {"id": 1, "v": "a"},
2557            {"id": 2, "v": "b"},
2558            {"id": 1, "v": "c"}
2559        ]);
2560        let expr = runtime.compile("unique_by_expr(&id, @)").unwrap();
2561        let result = expr.search(&data).unwrap();
2562        let arr = result.as_array().unwrap();
2563        assert_eq!(arr.len(), 2);
2564        // Keeps first occurrence
2565        assert_eq!(
2566            arr[0]
2567                .as_object()
2568                .unwrap()
2569                .get("v")
2570                .unwrap()
2571                .as_str()
2572                .unwrap(),
2573            "a"
2574        );
2575    }
2576
2577    // =========================================================================
2578    // Edge case tests
2579    // =========================================================================
2580
2581    #[test]
2582    fn test_any_expr_empty() {
2583        let runtime = setup_runtime();
2584        let data = json!([]);
2585        let expr = runtime.compile("any_expr(&(@ > `0`), @)").unwrap();
2586        let result = expr.search(&data).unwrap();
2587        assert!(!result.as_bool().unwrap());
2588    }
2589
2590    #[test]
2591    fn test_max_by_expr_empty() {
2592        let runtime = setup_runtime();
2593        let data = json!([]);
2594        let expr = runtime.compile("max_by_expr(&age, @)").unwrap();
2595        let result = expr.search(&data).unwrap();
2596        assert!(result.is_null());
2597    }
2598
2599    #[test]
2600    fn test_flat_map_expr_duplicate() {
2601        let runtime = setup_runtime();
2602        let data = json!([1, 2, 3]);
2603        // Duplicate each element
2604        let expr = runtime.compile("flat_map_expr(&[@, @], @)").unwrap();
2605        let result = expr.search(&data).unwrap();
2606        let arr = result.as_array().unwrap();
2607        assert_eq!(arr.len(), 6);
2608    }
2609
2610    #[test]
2611    fn test_reject_greater_than() {
2612        let runtime = setup_runtime();
2613        let data = json!([1, 2, 3, 4, 5, 6]);
2614        let expr = runtime.compile("reject(&(@ > `3`), @)").unwrap();
2615        let result = expr.search(&data).unwrap();
2616        let arr = result.as_array().unwrap();
2617        assert_eq!(arr.len(), 3); // 1, 2, 3
2618    }
2619
2620    #[test]
2621    fn test_every_false_case() {
2622        let runtime = setup_runtime();
2623        let data = json!([1, -1, 3]);
2624        let expr = runtime.compile("every(&(@ > `0`), @)").unwrap();
2625        let result = expr.search(&data).unwrap();
2626        assert!(!result.as_bool().unwrap());
2627    }
2628
2629    #[test]
2630    fn test_count_expr_all_match() {
2631        let runtime = setup_runtime();
2632        let data = json!([5, 10, 15, 20]);
2633        let expr = runtime.compile("count_expr(&(@ > `0`), @)").unwrap();
2634        let result = expr.search(&data).unwrap();
2635        assert_eq!(result.as_f64().unwrap(), 4.0);
2636    }
2637
2638    #[test]
2639    fn test_find_expr_first_match() {
2640        let runtime = setup_runtime();
2641        let data = json!([1, 5, 10, 15]);
2642        let expr = runtime.compile("find_expr(&(@ > `3`), @)").unwrap();
2643        let result = expr.search(&data).unwrap();
2644        assert_eq!(result.as_f64().unwrap(), 5.0);
2645    }
2646
2647    #[test]
2648    fn test_find_index_expr_first_match() {
2649        let runtime = setup_runtime();
2650        let data = json!([1, 5, 10, 15]);
2651        let expr = runtime.compile("find_index_expr(&(@ > `3`), @)").unwrap();
2652        let result = expr.search(&data).unwrap();
2653        assert_eq!(result.as_f64().unwrap(), 1.0);
2654    }
2655
2656    #[test]
2657    fn test_take_while_basic() {
2658        let runtime = setup_runtime();
2659        let data = json!([1, 2, 3, 5, 1, 2]);
2660        let expr = runtime.compile("take_while(&(@ < `4`), @)").unwrap();
2661        let result = expr.search(&data).unwrap();
2662        let arr = result.as_array().unwrap();
2663        assert_eq!(arr.len(), 3);
2664        assert_eq!(arr[0].as_f64().unwrap(), 1.0);
2665        assert_eq!(arr[1].as_f64().unwrap(), 2.0);
2666        assert_eq!(arr[2].as_f64().unwrap(), 3.0);
2667    }
2668
2669    #[test]
2670    fn test_take_while_all_match() {
2671        let runtime = setup_runtime();
2672        let data = json!([1, 2, 3]);
2673        let expr = runtime.compile("take_while(&(@ < `10`), @)").unwrap();
2674        let result = expr.search(&data).unwrap();
2675        let arr = result.as_array().unwrap();
2676        assert_eq!(arr.len(), 3);
2677    }
2678
2679    #[test]
2680    fn test_take_while_none_match() {
2681        let runtime = setup_runtime();
2682        let data = json!([5, 6, 7]);
2683        let expr = runtime.compile("take_while(&(@ < `4`), @)").unwrap();
2684        let result = expr.search(&data).unwrap();
2685        let arr = result.as_array().unwrap();
2686        assert_eq!(arr.len(), 0);
2687    }
2688
2689    #[test]
2690    fn test_drop_while_basic() {
2691        let runtime = setup_runtime();
2692        let data = json!([1, 2, 3, 5, 1, 2]);
2693        let expr = runtime.compile("drop_while(&(@ < `4`), @)").unwrap();
2694        let result = expr.search(&data).unwrap();
2695        let arr = result.as_array().unwrap();
2696        assert_eq!(arr.len(), 3);
2697        assert_eq!(arr[0].as_f64().unwrap(), 5.0);
2698        assert_eq!(arr[1].as_f64().unwrap(), 1.0);
2699        assert_eq!(arr[2].as_f64().unwrap(), 2.0);
2700    }
2701
2702    #[test]
2703    fn test_drop_while_all_match() {
2704        let runtime = setup_runtime();
2705        let data = json!([1, 2, 3]);
2706        let expr = runtime.compile("drop_while(&(@ < `10`), @)").unwrap();
2707        let result = expr.search(&data).unwrap();
2708        let arr = result.as_array().unwrap();
2709        assert_eq!(arr.len(), 0);
2710    }
2711
2712    #[test]
2713    fn test_drop_while_none_match() {
2714        let runtime = setup_runtime();
2715        let data = json!([5, 6, 7]);
2716        let expr = runtime.compile("drop_while(&(@ < `4`), @)").unwrap();
2717        let result = expr.search(&data).unwrap();
2718        let arr = result.as_array().unwrap();
2719        assert_eq!(arr.len(), 3);
2720    }
2721
2722    #[test]
2723    fn test_zip_with_add() {
2724        let runtime = setup_runtime();
2725        let data = json!(null);
2726        let expr = runtime
2727            .compile("zip_with(&add([0], [1]), `[1, 2, 3]`, `[10, 20, 30]`)")
2728            .unwrap();
2729        let result = expr.search(&data).unwrap();
2730        let arr = result.as_array().unwrap();
2731        assert_eq!(arr.len(), 3);
2732        assert_eq!(arr[0].as_f64().unwrap(), 11.0);
2733        assert_eq!(arr[1].as_f64().unwrap(), 22.0);
2734        assert_eq!(arr[2].as_f64().unwrap(), 33.0);
2735    }
2736
2737    #[test]
2738    fn test_zip_with_unequal_lengths() {
2739        let runtime = setup_runtime();
2740        let data = json!(null);
2741        let expr = runtime
2742            .compile("zip_with(&add([0], [1]), `[1, 2, 3, 4, 5]`, `[10, 20]`)")
2743            .unwrap();
2744        let result = expr.search(&data).unwrap();
2745        let arr = result.as_array().unwrap();
2746        assert_eq!(arr.len(), 2);
2747        assert_eq!(arr[0].as_f64().unwrap(), 11.0);
2748        assert_eq!(arr[1].as_f64().unwrap(), 22.0);
2749    }
2750
2751    #[test]
2752    fn test_zip_with_multiply() {
2753        let runtime = setup_runtime();
2754        let data = json!(null);
2755        let expr = runtime
2756            .compile("zip_with(&multiply([0], [1]), `[2, 3, 4]`, `[5, 6, 7]`)")
2757            .unwrap();
2758        let result = expr.search(&data).unwrap();
2759        let arr = result.as_array().unwrap();
2760        assert_eq!(arr.len(), 3);
2761        assert_eq!(arr[0].as_f64().unwrap(), 10.0);
2762        assert_eq!(arr[1].as_f64().unwrap(), 18.0);
2763        assert_eq!(arr[2].as_f64().unwrap(), 28.0);
2764    }
2765
2766    // =========================================================================
2767    // walk tests
2768    // =========================================================================
2769
2770    #[test]
2771    fn test_walk_identity() {
2772        let runtime = setup_runtime();
2773        let data = json!({"a": [1, 2, 3], "b": {"c": 4}});
2774        let expr = runtime.compile("walk(&@, @)").unwrap();
2775        let result = expr.search(&data).unwrap();
2776        // Identity should return the same structure
2777        assert!(result.is_object());
2778        let obj = result.as_object().unwrap();
2779        assert!(obj.contains_key("a"));
2780        assert!(obj.contains_key("b"));
2781    }
2782
2783    #[test]
2784    fn test_walk_type_of_all() {
2785        let runtime = setup_runtime();
2786        let data = json!({"a": 5, "b": [1, 2]});
2787        // type() works on everything - shows bottom-up processing
2788        let expr = runtime.compile("walk(&type(@), @)").unwrap();
2789        let result = expr.search(&data).unwrap();
2790        // After walking, everything becomes its type string, and the final result
2791        // is type of the top-level result
2792        assert_eq!(result.as_str().unwrap(), "object");
2793    }
2794
2795    #[test]
2796    fn test_walk_nested_arrays() {
2797        let runtime = setup_runtime();
2798        // Use only arrays (no scalars inside) so length works at every level
2799        let data = json!([[[], []], [[]]]);
2800        // length works on arrays - get lengths at each level
2801        let expr = runtime.compile("walk(&length(@), @)").unwrap();
2802        let result = expr.search(&data).unwrap();
2803        // Inner [] -> 0, outer arrays get lengths, top level has 2 elements
2804        assert_eq!(result.as_f64().unwrap(), 2.0);
2805    }
2806
2807    #[test]
2808    fn test_walk_scalar() {
2809        let runtime = setup_runtime();
2810        let data = json!(5);
2811        // Double the number
2812        let expr = runtime.compile("walk(&multiply(@, `2`), @)").unwrap();
2813        let result = expr.search(&data).unwrap();
2814        assert_eq!(result.as_f64().unwrap(), 10.0);
2815    }
2816
2817    #[test]
2818    fn test_walk_length_all() {
2819        let runtime = setup_runtime();
2820        let data = json!({"items": ["a", "bb", "ccc"]});
2821        // Get length of everything (works for strings, arrays, objects)
2822        let expr = runtime.compile("walk(&length(@), @)").unwrap();
2823        let result = expr.search(&data).unwrap();
2824        // Top level object has 1 key
2825        assert_eq!(result.as_f64().unwrap(), 1.0);
2826    }
2827
2828    #[test]
2829    fn test_walk_preserves_structure() {
2830        let runtime = setup_runtime();
2831        let data = json!({"a": [1, {"b": 2}], "c": "hello"});
2832        // Identity transform - should preserve structure
2833        let expr = runtime.compile("walk(&@, @)").unwrap();
2834        let result = expr.search(&data).unwrap();
2835
2836        let obj = result.as_object().unwrap();
2837        assert!(obj.contains_key("a"));
2838        assert!(obj.contains_key("c"));
2839        let arr = obj.get("a").unwrap().as_array().unwrap();
2840        assert_eq!(arr.len(), 2);
2841    }
2842
2843    #[test]
2844    fn test_walk_empty_structures() {
2845        let runtime = setup_runtime();
2846
2847        // Empty array
2848        let data = json!([]);
2849        let expr = runtime.compile("walk(&@, @)").unwrap();
2850        let result = expr.search(&data).unwrap();
2851        assert!(result.as_array().unwrap().is_empty());
2852
2853        // Empty object
2854        let data = json!({});
2855        let result = expr.search(&data).unwrap();
2856        assert!(result.as_object().unwrap().is_empty());
2857    }
2858
2859    // =========================================================================
2860    // recurse tests
2861    // =========================================================================
2862
2863    #[test]
2864    fn test_recurse_nested_object() {
2865        let runtime = setup_runtime();
2866        let data = json!({"a": {"b": 1}});
2867        let expr = runtime.compile("recurse(@)").unwrap();
2868        let result = expr.search(&data).unwrap();
2869        let arr = result.as_array().unwrap();
2870        // Should have: {a: {b: 1}}, {b: 1}, 1
2871        assert_eq!(arr.len(), 3);
2872    }
2873
2874    #[test]
2875    fn test_recurse_nested_array() {
2876        let runtime = setup_runtime();
2877        let data = json!([1, [2, 3]]);
2878        let expr = runtime.compile("recurse(@)").unwrap();
2879        let result = expr.search(&data).unwrap();
2880        let arr = result.as_array().unwrap();
2881        // Should have: [1, [2, 3]], 1, [2, 3], 2, 3
2882        assert_eq!(arr.len(), 5);
2883    }
2884
2885    #[test]
2886    fn test_recurse_scalar() {
2887        let runtime = setup_runtime();
2888        let data = json!(42);
2889        let expr = runtime.compile("recurse(@)").unwrap();
2890        let result = expr.search(&data).unwrap();
2891        let arr = result.as_array().unwrap();
2892        // Just the scalar itself
2893        assert_eq!(arr.len(), 1);
2894        assert_eq!(arr[0].as_f64().unwrap(), 42.0);
2895    }
2896
2897    #[test]
2898    fn test_recurse_with_field() {
2899        let runtime = setup_runtime();
2900        let data = json!({"a": {"a": {"a": null}}});
2901        let expr = runtime.compile("recurse_with(@, &a)").unwrap();
2902        let result = expr.search(&data).unwrap();
2903        let arr = result.as_array().unwrap();
2904        // Should have: {a: {a: {a: null}}}, {a: {a: null}}, {a: null}
2905        assert_eq!(arr.len(), 3);
2906    }
2907
2908    #[test]
2909    fn test_recurse_with_children() {
2910        let runtime = setup_runtime();
2911        let data = json!({
2912            "name": "root",
2913            "children": [
2914                {"name": "child1", "children": []},
2915                {"name": "child2", "children": []}
2916            ]
2917        });
2918        let expr = runtime.compile("recurse_with(@, &children[])").unwrap();
2919        let result = expr.search(&data).unwrap();
2920        let arr = result.as_array().unwrap();
2921        // Should have: root, child1, child2
2922        assert_eq!(arr.len(), 3);
2923    }
2924
2925    // =========================================================================
2926    // while_expr tests
2927    // =========================================================================
2928
2929    #[test]
2930    fn test_while_expr_doubling() {
2931        let runtime = setup_runtime();
2932        let data = json!(null);
2933        let expr = runtime
2934            .compile("while_expr(`1`, &(@ < `100`), &multiply(@, `2`))")
2935            .unwrap();
2936        let result = expr.search(&data).unwrap();
2937        // 1 -> 2 -> 4 -> 8 -> 16 -> 32 -> 64 -> 128 (stops because 128 >= 100)
2938        assert_eq!(result.as_f64().unwrap(), 128.0);
2939    }
2940
2941    #[test]
2942    fn test_while_expr_counter() {
2943        let runtime = setup_runtime();
2944        let data = json!(null);
2945        let expr = runtime
2946            .compile("while_expr(`0`, &(@ < `5`), &add(@, `1`))")
2947            .unwrap();
2948        let result = expr.search(&data).unwrap();
2949        // 0 -> 1 -> 2 -> 3 -> 4 -> 5 (stops because 5 >= 5)
2950        assert_eq!(result.as_f64().unwrap(), 5.0);
2951    }
2952
2953    #[test]
2954    fn test_while_expr_immediate_false() {
2955        let runtime = setup_runtime();
2956        let data = json!(null);
2957        let expr = runtime
2958            .compile("while_expr(`100`, &(@ < `10`), &add(@, `1`))")
2959            .unwrap();
2960        let result = expr.search(&data).unwrap();
2961        // Condition is false immediately, return initial value
2962        assert_eq!(result.as_f64().unwrap(), 100.0);
2963    }
2964
2965    // =========================================================================
2966    // until_expr tests
2967    // =========================================================================
2968
2969    #[test]
2970    fn test_until_expr_doubling() {
2971        let runtime = setup_runtime();
2972        let data = json!(null);
2973        let expr = runtime
2974            .compile("until_expr(`1`, &(@ >= `100`), &multiply(@, `2`))")
2975            .unwrap();
2976        let result = expr.search(&data).unwrap();
2977        // 1 -> 2 -> 4 -> 8 -> 16 -> 32 -> 64 -> 128 (stops because 128 >= 100)
2978        assert_eq!(result.as_f64().unwrap(), 128.0);
2979    }
2980
2981    #[test]
2982    fn test_until_expr_counter() {
2983        let runtime = setup_runtime();
2984        let data = json!(null);
2985        let expr = runtime
2986            .compile("until_expr(`0`, &(@ == `5`), &add(@, `1`))")
2987            .unwrap();
2988        let result = expr.search(&data).unwrap();
2989        // 0 -> 1 -> 2 -> 3 -> 4 -> 5 (stops because 5 == 5)
2990        assert_eq!(result.as_f64().unwrap(), 5.0);
2991    }
2992
2993    #[test]
2994    fn test_until_expr_immediate_true() {
2995        let runtime = setup_runtime();
2996        let data = json!(null);
2997        let expr = runtime
2998            .compile("until_expr(`100`, &(@ > `10`), &add(@, `1`))")
2999            .unwrap();
3000        let result = expr.search(&data).unwrap();
3001        // Condition is true immediately, return initial value
3002        assert_eq!(result.as_f64().unwrap(), 100.0);
3003    }
3004}