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 {
24 fn name(&self) -> &'static str {
25 "CHOOSE"
26 }
27
28 fn min_args(&self) -> usize {
29 2
30 }
31
32 func_caps!(PURE, LOOKUP);
33
34 fn arg_schema(&self) -> &'static [ArgSchema] {
35 use once_cell::sync::Lazy;
36 static SCHEMA: Lazy<Vec<ArgSchema>> = Lazy::new(|| {
37 vec![
38 ArgSchema {
40 kinds: smallvec::smallvec![ArgKind::Number],
41 required: true,
42 by_ref: false,
43 shape: ShapeKind::Scalar,
44 coercion: CoercionPolicy::NumberStrict,
45 max: None,
46 repeating: None,
47 default: None,
48 },
49 ArgSchema {
51 kinds: smallvec::smallvec![ArgKind::Any],
52 required: true,
53 by_ref: false, shape: ShapeKind::Scalar, coercion: CoercionPolicy::None,
56 max: None,
57 repeating: Some(1), default: None,
59 },
60 ]
61 });
62 &SCHEMA
63 }
64
65 fn eval<'a, 'b, 'c>(
66 &self,
67 args: &'c [ArgumentHandle<'a, 'b>],
68 _ctx: &dyn FunctionContext<'b>,
69 ) -> Result<crate::traits::CalcValue<'b>, ExcelError> {
70 if args.len() < 2 {
71 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
72 ExcelError::new(ExcelErrorKind::Value),
73 )));
74 }
75
76 let index_val = args[0].value()?.into_literal();
78 if let LiteralValue::Error(e) = index_val {
79 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(e)));
80 }
81
82 let index = match index_val {
83 LiteralValue::Number(n) => n as i64,
84 LiteralValue::Int(i) => i,
85 LiteralValue::Text(s) => s.parse::<f64>().map(|n| n as i64).unwrap_or(-1),
86 _ => {
87 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
88 ExcelError::new(ExcelErrorKind::Value),
89 )));
90 }
91 };
92
93 if index < 1 || index as usize > args.len() - 1 {
95 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
96 ExcelError::new(ExcelErrorKind::Value),
97 )));
98 }
99
100 let selected_arg = &args[index as usize];
102 selected_arg.value()
103 }
104}
105
106fn materialize_rows_2d<'b>(
109 arg: &ArgumentHandle<'_, 'b>,
110 ctx: &dyn FunctionContext<'b>,
111) -> Result<Vec<Vec<formualizer_common::LiteralValue>>, ExcelError> {
112 if let Ok(r) = arg.as_reference_or_eval() {
113 let mut rows: Vec<Vec<LiteralValue>> = Vec::new();
114 let sheet = ctx.current_sheet();
115 let rv = ctx.resolve_range_view(&r, sheet)?;
116 rv.for_each_row(&mut |row| {
117 rows.push(row.to_vec());
118 Ok(())
119 })?;
120 Ok(rows)
121 } else {
122 let v = arg.value()?.into_literal();
123 match v {
124 LiteralValue::Array(a) => Ok(a),
125 other => Ok(vec![vec![other]]),
126 }
127 }
128}
129
130impl Function for ChooseColsFn {
131 func_caps!(PURE, LOOKUP);
132 fn name(&self) -> &'static str {
133 "CHOOSECOLS"
134 }
135 fn min_args(&self) -> usize {
136 2
137 }
138 fn variadic(&self) -> bool {
139 true
140 }
141 fn arg_schema(&self) -> &'static [ArgSchema] {
142 use once_cell::sync::Lazy;
143 static SCHEMA: Lazy<Vec<ArgSchema>> = Lazy::new(|| {
144 vec![
145 ArgSchema {
147 kinds: smallvec::smallvec![ArgKind::Range, ArgKind::Any],
148 required: true,
149 by_ref: false,
150 shape: ShapeKind::Range,
151 coercion: CoercionPolicy::None,
152 max: None,
153 repeating: None,
154 default: None,
155 },
156 ArgSchema {
158 kinds: smallvec::smallvec![ArgKind::Number],
159 required: true,
160 by_ref: false,
161 shape: ShapeKind::Scalar,
162 coercion: CoercionPolicy::NumberLenientText,
163 max: None,
164 repeating: Some(1),
165 default: None,
166 },
167 ]
168 });
169 &SCHEMA
170 }
171 fn eval<'a, 'b, 'c>(
172 &self,
173 args: &'c [ArgumentHandle<'a, 'b>],
174 _ctx: &dyn FunctionContext<'b>,
175 ) -> Result<crate::traits::CalcValue<'b>, ExcelError> {
176 if args.len() < 2 {
177 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
178 ExcelError::new(ExcelErrorKind::Value),
179 )));
180 }
181 let view = args[0].range_view()?;
182 let (rows, cols) = view.dims();
183 if rows == 0 || cols == 0 {
184 return Ok(crate::traits::CalcValue::Range(
185 crate::engine::range_view::RangeView::from_owned_rows(vec![], _ctx.date_system()),
186 ));
187 }
188
189 let mut indices: Vec<usize> = Vec::new();
190 for a in &args[1..] {
191 let v = a.value()?.into_literal();
192 if let LiteralValue::Error(e) = v {
193 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(e)));
194 }
195 let raw = match v {
196 LiteralValue::Int(i) => i,
197 LiteralValue::Number(n) => n as i64,
198 _ => {
199 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
200 ExcelError::new(ExcelErrorKind::Value),
201 )));
202 }
203 };
204 if raw == 0 {
205 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
206 ExcelError::new(ExcelErrorKind::Value),
207 )));
208 }
209 let adj = if raw > 0 {
210 raw - 1
211 } else {
212 (cols as i64) + raw
213 };
214 if adj < 0 || adj as usize >= cols {
215 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
216 ExcelError::new(ExcelErrorKind::Value),
217 )));
218 }
219 indices.push(adj as usize);
220 }
221 let mut out: Vec<Vec<LiteralValue>> = Vec::with_capacity(rows);
222 for r in 0..rows {
223 let mut new_row = Vec::with_capacity(indices.len());
224 for &c in &indices {
225 new_row.push(view.get_cell(r, c));
226 }
227 out.push(new_row);
228 }
229
230 Ok(collapse_if_scalar(out, _ctx.date_system()))
231 }
232}
233
234impl Function for ChooseRowsFn {
235 func_caps!(PURE, LOOKUP);
236 fn name(&self) -> &'static str {
237 "CHOOSEROWS"
238 }
239 fn min_args(&self) -> usize {
240 2
241 }
242 fn variadic(&self) -> bool {
243 true
244 }
245 fn arg_schema(&self) -> &'static [ArgSchema] {
246 use once_cell::sync::Lazy;
247 static SCHEMA: Lazy<Vec<ArgSchema>> = Lazy::new(|| {
248 vec![
249 ArgSchema {
251 kinds: smallvec::smallvec![ArgKind::Range, ArgKind::Any],
252 required: true,
253 by_ref: false,
254 shape: ShapeKind::Range,
255 coercion: CoercionPolicy::None,
256 max: None,
257 repeating: None,
258 default: None,
259 },
260 ArgSchema {
262 kinds: smallvec::smallvec![ArgKind::Number],
263 required: true,
264 by_ref: false,
265 shape: ShapeKind::Scalar,
266 coercion: CoercionPolicy::NumberLenientText,
267 max: None,
268 repeating: Some(1),
269 default: None,
270 },
271 ]
272 });
273 &SCHEMA
274 }
275 fn eval<'a, 'b, 'c>(
276 &self,
277 args: &'c [ArgumentHandle<'a, 'b>],
278 _ctx: &dyn FunctionContext<'b>,
279 ) -> Result<crate::traits::CalcValue<'b>, ExcelError> {
280 if args.len() < 2 {
281 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
282 ExcelError::new(ExcelErrorKind::Value),
283 )));
284 }
285 let view = args[0].range_view()?;
286 let (rows, cols) = view.dims();
287 if rows == 0 || cols == 0 {
288 return Ok(crate::traits::CalcValue::Range(
289 crate::engine::range_view::RangeView::from_owned_rows(vec![], _ctx.date_system()),
290 ));
291 }
292
293 let mut indices: Vec<usize> = Vec::new();
294 for a in &args[1..] {
295 let v = a.value()?.into_literal();
296 if let LiteralValue::Error(e) = v {
297 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(e)));
298 }
299 let raw = match v {
300 LiteralValue::Int(i) => i,
301 LiteralValue::Number(n) => n as i64,
302 _ => {
303 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
304 ExcelError::new(ExcelErrorKind::Value),
305 )));
306 }
307 };
308 if raw == 0 {
309 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
310 ExcelError::new(ExcelErrorKind::Value),
311 )));
312 }
313 let adj = if raw > 0 {
314 raw - 1
315 } else {
316 (rows as i64) + raw
317 };
318 if adj < 0 || adj as usize >= rows {
319 return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
320 ExcelError::new(ExcelErrorKind::Value),
321 )));
322 }
323 indices.push(adj as usize);
324 }
325 let mut out: Vec<Vec<LiteralValue>> = Vec::with_capacity(indices.len());
326 for &r in &indices {
327 let mut row_vals = Vec::with_capacity(cols);
328 for c in 0..cols {
329 row_vals.push(view.get_cell(r, c));
330 }
331 out.push(row_vals);
332 }
333
334 Ok(collapse_if_scalar(out, _ctx.date_system()))
335 }
336}
337
338#[cfg(test)]
339mod tests {
340 use super::*;
341 use crate::test_workbook::TestWorkbook;
342 use formualizer_parse::parser::{ASTNode, ASTNodeType, ReferenceType};
343 use std::sync::Arc;
344
345 fn lit(v: LiteralValue) -> ASTNode {
346 ASTNode::new(ASTNodeType::Literal(v), None)
347 }
348
349 #[test]
350 fn choose_basic() {
351 let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
352 let ctx = wb.interpreter();
353 let f = ctx.context.get_function("", "CHOOSE").unwrap();
354
355 let two = lit(LiteralValue::Int(2));
357 let a = lit(LiteralValue::Text("A".into()));
358 let b = lit(LiteralValue::Text("B".into()));
359 let c = lit(LiteralValue::Text("C".into()));
360
361 let args = vec![
362 ArgumentHandle::new(&two, &ctx),
363 ArgumentHandle::new(&a, &ctx),
364 ArgumentHandle::new(&b, &ctx),
365 ArgumentHandle::new(&c, &ctx),
366 ];
367
368 let result = f
369 .dispatch(&args, &ctx.function_context(None))
370 .unwrap()
371 .into_literal();
372 assert_eq!(result, LiteralValue::Text("B".into()));
373 }
374
375 #[test]
376 fn choose_numeric_values() {
377 let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
378 let ctx = wb.interpreter();
379 let f = ctx.context.get_function("", "CHOOSE").unwrap();
380
381 let three = lit(LiteralValue::Int(3));
383 let ten = lit(LiteralValue::Int(10));
384 let twenty = lit(LiteralValue::Int(20));
385 let thirty = lit(LiteralValue::Int(30));
386 let forty = lit(LiteralValue::Int(40));
387
388 let args = vec![
389 ArgumentHandle::new(&three, &ctx),
390 ArgumentHandle::new(&ten, &ctx),
391 ArgumentHandle::new(&twenty, &ctx),
392 ArgumentHandle::new(&thirty, &ctx),
393 ArgumentHandle::new(&forty, &ctx),
394 ];
395
396 let result = f
397 .dispatch(&args, &ctx.function_context(None))
398 .unwrap()
399 .into_literal();
400 assert_eq!(result, LiteralValue::Int(30));
401 }
402
403 #[test]
404 fn choose_out_of_range() {
405 let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
406 let ctx = wb.interpreter();
407 let f = ctx.context.get_function("", "CHOOSE").unwrap();
408
409 let five = lit(LiteralValue::Int(5));
411 let a = lit(LiteralValue::Text("A".into()));
412 let b = lit(LiteralValue::Text("B".into()));
413 let c = lit(LiteralValue::Text("C".into()));
414
415 let args = vec![
416 ArgumentHandle::new(&five, &ctx),
417 ArgumentHandle::new(&a, &ctx),
418 ArgumentHandle::new(&b, &ctx),
419 ArgumentHandle::new(&c, &ctx),
420 ];
421
422 let result = f
423 .dispatch(&args, &ctx.function_context(None))
424 .unwrap()
425 .into_literal();
426 assert!(matches!(result, LiteralValue::Error(e) if e.kind == ExcelErrorKind::Value));
427
428 let zero = lit(LiteralValue::Int(0));
430 let args2 = vec![
431 ArgumentHandle::new(&zero, &ctx),
432 ArgumentHandle::new(&a, &ctx),
433 ArgumentHandle::new(&b, &ctx),
434 ];
435
436 let result2 = f
437 .dispatch(&args2, &ctx.function_context(None))
438 .unwrap()
439 .into_literal();
440 assert!(matches!(result2, LiteralValue::Error(e) if e.kind == ExcelErrorKind::Value));
441 }
442
443 #[test]
444 fn choose_decimal_index() {
445 let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
446 let ctx = wb.interpreter();
447 let f = ctx.context.get_function("", "CHOOSE").unwrap();
448
449 let two_seven = lit(LiteralValue::Number(2.7));
451 let a = lit(LiteralValue::Text("A".into()));
452 let b = lit(LiteralValue::Text("B".into()));
453 let c = lit(LiteralValue::Text("C".into()));
454
455 let args = vec![
456 ArgumentHandle::new(&two_seven, &ctx),
457 ArgumentHandle::new(&a, &ctx),
458 ArgumentHandle::new(&b, &ctx),
459 ArgumentHandle::new(&c, &ctx),
460 ];
461
462 let result = f
463 .dispatch(&args, &ctx.function_context(None))
464 .unwrap()
465 .into_literal();
466 assert_eq!(result, LiteralValue::Text("B".into()));
467 }
468
469 #[test]
470 fn choose_text_index_numeric_string() {
471 let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
472 let ctx = wb.interpreter();
473 let f = ctx.context.get_function("", "CHOOSE").unwrap();
474 let two_txt = lit(LiteralValue::Text("2".into()));
475 let a = lit(LiteralValue::Text("A".into()));
476 let b = lit(LiteralValue::Text("B".into()));
477 let c = lit(LiteralValue::Text("C".into()));
478 let args = vec![
479 ArgumentHandle::new(&two_txt, &ctx),
480 ArgumentHandle::new(&a, &ctx),
481 ArgumentHandle::new(&b, &ctx),
482 ArgumentHandle::new(&c, &ctx),
483 ];
484 let result = f
485 .dispatch(&args, &ctx.function_context(None))
486 .unwrap()
487 .into_literal();
488 assert!(matches!(result, LiteralValue::Error(e) if e.kind == ExcelErrorKind::Value));
490 }
491
492 #[test]
493 fn choose_decimal_less_than_one() {
494 let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
495 let ctx = wb.interpreter();
496 let f = ctx.context.get_function("", "CHOOSE").unwrap();
497 let zero_nine = lit(LiteralValue::Number(0.9));
498 let a = lit(LiteralValue::Text("A".into()));
499 let b = lit(LiteralValue::Text("B".into()));
500 let args = vec![
501 ArgumentHandle::new(&zero_nine, &ctx),
502 ArgumentHandle::new(&a, &ctx),
503 ArgumentHandle::new(&b, &ctx),
504 ];
505 let result = f
506 .dispatch(&args, &ctx.function_context(None))
507 .unwrap()
508 .into_literal();
509 assert!(matches!(result, LiteralValue::Error(e) if e.kind == ExcelErrorKind::Value));
510 }
511
512 fn range(r: &str, sr: u32, sc: u32, er: u32, ec: u32) -> ASTNode {
513 ASTNode::new(
514 ASTNodeType::Reference {
515 original: r.into(),
516 reference: ReferenceType::range(None, Some(sr), Some(sc), Some(er), Some(ec)),
517 },
518 None,
519 )
520 }
521
522 #[test]
523 fn choosecols_basic_and_negative_and_duplicates() {
524 let wb = TestWorkbook::new().with_function(Arc::new(ChooseColsFn));
525 let wb = wb
526 .with_cell_a1("Sheet1", "A1", LiteralValue::Int(1))
527 .with_cell_a1("Sheet1", "B1", LiteralValue::Int(2))
528 .with_cell_a1("Sheet1", "C1", LiteralValue::Int(3))
529 .with_cell_a1("Sheet1", "A2", LiteralValue::Int(10))
530 .with_cell_a1("Sheet1", "B2", LiteralValue::Int(20))
531 .with_cell_a1("Sheet1", "C2", LiteralValue::Int(30));
532 let ctx = wb.interpreter();
533 let arr = range("A1:C2", 1, 1, 2, 3);
534 let f = ctx.context.get_function("", "CHOOSECOLS").unwrap();
535 let one = lit(LiteralValue::Int(1));
536 let three = lit(LiteralValue::Int(3));
537 let neg_one = lit(LiteralValue::Int(-1));
538 let args = vec![
540 ArgumentHandle::new(&arr, &ctx),
541 ArgumentHandle::new(&one, &ctx),
542 ArgumentHandle::new(&three, &ctx),
543 ];
544 let v = f
545 .dispatch(&args, &ctx.function_context(None))
546 .unwrap()
547 .into_literal();
548 match v {
549 LiteralValue::Array(a) => {
550 assert_eq!(a.len(), 2);
551 assert_eq!(
552 a[0],
553 vec![LiteralValue::Number(1.0), LiteralValue::Number(3.0)]
554 );
555 }
556 other => panic!("expected array got {other:?}"),
557 }
558 let args_neg = vec![
560 ArgumentHandle::new(&arr, &ctx),
561 ArgumentHandle::new(&neg_one, &ctx),
562 ];
563 let v2 = f
564 .dispatch(&args_neg, &ctx.function_context(None))
565 .unwrap()
566 .into_literal();
567 match v2 {
568 LiteralValue::Array(a) => {
569 assert_eq!(a[0], vec![LiteralValue::Number(3.0)]);
570 }
571 other => panic!("expected array last col got {other:?}"),
572 }
573 let args_dup = vec![
575 ArgumentHandle::new(&arr, &ctx),
576 ArgumentHandle::new(&one, &ctx),
577 ArgumentHandle::new(&one, &ctx),
578 ];
579 let v3 = f
580 .dispatch(&args_dup, &ctx.function_context(None))
581 .unwrap()
582 .into_literal();
583 match v3 {
584 LiteralValue::Array(a) => {
585 assert_eq!(
586 a[0],
587 vec![LiteralValue::Number(1.0), LiteralValue::Number(1.0)]
588 );
589 }
590 other => panic!("expected dup cols got {other:?}"),
591 }
592 }
593
594 #[test]
595 fn choosecols_out_of_range() {
596 let wb = TestWorkbook::new().with_function(Arc::new(ChooseColsFn));
597 let wb = wb
598 .with_cell_a1("Sheet1", "A1", LiteralValue::Int(1))
599 .with_cell_a1("Sheet1", "B1", LiteralValue::Int(2));
600 let ctx = wb.interpreter();
601 let arr = range("A1:B1", 1, 1, 1, 2);
602 let f = ctx.context.get_function("", "CHOOSECOLS").unwrap();
603 let three = lit(LiteralValue::Int(3));
604 let args = vec![
605 ArgumentHandle::new(&arr, &ctx),
606 ArgumentHandle::new(&three, &ctx),
607 ];
608 let v = f
609 .dispatch(&args, &ctx.function_context(None))
610 .unwrap()
611 .into_literal();
612 match v {
613 LiteralValue::Error(e) => assert_eq!(e.kind, ExcelErrorKind::Value),
614 other => panic!("expected #VALUE! got {other:?}"),
615 }
616 }
617
618 #[test]
619 fn chooserows_basic_and_negative() {
620 let wb = TestWorkbook::new().with_function(Arc::new(ChooseRowsFn));
621 let wb = wb
622 .with_cell_a1("Sheet1", "A1", LiteralValue::Int(1))
623 .with_cell_a1("Sheet1", "B1", LiteralValue::Int(2))
624 .with_cell_a1("Sheet1", "A2", LiteralValue::Int(10))
625 .with_cell_a1("Sheet1", "B2", LiteralValue::Int(20))
626 .with_cell_a1("Sheet1", "A3", LiteralValue::Int(100))
627 .with_cell_a1("Sheet1", "B3", LiteralValue::Int(200));
628 let ctx = wb.interpreter();
629 let arr = range("A1:B3", 1, 1, 3, 2);
630 let f = ctx.context.get_function("", "CHOOSEROWS").unwrap();
631 let one = lit(LiteralValue::Int(1));
632 let neg_one = lit(LiteralValue::Int(-1));
633 let args = vec![
634 ArgumentHandle::new(&arr, &ctx),
635 ArgumentHandle::new(&one, &ctx),
636 ArgumentHandle::new(&neg_one, &ctx),
637 ];
638 let v = f
639 .dispatch(&args, &ctx.function_context(None))
640 .unwrap()
641 .into_literal();
642 match v {
643 LiteralValue::Array(a) => {
644 assert_eq!(a.len(), 2);
645 assert_eq!(a[0][0], LiteralValue::Number(1.0));
646 assert_eq!(a[1][0], LiteralValue::Number(100.0));
647 }
648 other => panic!("expected array got {other:?}"),
649 }
650 }
651
652 #[test]
653 fn chooserows_out_of_range() {
654 let wb = TestWorkbook::new().with_function(Arc::new(ChooseRowsFn));
655 let wb = wb.with_cell_a1("Sheet1", "A1", LiteralValue::Int(1));
656 let ctx = wb.interpreter();
657 let arr = range("A1:A1", 1, 1, 1, 1);
658 let f = ctx.context.get_function("", "CHOOSEROWS").unwrap();
659 let two = lit(LiteralValue::Int(2));
660 let args = vec![
661 ArgumentHandle::new(&arr, &ctx),
662 ArgumentHandle::new(&two, &ctx),
663 ];
664 let v = f
665 .dispatch(&args, &ctx.function_context(None))
666 .unwrap()
667 .into_literal();
668 match v {
669 LiteralValue::Error(e) => assert_eq!(e.kind, ExcelErrorKind::Value),
670 other => panic!("expected #VALUE! got {other:?}"),
671 }
672 }
673}