1use crate::eval::{evaluate_expr, EvalCtx};
4use crate::parser::ast::Expr;
5use crate::types::{ErrorKind, Value};
6
7use super::{check_arity, check_arity_len, FunctionMeta, Registry};
8
9pub fn to_2d(v: &Value) -> Vec<Vec<Value>> {
16 match v {
17 Value::Array(outer) => {
18 if outer.iter().any(|e| matches!(e, Value::Array(_))) {
19 outer
20 .iter()
21 .map(|row| match row {
22 Value::Array(cols) => cols.clone(),
23 other => vec![other.clone()],
24 })
25 .collect()
26 } else {
27 vec![outer.clone()] }
29 }
30 other => vec![vec![other.clone()]], }
32}
33
34pub fn from_2d(rows: Vec<Vec<Value>>) -> Value {
39 if rows.is_empty() {
40 return Value::Array(vec![]);
41 }
42 if rows.len() == 1 {
43 return Value::Array(rows.into_iter().next().unwrap());
44 }
45 Value::Array(rows.into_iter().map(Value::Array).collect())
46}
47
48pub fn flatten_val(v: &Value) -> Vec<Value> {
50 match v {
51 Value::Array(outer) => {
52 if outer.iter().any(|e| matches!(e, Value::Array(_))) {
53 outer
54 .iter()
55 .flat_map(|row| match row {
56 Value::Array(cols) => cols.clone(),
57 other => vec![other.clone()],
58 })
59 .collect()
60 } else {
61 outer.clone()
62 }
63 }
64 other => vec![other.clone()],
65 }
66}
67
68fn to_f64(v: &Value) -> Option<f64> {
70 match v {
71 Value::Number(n) => Some(*n),
72 Value::Bool(b) => Some(if *b { 1.0 } else { 0.0 }),
73 _ => None,
74 }
75}
76
77pub(crate) fn rows_fn(args: &[Value]) -> Value {
80 if let Some(e) = check_arity(args, 1, 1) {
81 return e;
82 }
83 let grid = to_2d(&args[0]);
84 Value::Number(grid.len() as f64)
85}
86
87pub(crate) fn columns_fn(args: &[Value]) -> Value {
90 if let Some(e) = check_arity(args, 1, 1) {
91 return e;
92 }
93 let grid = to_2d(&args[0]);
94 let cols = grid.first().map(|r| r.len()).unwrap_or(0);
95 Value::Number(cols as f64)
96}
97
98fn index_fn(args: &[Value]) -> Value {
102 if let Some(e) = check_arity(args, 1, 3) {
103 return e;
104 }
105 let grid = to_2d(&args[0]);
106 let nrows = grid.len();
107 let ncols = grid.first().map(|r| r.len()).unwrap_or(0);
108
109 let row_idx = if args.len() >= 2 {
110 match to_f64(&args[1]) {
111 Some(n) => {
112 let r = n as isize;
113 if r < 0 {
114 (nrows as isize + r + 1) as usize
115 } else {
116 r as usize
117 }
118 }
119 None => return Value::Error(ErrorKind::Value),
120 }
121 } else {
122 0 };
124
125 let col_idx = if args.len() >= 3 {
126 match to_f64(&args[2]) {
127 Some(n) => {
128 let c = n as isize;
129 if c < 0 {
130 (ncols as isize + c + 1) as usize
131 } else {
132 c as usize
133 }
134 }
135 None => return Value::Error(ErrorKind::Value),
136 }
137 } else {
138 0 };
140
141 if row_idx > 0 && col_idx > 0 {
143 if row_idx > nrows || col_idx > ncols {
144 return Value::Error(ErrorKind::Ref);
145 }
146 return grid[row_idx - 1][col_idx - 1].clone();
147 }
148
149 if row_idx > 0 && col_idx == 0 {
151 if nrows == 1 && ncols > 1 {
153 if row_idx > ncols {
154 return Value::Error(ErrorKind::Ref);
155 }
156 return grid[0][row_idx - 1].clone();
157 }
158 if row_idx > nrows {
159 return Value::Error(ErrorKind::Ref);
160 }
161 let row = grid[row_idx - 1].clone();
162 if row.len() == 1 {
163 return row.into_iter().next().unwrap();
164 }
165 return Value::Array(row);
166 }
167
168 if row_idx == 0 && col_idx > 0 {
170 if col_idx > ncols {
171 return Value::Error(ErrorKind::Ref);
172 }
173 let col: Vec<Value> = grid.iter().map(|r| r[col_idx - 1].clone()).collect();
174 if col.len() == 1 {
175 return col.into_iter().next().unwrap();
176 }
177 return from_2d(col.into_iter().map(|v| vec![v]).collect());
178 }
179
180 args[0].clone()
182}
183
184pub(crate) fn transpose_fn(args: &[Value]) -> Value {
187 if let Some(e) = check_arity(args, 1, 1) {
188 return e;
189 }
190 let grid = to_2d(&args[0]);
191 if grid.is_empty() {
192 return Value::Array(vec![]);
193 }
194 let nrows = grid.len();
195 let ncols = grid[0].len();
196 let transposed: Vec<Vec<Value>> = (0..ncols)
197 .map(|c| (0..nrows).map(|r| grid[r][c].clone()).collect())
198 .collect();
199 from_2d(transposed)
200}
201
202pub(crate) fn array_constrain_fn(args: &[Value]) -> Value {
205 if let Some(e) = check_arity(args, 3, 3) {
206 return e;
207 }
208 let grid = to_2d(&args[0]);
209 let num_rows = match to_f64(&args[1]) {
210 Some(n) if n >= 1.0 => n as usize,
211 _ => return Value::Error(ErrorKind::Value),
212 };
213 let num_cols = match to_f64(&args[2]) {
214 Some(n) if n >= 1.0 => n as usize,
215 _ => return Value::Error(ErrorKind::Value),
216 };
217 let rows_to_take = num_rows.min(grid.len());
218 let result: Vec<Vec<Value>> = grid[..rows_to_take]
219 .iter()
220 .map(|row| {
221 let cols_to_take = num_cols.min(row.len());
222 row[..cols_to_take].to_vec()
223 })
224 .collect();
225 from_2d(result)
226}
227
228fn choosecols_fn(args: &[Value]) -> Value {
231 if let Some(e) = check_arity(args, 2, usize::MAX) {
232 return e;
233 }
234 let grid = to_2d(&args[0]);
235 let ncols = grid.first().map(|r| r.len()).unwrap_or(0);
236 let mut selected_cols: Vec<usize> = Vec::new();
237 for col_arg in &args[1..] {
238 match to_f64(col_arg) {
239 Some(0.0) => return Value::Error(ErrorKind::Value),
240 Some(n) => {
241 let idx = if n < 0.0 {
242 let i = (ncols as isize + n as isize) as usize;
243 if n as isize + (ncols as isize) < 0 {
244 return Value::Error(ErrorKind::Value);
245 }
246 i
247 } else {
248 let i = n as usize - 1;
249 if i >= ncols {
250 return Value::Error(ErrorKind::Value);
251 }
252 i
253 };
254 selected_cols.push(idx);
255 }
256 None => return Value::Error(ErrorKind::Value),
257 }
258 }
259 let result: Vec<Vec<Value>> = grid
260 .iter()
261 .map(|row| {
262 selected_cols
263 .iter()
264 .map(|&c| row.get(c).cloned().unwrap_or(Value::Empty))
265 .collect()
266 })
267 .collect();
268 from_2d(result)
269}
270
271fn chooserows_fn(args: &[Value]) -> Value {
274 if let Some(e) = check_arity(args, 2, usize::MAX) {
275 return e;
276 }
277 let grid = to_2d(&args[0]);
278 let nrows = grid.len();
279 let mut selected_rows: Vec<usize> = Vec::new();
280 for row_arg in &args[1..] {
281 match to_f64(row_arg) {
282 Some(0.0) => return Value::Error(ErrorKind::Value),
283 Some(n) => {
284 let idx = if n < 0.0 {
285 let i = (nrows as isize + n as isize) as usize;
286 if n as isize + (nrows as isize) < 0 {
287 return Value::Error(ErrorKind::Value);
288 }
289 i
290 } else {
291 let i = n as usize - 1;
292 if i >= nrows {
293 return Value::Error(ErrorKind::Value);
294 }
295 i
296 };
297 selected_rows.push(idx);
298 }
299 None => return Value::Error(ErrorKind::Value),
300 }
301 }
302 let result: Vec<Vec<Value>> = selected_rows
303 .iter()
304 .map(|&r| grid.get(r).cloned().unwrap_or_default())
305 .collect();
306 from_2d(result)
307}
308
309pub(crate) fn flatten_fn(args: &[Value]) -> Value {
313 if let Some(e) = check_arity(args, 1, 1) {
314 return e;
315 }
316 let flat = flatten_val(&args[0]);
317 let col: Vec<Vec<Value>> = flat.into_iter().map(|v| vec![v]).collect();
319 from_2d(col)
320}
321
322fn hstack_fn(args: &[Value]) -> Value {
325 if let Some(e) = check_arity(args, 1, usize::MAX) {
326 return e;
327 }
328 let grids: Vec<Vec<Vec<Value>>> = args.iter().map(to_2d).collect();
329 let nrows = grids.iter().map(|g| g.len()).max().unwrap_or(0);
330 let result: Vec<Vec<Value>> = (0..nrows)
331 .map(|r| {
332 grids
333 .iter()
334 .flat_map(|g| {
335 g.get(r).cloned().unwrap_or_default()
336 })
337 .collect()
338 })
339 .collect();
340 from_2d(result)
341}
342
343fn vstack_fn(args: &[Value]) -> Value {
346 if let Some(e) = check_arity(args, 1, usize::MAX) {
347 return e;
348 }
349 let mut result: Vec<Vec<Value>> = Vec::new();
350 for arg in args {
351 let grid = to_2d(arg);
352 result.extend(grid);
353 }
354 from_2d(result)
355}
356
357fn tocol_fn(args: &[Value]) -> Value {
361 if let Some(e) = check_arity(args, 1, 3) {
362 return e;
363 }
364 let flat = flatten_val(&args[0]);
365 let col: Vec<Vec<Value>> = flat.into_iter().map(|v| vec![v]).collect();
366 from_2d(col)
367}
368
369fn torow_fn(args: &[Value]) -> Value {
373 if let Some(e) = check_arity(args, 1, 3) {
374 return e;
375 }
376 let flat = flatten_val(&args[0]);
377 Value::Array(flat)
378}
379
380fn wrapcols_fn(args: &[Value]) -> Value {
385 if let Some(e) = check_arity(args, 2, 3) {
386 return e;
387 }
388 let flat = flatten_val(&args[0]);
389 let wrap_count = match to_f64(&args[1]) {
390 Some(n) if n >= 1.0 => n as usize,
391 _ => return Value::Error(ErrorKind::Value),
392 };
393 let pad = args.get(2).cloned().unwrap_or(Value::Empty);
394
395 let ncols = flat.len().div_ceil(wrap_count);
397 let nrows = wrap_count;
398
399 let grid: Vec<Vec<Value>> = (0..nrows)
401 .map(|r| {
402 (0..ncols)
403 .map(|c| {
404 let idx = c * wrap_count + r;
405 flat.get(idx).cloned().unwrap_or_else(|| pad.clone())
406 })
407 .collect()
408 })
409 .collect();
410 from_2d(grid)
411}
412
413fn wraprows_fn(args: &[Value]) -> Value {
417 if let Some(e) = check_arity(args, 2, 3) {
418 return e;
419 }
420 let flat = flatten_val(&args[0]);
421 let wrap_count = match to_f64(&args[1]) {
422 Some(n) if n >= 1.0 => n as usize,
423 _ => return Value::Error(ErrorKind::Value),
424 };
425 let pad = args.get(2).cloned().unwrap_or(Value::Empty);
426
427 let nrows = flat.len().div_ceil(wrap_count);
428 let grid: Vec<Vec<Value>> = (0..nrows)
429 .map(|r| {
430 (0..wrap_count)
431 .map(|c| {
432 let idx = r * wrap_count + c;
433 flat.get(idx).cloned().unwrap_or_else(|| pad.clone())
434 })
435 .collect()
436 })
437 .collect();
438 from_2d(grid)
439}
440
441pub(crate) fn sort_fn(args: &[Value]) -> Value {
444 if let Some(e) = check_arity(args, 1, 4) {
445 return e;
446 }
447 let mut grid = to_2d(&args[0]);
448 let sort_col = if args.len() >= 2 {
449 match to_f64(&args[1]) {
450 Some(n) => n as usize - 1,
451 None => 0,
452 }
453 } else {
454 0
455 };
456 let ascending = if args.len() >= 3 {
457 match &args[2] {
458 Value::Number(n) => *n >= 0.0,
459 Value::Bool(b) => *b,
460 _ => true,
461 }
462 } else {
463 true
464 };
465
466 grid.sort_by(|a, b| {
467 let va = a.get(sort_col).unwrap_or(&Value::Empty);
468 let vb = b.get(sort_col).unwrap_or(&Value::Empty);
469 let cmp = compare_values_sort(va, vb);
470 if ascending { cmp } else { cmp.reverse() }
471 });
472 from_2d(grid)
473}
474
475fn compare_values_sort(a: &Value, b: &Value) -> std::cmp::Ordering {
476 match (a, b) {
477 (Value::Number(x), Value::Number(y)) => x.partial_cmp(y).unwrap_or(std::cmp::Ordering::Equal),
478 (Value::Text(x), Value::Text(y)) => x.cmp(y),
479 (Value::Bool(x), Value::Bool(y)) => x.cmp(y),
480 _ => std::cmp::Ordering::Equal,
481 }
482}
483
484fn sortby_fn(args: &[Value]) -> Value {
487 if let Some(e) = check_arity(args, 2, usize::MAX) {
488 return e;
489 }
490 let grid = to_2d(&args[0]);
491 let nrows = grid.len();
492
493 let mut sort_keys: Vec<(Vec<Value>, bool)> = Vec::new();
495 let mut i = 1;
496 while i < args.len() {
497 let key_vals = flatten_val(&args[i]);
498 if key_vals.len() != nrows && nrows > 1 {
499 return Value::Error(ErrorKind::Value);
500 }
501 let ascending = if i + 1 < args.len() {
502 match to_f64(&args[i + 1]) {
503 Some(n) => n >= 0.0,
504 None => true,
505 }
506 } else {
507 true
508 };
509 sort_keys.push((key_vals, ascending));
510 i += 2;
511 }
512
513 let mut indices: Vec<usize> = (0..nrows).collect();
514 indices.sort_by(|&ra, &rb| {
515 for (keys, asc) in &sort_keys {
516 let va = keys.get(ra).unwrap_or(&Value::Empty);
517 let vb = keys.get(rb).unwrap_or(&Value::Empty);
518 let cmp = compare_values_sort(va, vb);
519 if cmp != std::cmp::Ordering::Equal {
520 return if *asc { cmp } else { cmp.reverse() };
521 }
522 }
523 std::cmp::Ordering::Equal
524 });
525
526 let sorted: Vec<Vec<Value>> = indices.iter().map(|&r| grid[r].clone()).collect();
527 drop(grid);
528 from_2d(sorted)
529}
530
531pub(crate) fn unique_fn(args: &[Value]) -> Value {
534 if let Some(e) = check_arity(args, 1, 3) {
535 return e;
536 }
537 let grid = to_2d(&args[0]);
538 let by_col = args.get(1).map(|v| matches!(v, Value::Bool(true))).unwrap_or(false);
540 let exactly_once = args.get(2).map(|v| matches!(v, Value::Bool(true))).unwrap_or(false);
541
542 if by_col {
543 let nrows = grid.len();
545 if nrows == 0 {
546 return from_2d(vec![]);
547 }
548 let ncols = grid[0].len();
549 let columns: Vec<Vec<Value>> = (0..ncols)
551 .map(|c| grid.iter().map(|row| row[c].clone()).collect())
552 .collect();
553 let mut seen_cols: Vec<Vec<Value>> = Vec::new();
554 let mut counts: Vec<usize> = Vec::new();
555 for col in columns {
556 if let Some(pos) = seen_cols.iter().position(|sc| sc == &col) {
557 counts[pos] += 1;
558 } else {
559 seen_cols.push(col);
560 counts.push(1);
561 }
562 }
563 let result_cols: Vec<Vec<Value>> = seen_cols
564 .into_iter()
565 .zip(counts)
566 .filter(|(_, cnt)| !exactly_once || *cnt == 1)
567 .map(|(col, _)| col)
568 .collect();
569 let ncols2 = result_cols.len();
571 let result: Vec<Vec<Value>> = (0..nrows)
572 .map(|r| (0..ncols2).map(|c| result_cols[c][r].clone()).collect())
573 .collect();
574 return from_2d(result);
575 }
576
577 let mut seen_rows: Vec<Vec<Value>> = Vec::new();
579 let mut counts: Vec<usize> = Vec::new();
580 for row in &grid {
581 if let Some(pos) = seen_rows.iter().position(|sr| sr == row) {
582 counts[pos] += 1;
583 } else {
584 seen_rows.push(row.clone());
585 counts.push(1);
586 }
587 }
588 let result: Vec<Vec<Value>> = seen_rows
589 .into_iter()
590 .zip(counts)
591 .filter(|(_, cnt)| !exactly_once || *cnt == 1)
592 .map(|(row, _)| row)
593 .collect();
594 from_2d(result)
595}
596
597pub(crate) fn sumproduct_fn(args: &[Value]) -> Value {
600 if let Some(e) = check_arity(args, 1, usize::MAX) {
601 return e;
602 }
603 let arrays: Vec<Vec<Value>> = args.iter().map(flatten_val).collect();
604 let len = arrays[0].len();
605 for arr in &arrays[1..] {
607 if arr.len() != len {
608 return Value::Error(ErrorKind::Value);
609 }
610 }
611 let mut sum = 0.0;
612 for i in 0..len {
613 let mut prod = 1.0;
614 for arr in &arrays {
615 match to_f64(&arr[i]) {
616 Some(n) => prod *= n,
617 None => return Value::Error(ErrorKind::Value),
618 }
619 }
620 sum += prod;
621 }
622 Value::Number(sum)
623}
624
625fn sumxmy2_fn(args: &[Value]) -> Value {
628 if let Some(e) = check_arity(args, 2, 2) {
629 return e;
630 }
631 let xs = flatten_val(&args[0]);
632 let ys = flatten_val(&args[1]);
633 if xs.len() != ys.len() {
634 return Value::Error(ErrorKind::Value);
635 }
636 let mut sum = 0.0;
637 for (x, y) in xs.iter().zip(ys.iter()) {
638 let xn = match to_f64(x) { Some(n) => n, None => return Value::Error(ErrorKind::Value) };
639 let yn = match to_f64(y) { Some(n) => n, None => return Value::Error(ErrorKind::Value) };
640 sum += (xn - yn).powi(2);
641 }
642 Value::Number(sum)
643}
644
645fn sumx2my2_fn(args: &[Value]) -> Value {
648 if let Some(e) = check_arity(args, 2, 2) {
649 return e;
650 }
651 let xs = flatten_val(&args[0]);
652 let ys = flatten_val(&args[1]);
653 if xs.len() != ys.len() {
654 return Value::Error(ErrorKind::Value);
655 }
656 let mut sum = 0.0;
657 for (x, y) in xs.iter().zip(ys.iter()) {
658 let xn = match to_f64(x) { Some(n) => n, None => return Value::Error(ErrorKind::Value) };
659 let yn = match to_f64(y) { Some(n) => n, None => return Value::Error(ErrorKind::Value) };
660 sum += xn * xn - yn * yn;
661 }
662 Value::Number(sum)
663}
664
665fn sumx2py2_fn(args: &[Value]) -> Value {
668 if let Some(e) = check_arity(args, 2, 2) {
669 return e;
670 }
671 let xs = flatten_val(&args[0]);
672 let ys = flatten_val(&args[1]);
673 if xs.len() != ys.len() {
674 return Value::Error(ErrorKind::Value);
675 }
676 let mut sum = 0.0;
677 for (x, y) in xs.iter().zip(ys.iter()) {
678 let xn = match to_f64(x) { Some(n) => n, None => return Value::Error(ErrorKind::Value) };
679 let yn = match to_f64(y) { Some(n) => n, None => return Value::Error(ErrorKind::Value) };
680 sum += xn * xn + yn * yn;
681 }
682 Value::Number(sum)
683}
684
685fn mmult_fn(args: &[Value]) -> Value {
688 if let Some(e) = check_arity(args, 2, 2) {
689 return e;
690 }
691 let a = to_2d(&args[0]);
692 let b = to_2d(&args[1]);
693 let n = a.first().map(|r| r.len()).unwrap_or(0);
694 let p = b.first().map(|r| r.len()).unwrap_or(0);
695 if b.len() != n {
696 return Value::Error(ErrorKind::Value);
697 }
698 let af: Vec<Vec<f64>> = a.iter().map(|row| {
700 row.iter().map(|v| to_f64(v).unwrap_or(f64::NAN)).collect()
701 }).collect();
702 let bf: Vec<Vec<f64>> = b.iter().map(|row| {
703 row.iter().map(|v| to_f64(v).unwrap_or(f64::NAN)).collect()
704 }).collect();
705 if af.iter().any(|r| r.iter().any(|v| v.is_nan())) || bf.iter().any(|r| r.iter().any(|v| v.is_nan())) {
706 return Value::Error(ErrorKind::Value);
707 }
708 let result: Vec<Vec<Value>> = af.iter().map(|row_a| {
709 (0..p).map(|j| {
710 let sum: f64 = row_a.iter().enumerate().map(|(k, &av)| av * bf[k][j]).sum();
711 Value::Number(sum)
712 }).collect()
713 }).collect();
714 from_2d(result)
715}
716
717fn mdeterm_fn(args: &[Value]) -> Value {
720 if let Some(e) = check_arity(args, 1, 1) {
721 return e;
722 }
723 let grid = to_2d(&args[0]);
724 let n = grid.len();
725 if n == 0 {
726 return Value::Error(ErrorKind::Value);
727 }
728 for row in &grid {
729 if row.len() != n {
730 return Value::Error(ErrorKind::Value);
731 }
732 }
733 let mut mat: Vec<Vec<f64>> = Vec::with_capacity(n);
735 for row in &grid {
736 let mut r = Vec::with_capacity(n);
737 for v in row {
738 match to_f64(v) {
739 Some(x) => r.push(x),
740 None => return Value::Error(ErrorKind::Value),
741 }
742 }
743 mat.push(r);
744 }
745 Value::Number(determinant(&mat))
746}
747
748fn determinant(mat: &[Vec<f64>]) -> f64 {
749 let n = mat.len();
750 if n == 1 {
751 return mat[0][0];
752 }
753 if n == 2 {
754 return mat[0][0] * mat[1][1] - mat[0][1] * mat[1][0];
755 }
756 let mut det = 0.0;
757 for c in 0..n {
758 let minor: Vec<Vec<f64>> = (1..n)
759 .map(|r| {
760 (0..n)
761 .filter(|&cc| cc != c)
762 .map(|cc| mat[r][cc])
763 .collect()
764 })
765 .collect();
766 let sign = if c % 2 == 0 { 1.0 } else { -1.0 };
767 det += sign * mat[0][c] * determinant(&minor);
768 }
769 det
770}
771
772fn minverse_fn(args: &[Value]) -> Value {
775 if let Some(e) = check_arity(args, 1, 1) {
776 return e;
777 }
778 let grid = to_2d(&args[0]);
779 let n = grid.len();
780 if n == 0 {
781 return Value::Error(ErrorKind::Value);
782 }
783 for row in &grid {
784 if row.len() != n {
785 return Value::Error(ErrorKind::Value);
786 }
787 }
788 let mut mat: Vec<Vec<f64>> = Vec::with_capacity(n);
789 for row in &grid {
790 let mut r = Vec::with_capacity(n);
791 for v in row {
792 match to_f64(v) {
793 Some(x) => r.push(x),
794 None => return Value::Error(ErrorKind::Value),
795 }
796 }
797 mat.push(r);
798 }
799 match invert_matrix(mat) {
800 Some(inv) => from_2d(inv.into_iter().map(|r| r.into_iter().map(Value::Number).collect()).collect()),
801 None => Value::Error(ErrorKind::Value),
802 }
803}
804
805fn invert_matrix(mut mat: Vec<Vec<f64>>) -> Option<Vec<Vec<f64>>> {
806 let n = mat.len();
807 let mut inv: Vec<Vec<f64>> = (0..n)
809 .map(|i| (0..n).map(|j| if i == j { 1.0 } else { 0.0 }).collect())
810 .collect();
811 for col in 0..n {
812 let pivot = (col..n).max_by(|&a, &b| mat[a][col].abs().partial_cmp(&mat[b][col].abs()).unwrap_or(std::cmp::Ordering::Equal))?;
814 if mat[pivot][col].abs() < 1e-12 {
815 return None; }
817 mat.swap(col, pivot);
818 inv.swap(col, pivot);
819 let div = mat[col][col];
820 for j in 0..n {
821 mat[col][j] /= div;
822 inv[col][j] /= div;
823 }
824 for r in 0..n {
825 if r != col {
826 let factor = mat[r][col];
827 for j in 0..n {
828 mat[r][j] -= factor * mat[col][j];
829 inv[r][j] -= factor * inv[col][j];
830 }
831 }
832 }
833 }
834 Some(inv)
835}
836
837fn frequency_fn(args: &[Value]) -> Value {
840 if let Some(e) = check_arity(args, 2, 2) {
841 return e;
842 }
843 let data = flatten_val(&args[0]);
844 let bins = flatten_val(&args[1]);
845
846 let mut bin_vals: Vec<f64> = bins
847 .iter()
848 .filter_map(to_f64)
849 .collect();
850 bin_vals.sort_by(|a, b| a.partial_cmp(b).unwrap_or(std::cmp::Ordering::Equal));
851
852 let mut counts = vec![0usize; bin_vals.len() + 1];
853 for d in &data {
854 if let Some(n) = to_f64(d) {
855 let bin = bin_vals.partition_point(|&b| b < n);
856 counts[bin] += 1;
857 }
858 }
859 let col: Vec<Vec<Value>> = counts.into_iter().map(|c| vec![Value::Number(c as f64)]).collect();
861 from_2d(col)
862}
863
864fn linest_fn(args: &[Value]) -> Value {
868 if let Some(e) = check_arity(args, 1, 4) {
869 return e;
870 }
871 let ys = flatten_val(&args[0]);
872 let n = ys.len();
873 let xs: Vec<f64> = if args.len() >= 2 {
874 flatten_val(&args[1]).iter().filter_map(to_f64).collect()
875 } else {
876 (1..=n).map(|i| i as f64).collect()
877 };
878 if xs.len() != n || n < 2 {
879 return Value::Error(ErrorKind::Value);
880 }
881 let y_vals: Vec<f64> = ys.iter().filter_map(to_f64).collect();
882 if y_vals.len() != n {
883 return Value::Error(ErrorKind::Value);
884 }
885 let (slope, intercept) = simple_linear_regression(&xs, &y_vals);
886 Value::Array(vec![Value::Number(slope), Value::Number(intercept)])
887}
888
889fn simple_linear_regression(xs: &[f64], ys: &[f64]) -> (f64, f64) {
890 let n = xs.len() as f64;
891 let sum_x: f64 = xs.iter().sum();
892 let sum_y: f64 = ys.iter().sum();
893 let sum_xy: f64 = xs.iter().zip(ys.iter()).map(|(x, y)| x * y).sum();
894 let sum_xx: f64 = xs.iter().map(|x| x * x).sum();
895 let denom = n * sum_xx - sum_x * sum_x;
896 if denom.abs() < 1e-15 {
897 let intercept = sum_y / n;
898 return (0.0, intercept);
899 }
900 let slope = (n * sum_xy - sum_x * sum_y) / denom;
901 let intercept = (sum_y - slope * sum_x) / n;
902 (slope, intercept)
903}
904
905fn logest_fn(args: &[Value]) -> Value {
909 if let Some(e) = check_arity(args, 1, 4) {
910 return e;
911 }
912 let ys = flatten_val(&args[0]);
913 let n = ys.len();
914 let xs: Vec<f64> = if args.len() >= 2 {
915 flatten_val(&args[1]).iter().filter_map(to_f64).collect()
916 } else {
917 (1..=n).map(|i| i as f64).collect()
918 };
919 if xs.len() != n || n < 2 {
920 return Value::Error(ErrorKind::Value);
921 }
922 let y_vals: Vec<f64> = ys.iter().filter_map(to_f64).collect();
923 if y_vals.len() != n {
924 return Value::Error(ErrorKind::Value);
925 }
926 let log_y: Vec<f64> = y_vals.iter().map(|&y| y.ln()).collect();
928 if log_y.iter().any(|v| v.is_nan() || v.is_infinite()) {
929 return Value::Error(ErrorKind::Num);
930 }
931 let (log_base, log_intercept) = simple_linear_regression(&xs, &log_y);
932 let base = log_base.exp();
933 let intercept = log_intercept.exp();
934 Value::Array(vec![Value::Number(base), Value::Number(intercept)])
935}
936
937fn trend_fn(args: &[Value]) -> Value {
941 if let Some(e) = check_arity(args, 1, 4) {
942 return e;
943 }
944 let ys = flatten_val(&args[0]);
945 let n = ys.len();
946 let xs: Vec<f64> = if args.len() >= 2 {
947 flatten_val(&args[1]).iter().filter_map(to_f64).collect()
948 } else {
949 (1..=n).map(|i| i as f64).collect()
950 };
951 if xs.len() != n || n < 2 {
952 return Value::Error(ErrorKind::Value);
953 }
954 let y_vals: Vec<f64> = ys.iter().filter_map(to_f64).collect();
955 if y_vals.len() != n {
956 return Value::Error(ErrorKind::Value);
957 }
958 let new_xs: Vec<f64> = if args.len() >= 3 {
959 flatten_val(&args[2]).iter().filter_map(to_f64).collect()
960 } else {
961 xs.clone()
962 };
963 let (slope, intercept) = simple_linear_regression(&xs, &y_vals);
964 let result: Vec<Value> = new_xs.iter().map(|&x| Value::Number(slope * x + intercept)).collect();
965 Value::Array(result)
966}
967
968fn growth_fn(args: &[Value]) -> Value {
972 if let Some(e) = check_arity(args, 1, 4) {
973 return e;
974 }
975 let ys = flatten_val(&args[0]);
976 let n = ys.len();
977 let xs: Vec<f64> = if args.len() >= 2 {
978 flatten_val(&args[1]).iter().filter_map(to_f64).collect()
979 } else {
980 (1..=n).map(|i| i as f64).collect()
981 };
982 if xs.len() != n || n < 2 {
983 return Value::Error(ErrorKind::Value);
984 }
985 let y_vals: Vec<f64> = ys.iter().filter_map(to_f64).collect();
986 if y_vals.len() != n {
987 return Value::Error(ErrorKind::Value);
988 }
989 let log_y: Vec<f64> = y_vals.iter().map(|&y| y.ln()).collect();
990 if log_y.iter().any(|v| v.is_nan() || v.is_infinite()) {
991 return Value::Error(ErrorKind::Num);
992 }
993 let new_xs: Vec<f64> = if args.len() >= 3 && !matches!(args[2], Value::Empty) {
994 let vals: Vec<f64> = flatten_val(&args[2]).iter().filter_map(to_f64).collect();
995 if vals.is_empty() { xs.clone() } else { vals }
996 } else {
997 xs.clone()
998 };
999 if args.len() >= 4 {
1001 let b_false = match &args[3] {
1002 Value::Bool(b) => !b,
1003 Value::Number(n) => *n == 0.0,
1004 _ => false,
1005 };
1006 if b_false {
1007 return Value::Error(ErrorKind::Value);
1008 }
1009 }
1010 let (log_base, log_intercept) = simple_linear_regression(&xs, &log_y);
1011 let result: Vec<Value> = new_xs
1012 .iter()
1013 .map(|&x| Value::Number((log_base * x + log_intercept).exp()))
1014 .collect();
1015 Value::Array(result)
1016}
1017
1018fn apply_lambda(lambda_expr: &Expr, bound_args: &[Value], ctx: &mut EvalCtx<'_>) -> Option<Value> {
1024 match lambda_expr {
1025 Expr::FunctionCall { name, args, .. } if name == "LAMBDA" => {
1026 if args.is_empty() {
1027 return None;
1028 }
1029 let body = &args[args.len() - 1];
1030 let params = &args[..args.len() - 1];
1031 if params.len() != bound_args.len() {
1032 return None;
1033 }
1034 let mut saved: Vec<(String, Value)> = Vec::new();
1036 for (param_expr, val) in params.iter().zip(bound_args.iter()) {
1037 if let Expr::Variable(name, _) = param_expr {
1038 let old = ctx.ctx.get(name);
1039 saved.push((name.clone(), old));
1040 ctx.ctx.set(name.clone(), val.clone());
1041 } else {
1042 return None;
1043 }
1044 }
1045 let result = evaluate_expr(body, ctx);
1046 for (name, old_val) in saved {
1048 ctx.ctx.set(name, old_val);
1049 }
1050 Some(result)
1051 }
1052 _ => None,
1053 }
1054}
1055
1056pub fn byrow_lazy_fn(args: &[Expr], ctx: &mut EvalCtx<'_>) -> Value {
1059 if let Some(e) = check_arity_len(args.len(), 2, 2) {
1060 return e;
1061 }
1062 let arr_val = evaluate_expr(&args[0], ctx);
1063 if matches!(arr_val, Value::Error(_)) {
1064 return arr_val;
1065 }
1066 let grid = to_2d(&arr_val);
1067 let lambda_expr = &args[1];
1068 let mut results: Vec<Value> = Vec::with_capacity(grid.len());
1069 for row in &grid {
1070 let row_val = Value::Array(row.clone());
1071 match apply_lambda(lambda_expr, &[row_val], ctx) {
1072 Some(v) => results.push(v),
1073 None => return Value::Error(ErrorKind::Value),
1074 }
1075 }
1076 let col: Vec<Vec<Value>> = results.into_iter().map(|v| vec![v]).collect();
1078 from_2d(col)
1079}
1080
1081pub fn bycol_lazy_fn(args: &[Expr], ctx: &mut EvalCtx<'_>) -> Value {
1084 if let Some(e) = check_arity_len(args.len(), 2, 2) {
1085 return e;
1086 }
1087 let arr_val = evaluate_expr(&args[0], ctx);
1088 if matches!(arr_val, Value::Error(_)) {
1089 return arr_val;
1090 }
1091 let grid = to_2d(&arr_val);
1092 let ncols = grid.first().map(|r| r.len()).unwrap_or(0);
1093 let columns: Vec<Vec<Value>> = (0..ncols)
1095 .map(|c| grid.iter().map(|row| row[c].clone()).collect())
1096 .collect();
1097 let lambda_expr = &args[1];
1098 let mut results: Vec<Value> = Vec::with_capacity(ncols);
1099 for col in columns {
1100 let col_val = Value::Array(col);
1102 match apply_lambda(lambda_expr, &[col_val], ctx) {
1103 Some(v) => results.push(v),
1104 None => return Value::Error(ErrorKind::Value),
1105 }
1106 }
1107 Value::Array(results)
1109}
1110
1111pub fn map_lazy_fn(args: &[Expr], ctx: &mut EvalCtx<'_>) -> Value {
1114 if let Some(e) = check_arity_len(args.len(), 2, usize::MAX) {
1115 return e;
1116 }
1117 let lambda_expr = &args[args.len() - 1];
1119 let arr_count = args.len() - 1;
1120 let arrays: Vec<Vec<Value>> = args[..arr_count]
1121 .iter()
1122 .map(|a| {
1123 let v = evaluate_expr(a, ctx);
1124 flatten_val(&v)
1125 })
1126 .collect();
1127 let len = arrays[0].len();
1128 for arr in &arrays[1..] {
1129 if arr.len() != len {
1130 return Value::Error(ErrorKind::Value);
1131 }
1132 }
1133 let mut results: Vec<Value> = Vec::with_capacity(len);
1134 for i in 0..len {
1135 let bound: Vec<Value> = arrays.iter().map(|a| a[i].clone()).collect();
1136 match apply_lambda(lambda_expr, &bound, ctx) {
1137 Some(v) => results.push(v),
1138 None => return Value::Error(ErrorKind::Value),
1139 }
1140 }
1141 let first_grid = to_2d(&evaluate_expr(&args[0], ctx));
1143 if first_grid.len() > 1 {
1144 let ncols = first_grid[0].len();
1146 let nrows = first_grid.len();
1147 let grid: Vec<Vec<Value>> = (0..nrows)
1148 .map(|r| (0..ncols).map(|c| results[r * ncols + c].clone()).collect())
1149 .collect();
1150 from_2d(grid)
1151 } else {
1152 Value::Array(results)
1153 }
1154}
1155
1156pub fn reduce_lazy_fn(args: &[Expr], ctx: &mut EvalCtx<'_>) -> Value {
1159 if let Some(e) = check_arity_len(args.len(), 3, 3) {
1160 return e;
1161 }
1162 let initial = evaluate_expr(&args[0], ctx);
1163 if matches!(initial, Value::Error(_)) {
1164 return initial;
1165 }
1166 let arr_val = evaluate_expr(&args[1], ctx);
1167 if matches!(arr_val, Value::Error(_)) {
1168 return arr_val;
1169 }
1170 let items = flatten_val(&arr_val);
1171 let lambda_expr = &args[2];
1172 let mut acc = initial;
1173 for item in &items {
1174 match apply_lambda(lambda_expr, &[acc.clone(), item.clone()], ctx) {
1175 Some(v) => acc = v,
1176 None => return Value::Error(ErrorKind::Value),
1177 }
1178 }
1179 acc
1180}
1181
1182pub fn scan_lazy_fn(args: &[Expr], ctx: &mut EvalCtx<'_>) -> Value {
1185 if let Some(e) = check_arity_len(args.len(), 3, 3) {
1186 return e;
1187 }
1188 let initial = evaluate_expr(&args[0], ctx);
1189 if matches!(initial, Value::Error(_)) {
1190 return initial;
1191 }
1192 let arr_val = evaluate_expr(&args[1], ctx);
1193 if matches!(arr_val, Value::Error(_)) {
1194 return arr_val;
1195 }
1196 let grid = to_2d(&arr_val);
1197 let items = flatten_val(&arr_val);
1198 let lambda_expr = &args[2];
1199 let mut acc = initial;
1200 let mut results: Vec<Value> = Vec::with_capacity(items.len());
1201 for item in &items {
1202 match apply_lambda(lambda_expr, &[acc.clone(), item.clone()], ctx) {
1203 Some(v) => {
1204 acc = v.clone();
1205 results.push(v);
1206 }
1207 None => return Value::Error(ErrorKind::Value),
1208 }
1209 }
1210 if grid.len() > 1 {
1212 let ncols = grid[0].len();
1213 let nrows = grid.len();
1214 let result_grid: Vec<Vec<Value>> = (0..nrows)
1215 .map(|r| (0..ncols).map(|c| results[r * ncols + c].clone()).collect())
1216 .collect();
1217 from_2d(result_grid)
1218 } else {
1219 Value::Array(results)
1220 }
1221}
1222
1223pub fn makearray_lazy_fn(args: &[Expr], ctx: &mut EvalCtx<'_>) -> Value {
1226 if let Some(e) = check_arity_len(args.len(), 3, 3) {
1227 return e;
1228 }
1229 let rows_val = evaluate_expr(&args[0], ctx);
1230 let cols_val = evaluate_expr(&args[1], ctx);
1231 if matches!(rows_val, Value::Error(_)) {
1232 return rows_val;
1233 }
1234 if matches!(cols_val, Value::Error(_)) {
1235 return cols_val;
1236 }
1237 let nrows = match to_f64(&rows_val) {
1238 Some(n) if n >= 1.0 => n as usize,
1239 _ => return Value::Error(ErrorKind::Value),
1240 };
1241 let ncols = match to_f64(&cols_val) {
1242 Some(n) if n >= 1.0 => n as usize,
1243 _ => return Value::Error(ErrorKind::Value),
1244 };
1245 let lambda_expr = &args[2];
1246 let mut grid: Vec<Vec<Value>> = Vec::with_capacity(nrows);
1247 for r in 1..=nrows {
1248 let mut row = Vec::with_capacity(ncols);
1249 for c in 1..=ncols {
1250 let rv = Value::Number(r as f64);
1251 let cv = Value::Number(c as f64);
1252 match apply_lambda(lambda_expr, &[rv, cv], ctx) {
1253 Some(v) => row.push(v),
1254 None => return Value::Error(ErrorKind::Value),
1255 }
1256 }
1257 grid.push(row);
1258 }
1259 from_2d(grid)
1260}
1261
1262pub fn register_array(registry: &mut Registry) {
1265 registry.register_eager("ROWS", rows_fn, FunctionMeta {
1266 category: "array",
1267 signature: "ROWS(array)",
1268 description: "Returns the number of rows in an array or range",
1269 });
1270 registry.register_eager("COLUMNS", columns_fn, FunctionMeta {
1271 category: "array",
1272 signature: "COLUMNS(array)",
1273 description: "Returns the number of columns in an array or range",
1274 });
1275 registry.register_eager("INDEX", index_fn, FunctionMeta {
1276 category: "array",
1277 signature: "INDEX(array, row, [col])",
1278 description: "Returns the value at the given row and column of an array",
1279 });
1280 registry.register_eager("TRANSPOSE", transpose_fn, FunctionMeta {
1281 category: "array",
1282 signature: "TRANSPOSE(array)",
1283 description: "Transposes the rows and columns of an array",
1284 });
1285 registry.register_eager("ARRAY_CONSTRAIN", array_constrain_fn, FunctionMeta {
1286 category: "array",
1287 signature: "ARRAY_CONSTRAIN(input, num_rows, num_cols)",
1288 description: "Constrains an array to a given number of rows and columns",
1289 });
1290 registry.register_eager("CHOOSECOLS", choosecols_fn, FunctionMeta {
1291 category: "array",
1292 signature: "CHOOSECOLS(array, col_num1, ...)",
1293 description: "Returns selected columns from an array",
1294 });
1295 registry.register_eager("CHOOSEROWS", chooserows_fn, FunctionMeta {
1296 category: "array",
1297 signature: "CHOOSEROWS(array, row_num1, ...)",
1298 description: "Returns selected rows from an array",
1299 });
1300 registry.register_eager("FLATTEN", flatten_fn, FunctionMeta {
1301 category: "array",
1302 signature: "FLATTEN(array)",
1303 description: "Flattens an array into a single column",
1304 });
1305 registry.register_eager("HSTACK", hstack_fn, FunctionMeta {
1306 category: "array",
1307 signature: "HSTACK(array1, ...)",
1308 description: "Horizontally stacks arrays",
1309 });
1310 registry.register_eager("VSTACK", vstack_fn, FunctionMeta {
1311 category: "array",
1312 signature: "VSTACK(array1, ...)",
1313 description: "Vertically stacks arrays",
1314 });
1315 registry.register_eager("TOCOL", tocol_fn, FunctionMeta {
1316 category: "array",
1317 signature: "TOCOL(array, [ignore], [scan_by_col])",
1318 description: "Converts an array to a single column",
1319 });
1320 registry.register_eager("TOROW", torow_fn, FunctionMeta {
1321 category: "array",
1322 signature: "TOROW(array, [ignore], [scan_by_col])",
1323 description: "Converts an array to a single row",
1324 });
1325 registry.register_eager("WRAPCOLS", wrapcols_fn, FunctionMeta {
1326 category: "array",
1327 signature: "WRAPCOLS(vector, wrap_count, [pad_with])",
1328 description: "Wraps a vector into columns of the given length",
1329 });
1330 registry.register_eager("WRAPROWS", wraprows_fn, FunctionMeta {
1331 category: "array",
1332 signature: "WRAPROWS(vector, wrap_count, [pad_with])",
1333 description: "Wraps a vector into rows of the given length",
1334 });
1335 registry.register_eager("SORT", sort_fn, FunctionMeta {
1336 category: "array",
1337 signature: "SORT(array, [sort_index], [sort_order], [by_col])",
1338 description: "Sorts an array",
1339 });
1340 registry.register_eager("SORTBY", sortby_fn, FunctionMeta {
1341 category: "array",
1342 signature: "SORTBY(array, by_array1, [sort_order1], ...)",
1343 description: "Sorts an array based on the values in corresponding arrays",
1344 });
1345 registry.register_eager("UNIQUE", unique_fn, FunctionMeta {
1346 category: "array",
1347 signature: "UNIQUE(array, [by_col], [exactly_once])",
1348 description: "Returns unique rows or columns from an array",
1349 });
1350 registry.register_eager("SUMPRODUCT", sumproduct_fn, FunctionMeta {
1351 category: "array",
1352 signature: "SUMPRODUCT(array1, [array2], ...)",
1353 description: "Returns the sum of products of corresponding elements",
1354 });
1355 registry.register_eager("SUMXMY2", sumxmy2_fn, FunctionMeta {
1356 category: "array",
1357 signature: "SUMXMY2(array_x, array_y)",
1358 description: "Returns sum of squares of differences",
1359 });
1360 registry.register_eager("SUMX2MY2", sumx2my2_fn, FunctionMeta {
1361 category: "array",
1362 signature: "SUMX2MY2(array_x, array_y)",
1363 description: "Returns sum of (x^2 - y^2)",
1364 });
1365 registry.register_eager("SUMX2PY2", sumx2py2_fn, FunctionMeta {
1366 category: "array",
1367 signature: "SUMX2PY2(array_x, array_y)",
1368 description: "Returns sum of (x^2 + y^2)",
1369 });
1370 registry.register_eager("MMULT", mmult_fn, FunctionMeta {
1371 category: "array",
1372 signature: "MMULT(array1, array2)",
1373 description: "Returns the matrix product of two arrays",
1374 });
1375 registry.register_eager("MDETERM", mdeterm_fn, FunctionMeta {
1376 category: "array",
1377 signature: "MDETERM(array)",
1378 description: "Returns the matrix determinant",
1379 });
1380 registry.register_eager("MINVERSE", minverse_fn, FunctionMeta {
1381 category: "array",
1382 signature: "MINVERSE(array)",
1383 description: "Returns the matrix inverse",
1384 });
1385 registry.register_eager("FREQUENCY", frequency_fn, FunctionMeta {
1386 category: "array",
1387 signature: "FREQUENCY(data, bins)",
1388 description: "Calculates the frequency distribution of values",
1389 });
1390 registry.register_eager("LINEST", linest_fn, FunctionMeta {
1391 category: "array",
1392 signature: "LINEST(known_y, [known_x], [const], [stats])",
1393 description: "Returns linear regression statistics",
1394 });
1395 registry.register_eager("LOGEST", logest_fn, FunctionMeta {
1396 category: "array",
1397 signature: "LOGEST(known_y, [known_x], [const], [stats])",
1398 description: "Returns exponential regression statistics",
1399 });
1400 registry.register_eager("TREND", trend_fn, FunctionMeta {
1401 category: "array",
1402 signature: "TREND(known_y, [known_x], [new_x], [const])",
1403 description: "Returns values along a linear trend",
1404 });
1405 registry.register_eager("GROWTH", growth_fn, FunctionMeta {
1406 category: "array",
1407 signature: "GROWTH(known_y, [known_x], [new_x], [const])",
1408 description: "Returns values along an exponential trend",
1409 });
1410 registry.register_lazy("BYROW", byrow_lazy_fn, FunctionMeta {
1411 category: "array",
1412 signature: "BYROW(array, lambda)",
1413 description: "Applies a LAMBDA to each row of an array",
1414 });
1415 registry.register_lazy("BYCOL", bycol_lazy_fn, FunctionMeta {
1416 category: "array",
1417 signature: "BYCOL(array, lambda)",
1418 description: "Applies a LAMBDA to each column of an array",
1419 });
1420 registry.register_lazy("MAP", map_lazy_fn, FunctionMeta {
1421 category: "array",
1422 signature: "MAP(array1, [array2, ...], lambda)",
1423 description: "Maps a LAMBDA over one or more arrays",
1424 });
1425 registry.register_lazy("REDUCE", reduce_lazy_fn, FunctionMeta {
1426 category: "array",
1427 signature: "REDUCE(initial_value, array, lambda)",
1428 description: "Reduces an array to a single value using a LAMBDA",
1429 });
1430 registry.register_lazy("SCAN", scan_lazy_fn, FunctionMeta {
1431 category: "array",
1432 signature: "SCAN(initial_value, array, lambda)",
1433 description: "Returns running accumulation using a LAMBDA",
1434 });
1435 registry.register_lazy("MAKEARRAY", makearray_lazy_fn, FunctionMeta {
1436 category: "array",
1437 signature: "MAKEARRAY(rows, cols, lambda)",
1438 description: "Creates an array using a LAMBDA for each cell value",
1439 });
1440}
1441
1442#[cfg(test)]
1443mod tests;