1use crate::args::{ArgSchema, CoercionPolicy, ShapeKind};
10use crate::builtins::utils::collapse_if_scalar;
11use crate::function::Function;
12use crate::traits::{ArgumentHandle, FunctionContext};
13use formualizer_common::{ArgKind, ExcelError, ExcelErrorKind, LiteralValue};
14use formualizer_macros::func_caps;
15
16#[derive(Debug)]
17pub struct ChooseFn;
18#[derive(Debug)]
19pub struct ChooseColsFn;
20#[derive(Debug)]
21pub struct ChooseRowsFn;
22
23impl Function for ChooseFn {
68 fn name(&self) -> &'static str {
69 "CHOOSE"
70 }
71
72 fn min_args(&self) -> usize {
73 2
74 }
75
76 func_caps!(PURE, LOOKUP);
77
78 fn arg_schema(&self) -> &'static [ArgSchema] {
79 use once_cell::sync::Lazy;
80 static SCHEMA: Lazy<Vec<ArgSchema>> = Lazy::new(|| {
81 vec![
82 ArgSchema {
84 kinds: smallvec::smallvec![ArgKind::Number],
85 required: true,
86 by_ref: false,
87 shape: ShapeKind::Scalar,
88 coercion: CoercionPolicy::NumberStrict,
89 max: None,
90 repeating: None,
91 default: None,
92 },
93 ArgSchema {
95 kinds: smallvec::smallvec![ArgKind::Any],
96 required: true,
97 by_ref: false, shape: ShapeKind::Scalar, coercion: CoercionPolicy::None,
100 max: None,
101 repeating: Some(1), default: None,
103 },
104 ]
105 });
106 &SCHEMA
107 }
108
109 fn eval<'a, 'b, 'c>(
110 &self,
111 args: &'c [ArgumentHandle<'a, 'b>],
112 _ctx: &dyn FunctionContext<'b>,
113 ) -> Result<crate::traits::CalcValue<'b>, ExcelError> {
114 if args.len() < 2 {
115 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
116 ExcelError::new(ExcelErrorKind::Value),
117 )));
118 }
119
120 let index_val = args[0].value()?.into_literal();
122 if let LiteralValue::Error(e) = index_val {
123 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(e)));
124 }
125
126 let index = match index_val {
127 LiteralValue::Number(n) => n as i64,
128 LiteralValue::Int(i) => i,
129 LiteralValue::Text(s) => s.parse::<f64>().map(|n| n as i64).unwrap_or(-1),
130 _ => {
131 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
132 ExcelError::new(ExcelErrorKind::Value),
133 )));
134 }
135 };
136
137 if index < 1 || index as usize > args.len() - 1 {
139 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
140 ExcelError::new(ExcelErrorKind::Value),
141 )));
142 }
143
144 let selected_arg = &args[index as usize];
146 selected_arg.value()
147 }
148}
149
150fn materialize_rows_2d<'b>(
153 arg: &ArgumentHandle<'_, 'b>,
154 ctx: &dyn FunctionContext<'b>,
155) -> Result<Vec<Vec<formualizer_common::LiteralValue>>, ExcelError> {
156 if let Ok(r) = arg.as_reference_or_eval() {
157 let mut rows: Vec<Vec<LiteralValue>> = Vec::new();
158 let sheet = ctx.current_sheet();
159 let rv = ctx.resolve_range_view(&r, sheet)?;
160 rv.for_each_row(&mut |row| {
161 rows.push(row.to_vec());
162 Ok(())
163 })?;
164 Ok(rows)
165 } else {
166 let v = arg.value()?.into_literal();
167 match v {
168 LiteralValue::Array(a) => Ok(a),
169 other => Ok(vec![vec![other]]),
170 }
171 }
172}
173
174impl Function for ChooseColsFn {
231 func_caps!(PURE, LOOKUP);
232 fn name(&self) -> &'static str {
233 "CHOOSECOLS"
234 }
235 fn min_args(&self) -> usize {
236 2
237 }
238 fn variadic(&self) -> bool {
239 true
240 }
241 fn arg_schema(&self) -> &'static [ArgSchema] {
242 use once_cell::sync::Lazy;
243 static SCHEMA: Lazy<Vec<ArgSchema>> = Lazy::new(|| {
244 vec![
245 ArgSchema {
247 kinds: smallvec::smallvec![ArgKind::Range, ArgKind::Any],
248 required: true,
249 by_ref: false,
250 shape: ShapeKind::Range,
251 coercion: CoercionPolicy::None,
252 max: None,
253 repeating: None,
254 default: None,
255 },
256 ArgSchema {
258 kinds: smallvec::smallvec![ArgKind::Number],
259 required: true,
260 by_ref: false,
261 shape: ShapeKind::Scalar,
262 coercion: CoercionPolicy::NumberLenientText,
263 max: None,
264 repeating: Some(1),
265 default: None,
266 },
267 ]
268 });
269 &SCHEMA
270 }
271 fn eval<'a, 'b, 'c>(
272 &self,
273 args: &'c [ArgumentHandle<'a, 'b>],
274 _ctx: &dyn FunctionContext<'b>,
275 ) -> Result<crate::traits::CalcValue<'b>, ExcelError> {
276 if args.len() < 2 {
277 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
278 ExcelError::new(ExcelErrorKind::Value),
279 )));
280 }
281 let view = args[0].range_view()?;
282 let (rows, cols) = view.dims();
283 if rows == 0 || cols == 0 {
284 return Ok(crate::traits::CalcValue::Range(
285 crate::engine::range_view::RangeView::from_owned_rows(vec![], _ctx.date_system()),
286 ));
287 }
288
289 let mut indices: Vec<usize> = Vec::new();
290 for a in &args[1..] {
291 let v = a.value()?.into_literal();
292 if let LiteralValue::Error(e) = v {
293 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(e)));
294 }
295 let raw = match v {
296 LiteralValue::Int(i) => i,
297 LiteralValue::Number(n) => n as i64,
298 _ => {
299 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
300 ExcelError::new(ExcelErrorKind::Value),
301 )));
302 }
303 };
304 if raw == 0 {
305 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
306 ExcelError::new(ExcelErrorKind::Value),
307 )));
308 }
309 let adj = if raw > 0 {
310 raw - 1
311 } else {
312 (cols as i64) + raw
313 };
314 if adj < 0 || adj as usize >= cols {
315 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
316 ExcelError::new(ExcelErrorKind::Value),
317 )));
318 }
319 indices.push(adj as usize);
320 }
321 let mut out: Vec<Vec<LiteralValue>> = Vec::with_capacity(rows);
322 for r in 0..rows {
323 let mut new_row = Vec::with_capacity(indices.len());
324 for &c in &indices {
325 new_row.push(view.get_cell(r, c));
326 }
327 out.push(new_row);
328 }
329
330 Ok(collapse_if_scalar(out, _ctx.date_system()))
331 }
332}
333
334impl Function for ChooseRowsFn {
387 func_caps!(PURE, LOOKUP);
388 fn name(&self) -> &'static str {
389 "CHOOSEROWS"
390 }
391 fn min_args(&self) -> usize {
392 2
393 }
394 fn variadic(&self) -> bool {
395 true
396 }
397 fn arg_schema(&self) -> &'static [ArgSchema] {
398 use once_cell::sync::Lazy;
399 static SCHEMA: Lazy<Vec<ArgSchema>> = Lazy::new(|| {
400 vec![
401 ArgSchema {
403 kinds: smallvec::smallvec![ArgKind::Range, ArgKind::Any],
404 required: true,
405 by_ref: false,
406 shape: ShapeKind::Range,
407 coercion: CoercionPolicy::None,
408 max: None,
409 repeating: None,
410 default: None,
411 },
412 ArgSchema {
414 kinds: smallvec::smallvec![ArgKind::Number],
415 required: true,
416 by_ref: false,
417 shape: ShapeKind::Scalar,
418 coercion: CoercionPolicy::NumberLenientText,
419 max: None,
420 repeating: Some(1),
421 default: None,
422 },
423 ]
424 });
425 &SCHEMA
426 }
427 fn eval<'a, 'b, 'c>(
428 &self,
429 args: &'c [ArgumentHandle<'a, 'b>],
430 _ctx: &dyn FunctionContext<'b>,
431 ) -> Result<crate::traits::CalcValue<'b>, ExcelError> {
432 if args.len() < 2 {
433 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
434 ExcelError::new(ExcelErrorKind::Value),
435 )));
436 }
437 let view = args[0].range_view()?;
438 let (rows, cols) = view.dims();
439 if rows == 0 || cols == 0 {
440 return Ok(crate::traits::CalcValue::Range(
441 crate::engine::range_view::RangeView::from_owned_rows(vec![], _ctx.date_system()),
442 ));
443 }
444
445 let mut indices: Vec<usize> = Vec::new();
446 for a in &args[1..] {
447 let v = a.value()?.into_literal();
448 if let LiteralValue::Error(e) = v {
449 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(e)));
450 }
451 let raw = match v {
452 LiteralValue::Int(i) => i,
453 LiteralValue::Number(n) => n as i64,
454 _ => {
455 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
456 ExcelError::new(ExcelErrorKind::Value),
457 )));
458 }
459 };
460 if raw == 0 {
461 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
462 ExcelError::new(ExcelErrorKind::Value),
463 )));
464 }
465 let adj = if raw > 0 {
466 raw - 1
467 } else {
468 (rows as i64) + raw
469 };
470 if adj < 0 || adj as usize >= rows {
471 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
472 ExcelError::new(ExcelErrorKind::Value),
473 )));
474 }
475 indices.push(adj as usize);
476 }
477 let mut out: Vec<Vec<LiteralValue>> = Vec::with_capacity(indices.len());
478 for &r in &indices {
479 let mut row_vals = Vec::with_capacity(cols);
480 for c in 0..cols {
481 row_vals.push(view.get_cell(r, c));
482 }
483 out.push(row_vals);
484 }
485
486 Ok(collapse_if_scalar(out, _ctx.date_system()))
487 }
488}
489
490#[cfg(test)]
491mod tests {
492 use super::*;
493 use crate::test_workbook::TestWorkbook;
494 use formualizer_parse::parser::{ASTNode, ASTNodeType, ReferenceType};
495 use std::sync::Arc;
496
497 fn lit(v: LiteralValue) -> ASTNode {
498 ASTNode::new(ASTNodeType::Literal(v), None)
499 }
500
501 #[test]
502 fn choose_basic() {
503 let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
504 let ctx = wb.interpreter();
505 let f = ctx.context.get_function("", "CHOOSE").unwrap();
506
507 let two = lit(LiteralValue::Int(2));
509 let a = lit(LiteralValue::Text("A".into()));
510 let b = lit(LiteralValue::Text("B".into()));
511 let c = lit(LiteralValue::Text("C".into()));
512
513 let args = vec![
514 ArgumentHandle::new(&two, &ctx),
515 ArgumentHandle::new(&a, &ctx),
516 ArgumentHandle::new(&b, &ctx),
517 ArgumentHandle::new(&c, &ctx),
518 ];
519
520 let result = f
521 .dispatch(&args, &ctx.function_context(None))
522 .unwrap()
523 .into_literal();
524 assert_eq!(result, LiteralValue::Text("B".into()));
525 }
526
527 #[test]
528 fn choose_numeric_values() {
529 let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
530 let ctx = wb.interpreter();
531 let f = ctx.context.get_function("", "CHOOSE").unwrap();
532
533 let three = lit(LiteralValue::Int(3));
535 let ten = lit(LiteralValue::Int(10));
536 let twenty = lit(LiteralValue::Int(20));
537 let thirty = lit(LiteralValue::Int(30));
538 let forty = lit(LiteralValue::Int(40));
539
540 let args = vec![
541 ArgumentHandle::new(&three, &ctx),
542 ArgumentHandle::new(&ten, &ctx),
543 ArgumentHandle::new(&twenty, &ctx),
544 ArgumentHandle::new(&thirty, &ctx),
545 ArgumentHandle::new(&forty, &ctx),
546 ];
547
548 let result = f
549 .dispatch(&args, &ctx.function_context(None))
550 .unwrap()
551 .into_literal();
552 assert_eq!(result, LiteralValue::Int(30));
553 }
554
555 #[test]
556 fn choose_out_of_range() {
557 let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
558 let ctx = wb.interpreter();
559 let f = ctx.context.get_function("", "CHOOSE").unwrap();
560
561 let five = lit(LiteralValue::Int(5));
563 let a = lit(LiteralValue::Text("A".into()));
564 let b = lit(LiteralValue::Text("B".into()));
565 let c = lit(LiteralValue::Text("C".into()));
566
567 let args = vec![
568 ArgumentHandle::new(&five, &ctx),
569 ArgumentHandle::new(&a, &ctx),
570 ArgumentHandle::new(&b, &ctx),
571 ArgumentHandle::new(&c, &ctx),
572 ];
573
574 let result = f
575 .dispatch(&args, &ctx.function_context(None))
576 .unwrap()
577 .into_literal();
578 assert!(matches!(result, LiteralValue::Error(e) if e.kind == ExcelErrorKind::Value));
579
580 let zero = lit(LiteralValue::Int(0));
582 let args2 = vec![
583 ArgumentHandle::new(&zero, &ctx),
584 ArgumentHandle::new(&a, &ctx),
585 ArgumentHandle::new(&b, &ctx),
586 ];
587
588 let result2 = f
589 .dispatch(&args2, &ctx.function_context(None))
590 .unwrap()
591 .into_literal();
592 assert!(matches!(result2, LiteralValue::Error(e) if e.kind == ExcelErrorKind::Value));
593 }
594
595 #[test]
596 fn choose_decimal_index() {
597 let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
598 let ctx = wb.interpreter();
599 let f = ctx.context.get_function("", "CHOOSE").unwrap();
600
601 let two_seven = lit(LiteralValue::Number(2.7));
603 let a = lit(LiteralValue::Text("A".into()));
604 let b = lit(LiteralValue::Text("B".into()));
605 let c = lit(LiteralValue::Text("C".into()));
606
607 let args = vec![
608 ArgumentHandle::new(&two_seven, &ctx),
609 ArgumentHandle::new(&a, &ctx),
610 ArgumentHandle::new(&b, &ctx),
611 ArgumentHandle::new(&c, &ctx),
612 ];
613
614 let result = f
615 .dispatch(&args, &ctx.function_context(None))
616 .unwrap()
617 .into_literal();
618 assert_eq!(result, LiteralValue::Text("B".into()));
619 }
620
621 #[test]
622 fn choose_text_index_numeric_string() {
623 let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
624 let ctx = wb.interpreter();
625 let f = ctx.context.get_function("", "CHOOSE").unwrap();
626 let two_txt = lit(LiteralValue::Text("2".into()));
627 let a = lit(LiteralValue::Text("A".into()));
628 let b = lit(LiteralValue::Text("B".into()));
629 let c = lit(LiteralValue::Text("C".into()));
630 let args = vec![
631 ArgumentHandle::new(&two_txt, &ctx),
632 ArgumentHandle::new(&a, &ctx),
633 ArgumentHandle::new(&b, &ctx),
634 ArgumentHandle::new(&c, &ctx),
635 ];
636 let result = f
637 .dispatch(&args, &ctx.function_context(None))
638 .unwrap()
639 .into_literal();
640 assert!(matches!(result, LiteralValue::Error(e) if e.kind == ExcelErrorKind::Value));
642 }
643
644 #[test]
645 fn choose_decimal_less_than_one() {
646 let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
647 let ctx = wb.interpreter();
648 let f = ctx.context.get_function("", "CHOOSE").unwrap();
649 let zero_nine = lit(LiteralValue::Number(0.9));
650 let a = lit(LiteralValue::Text("A".into()));
651 let b = lit(LiteralValue::Text("B".into()));
652 let args = vec![
653 ArgumentHandle::new(&zero_nine, &ctx),
654 ArgumentHandle::new(&a, &ctx),
655 ArgumentHandle::new(&b, &ctx),
656 ];
657 let result = f
658 .dispatch(&args, &ctx.function_context(None))
659 .unwrap()
660 .into_literal();
661 assert!(matches!(result, LiteralValue::Error(e) if e.kind == ExcelErrorKind::Value));
662 }
663
664 fn range(r: &str, sr: u32, sc: u32, er: u32, ec: u32) -> ASTNode {
665 ASTNode::new(
666 ASTNodeType::Reference {
667 original: r.into(),
668 reference: ReferenceType::range(None, Some(sr), Some(sc), Some(er), Some(ec)),
669 },
670 None,
671 )
672 }
673
674 #[test]
675 fn choosecols_basic_and_negative_and_duplicates() {
676 let wb = TestWorkbook::new().with_function(Arc::new(ChooseColsFn));
677 let wb = wb
678 .with_cell_a1("Sheet1", "A1", LiteralValue::Int(1))
679 .with_cell_a1("Sheet1", "B1", LiteralValue::Int(2))
680 .with_cell_a1("Sheet1", "C1", LiteralValue::Int(3))
681 .with_cell_a1("Sheet1", "A2", LiteralValue::Int(10))
682 .with_cell_a1("Sheet1", "B2", LiteralValue::Int(20))
683 .with_cell_a1("Sheet1", "C2", LiteralValue::Int(30));
684 let ctx = wb.interpreter();
685 let arr = range("A1:C2", 1, 1, 2, 3);
686 let f = ctx.context.get_function("", "CHOOSECOLS").unwrap();
687 let one = lit(LiteralValue::Int(1));
688 let three = lit(LiteralValue::Int(3));
689 let neg_one = lit(LiteralValue::Int(-1));
690 let args = vec![
692 ArgumentHandle::new(&arr, &ctx),
693 ArgumentHandle::new(&one, &ctx),
694 ArgumentHandle::new(&three, &ctx),
695 ];
696 let v = f
697 .dispatch(&args, &ctx.function_context(None))
698 .unwrap()
699 .into_literal();
700 match v {
701 LiteralValue::Array(a) => {
702 assert_eq!(a.len(), 2);
703 assert_eq!(
704 a[0],
705 vec![LiteralValue::Number(1.0), LiteralValue::Number(3.0)]
706 );
707 }
708 other => panic!("expected array got {other:?}"),
709 }
710 let args_neg = vec![
712 ArgumentHandle::new(&arr, &ctx),
713 ArgumentHandle::new(&neg_one, &ctx),
714 ];
715 let v2 = f
716 .dispatch(&args_neg, &ctx.function_context(None))
717 .unwrap()
718 .into_literal();
719 match v2 {
720 LiteralValue::Array(a) => {
721 assert_eq!(a[0], vec![LiteralValue::Number(3.0)]);
722 }
723 other => panic!("expected array last col got {other:?}"),
724 }
725 let args_dup = vec![
727 ArgumentHandle::new(&arr, &ctx),
728 ArgumentHandle::new(&one, &ctx),
729 ArgumentHandle::new(&one, &ctx),
730 ];
731 let v3 = f
732 .dispatch(&args_dup, &ctx.function_context(None))
733 .unwrap()
734 .into_literal();
735 match v3 {
736 LiteralValue::Array(a) => {
737 assert_eq!(
738 a[0],
739 vec![LiteralValue::Number(1.0), LiteralValue::Number(1.0)]
740 );
741 }
742 other => panic!("expected dup cols got {other:?}"),
743 }
744 }
745
746 #[test]
747 fn choosecols_out_of_range() {
748 let wb = TestWorkbook::new().with_function(Arc::new(ChooseColsFn));
749 let wb = wb
750 .with_cell_a1("Sheet1", "A1", LiteralValue::Int(1))
751 .with_cell_a1("Sheet1", "B1", LiteralValue::Int(2));
752 let ctx = wb.interpreter();
753 let arr = range("A1:B1", 1, 1, 1, 2);
754 let f = ctx.context.get_function("", "CHOOSECOLS").unwrap();
755 let three = lit(LiteralValue::Int(3));
756 let args = vec![
757 ArgumentHandle::new(&arr, &ctx),
758 ArgumentHandle::new(&three, &ctx),
759 ];
760 let v = f
761 .dispatch(&args, &ctx.function_context(None))
762 .unwrap()
763 .into_literal();
764 match v {
765 LiteralValue::Error(e) => assert_eq!(e.kind, ExcelErrorKind::Value),
766 other => panic!("expected #VALUE! got {other:?}"),
767 }
768 }
769
770 #[test]
771 fn chooserows_basic_and_negative() {
772 let wb = TestWorkbook::new().with_function(Arc::new(ChooseRowsFn));
773 let wb = wb
774 .with_cell_a1("Sheet1", "A1", LiteralValue::Int(1))
775 .with_cell_a1("Sheet1", "B1", LiteralValue::Int(2))
776 .with_cell_a1("Sheet1", "A2", LiteralValue::Int(10))
777 .with_cell_a1("Sheet1", "B2", LiteralValue::Int(20))
778 .with_cell_a1("Sheet1", "A3", LiteralValue::Int(100))
779 .with_cell_a1("Sheet1", "B3", LiteralValue::Int(200));
780 let ctx = wb.interpreter();
781 let arr = range("A1:B3", 1, 1, 3, 2);
782 let f = ctx.context.get_function("", "CHOOSEROWS").unwrap();
783 let one = lit(LiteralValue::Int(1));
784 let neg_one = lit(LiteralValue::Int(-1));
785 let args = vec![
786 ArgumentHandle::new(&arr, &ctx),
787 ArgumentHandle::new(&one, &ctx),
788 ArgumentHandle::new(&neg_one, &ctx),
789 ];
790 let v = f
791 .dispatch(&args, &ctx.function_context(None))
792 .unwrap()
793 .into_literal();
794 match v {
795 LiteralValue::Array(a) => {
796 assert_eq!(a.len(), 2);
797 assert_eq!(a[0][0], LiteralValue::Number(1.0));
798 assert_eq!(a[1][0], LiteralValue::Number(100.0));
799 }
800 other => panic!("expected array got {other:?}"),
801 }
802 }
803
804 #[test]
805 fn chooserows_out_of_range() {
806 let wb = TestWorkbook::new().with_function(Arc::new(ChooseRowsFn));
807 let wb = wb.with_cell_a1("Sheet1", "A1", LiteralValue::Int(1));
808 let ctx = wb.interpreter();
809 let arr = range("A1:A1", 1, 1, 1, 1);
810 let f = ctx.context.get_function("", "CHOOSEROWS").unwrap();
811 let two = lit(LiteralValue::Int(2));
812 let args = vec![
813 ArgumentHandle::new(&arr, &ctx),
814 ArgumentHandle::new(&two, &ctx),
815 ];
816 let v = f
817 .dispatch(&args, &ctx.function_context(None))
818 .unwrap()
819 .into_literal();
820 match v {
821 LiteralValue::Error(e) => assert_eq!(e.kind, ExcelErrorKind::Value),
822 other => panic!("expected #VALUE! got {other:?}"),
823 }
824 }
825}