formualizer_eval/builtins/lookup/
choose.rs

1//! CHOOSE function - selects a value from a list based on an index
2//!
3//! Excel semantics:
4//! - CHOOSE(index_num, value1, [value2], ...)
5//! - index_num must be between 1 and the number of values
6//! - Returns #VALUE! if index is out of range or not numeric
7//! - Can return references, not just values
8
9use crate::args::{ArgSchema, CoercionPolicy, ShapeKind};
10use crate::function::Function;
11use crate::traits::{ArgumentHandle, FunctionContext};
12use formualizer_common::{ArgKind, ExcelError, ExcelErrorKind, LiteralValue};
13use formualizer_macros::func_caps;
14
15#[derive(Debug)]
16pub struct ChooseFn;
17#[derive(Debug)]
18pub struct ChooseColsFn;
19#[derive(Debug)]
20pub struct ChooseRowsFn;
21
22impl Function for ChooseFn {
23    fn name(&self) -> &'static str {
24        "CHOOSE"
25    }
26
27    fn min_args(&self) -> usize {
28        2
29    }
30
31    func_caps!(PURE, LOOKUP);
32
33    fn arg_schema(&self) -> &'static [ArgSchema] {
34        use once_cell::sync::Lazy;
35        static SCHEMA: Lazy<Vec<ArgSchema>> = Lazy::new(|| {
36            vec![
37                // index_num (strict numeric)
38                ArgSchema {
39                    kinds: smallvec::smallvec![ArgKind::Number],
40                    required: true,
41                    by_ref: false,
42                    shape: ShapeKind::Scalar,
43                    coercion: CoercionPolicy::NumberStrict,
44                    max: None,
45                    repeating: None,
46                    default: None,
47                },
48                // value1, value2, ... (variadic, any type)
49                ArgSchema {
50                    kinds: smallvec::smallvec![ArgKind::Any],
51                    required: true,
52                    by_ref: false, // Could be reference but we'll unwrap value or pass through
53                    shape: ShapeKind::Scalar, // Treat each choice as scalar (top-left if range)
54                    coercion: CoercionPolicy::None,
55                    max: None,
56                    repeating: Some(1), // any number of choices after index
57                    default: None,
58                },
59            ]
60        });
61        &SCHEMA
62    }
63
64    fn eval_scalar<'a, 'b>(
65        &self,
66        args: &'a [ArgumentHandle<'a, 'b>],
67        _ctx: &dyn FunctionContext,
68    ) -> Result<LiteralValue, ExcelError> {
69        if args.len() < 2 {
70            return Ok(LiteralValue::Error(ExcelError::new(ExcelErrorKind::Value)));
71        }
72
73        // Get index
74        let index_val = args[0].value()?;
75        if let LiteralValue::Error(e) = index_val.as_ref() {
76            return Ok(LiteralValue::Error(e.clone()));
77        }
78
79        let index = match index_val.as_ref() {
80            LiteralValue::Number(n) => *n as i64,
81            LiteralValue::Int(i) => *i,
82            LiteralValue::Text(s) => s.parse::<f64>().map(|n| n as i64).unwrap_or(-1),
83            _ => return Ok(LiteralValue::Error(ExcelError::new(ExcelErrorKind::Value))),
84        };
85
86        // Check bounds
87        if index < 1 || index as usize > args.len() - 1 {
88            return Ok(LiteralValue::Error(ExcelError::new(ExcelErrorKind::Value)));
89        }
90
91        // Return the selected value (1-based indexing for the choice)
92        let selected_arg = &args[index as usize];
93        match selected_arg.value() {
94            Ok(v) => Ok(v.as_ref().clone()),
95            Err(e) => Ok(LiteralValue::Error(e)),
96        }
97    }
98}
99
100/* ───────────────────────── CHOOSECOLS() / CHOOSEROWS() ───────────────────────── */
101
102fn materialize_rows_2d(
103    arg: &ArgumentHandle,
104    ctx: &dyn FunctionContext,
105) -> Result<Vec<Vec<formualizer_common::LiteralValue>>, ExcelError> {
106    if let Ok(r) = arg.as_reference_or_eval() {
107        let mut rows: Vec<Vec<LiteralValue>> = Vec::new();
108        let sheet = "Sheet1"; // TODO propagate sheet
109        let rv = ctx.resolve_range_view(&r, sheet)?;
110        rv.for_each_row(&mut |row| {
111            rows.push(row.to_vec());
112            Ok(())
113        })?;
114        Ok(rows)
115    } else {
116        match arg.value()?.as_ref() {
117            LiteralValue::Array(a) => Ok(a.clone()),
118            v => Ok(vec![vec![v.clone()]]),
119        }
120    }
121}
122
123impl Function for ChooseColsFn {
124    func_caps!(PURE, LOOKUP);
125    fn name(&self) -> &'static str {
126        "CHOOSECOLS"
127    }
128    fn min_args(&self) -> usize {
129        2
130    }
131    fn variadic(&self) -> bool {
132        true
133    }
134    fn arg_schema(&self) -> &'static [ArgSchema] {
135        use once_cell::sync::Lazy;
136        static SCHEMA: Lazy<Vec<ArgSchema>> = Lazy::new(|| {
137            vec![
138                // array
139                ArgSchema {
140                    kinds: smallvec::smallvec![ArgKind::Range, ArgKind::Any],
141                    required: true,
142                    by_ref: false,
143                    shape: ShapeKind::Range,
144                    coercion: CoercionPolicy::None,
145                    max: None,
146                    repeating: None,
147                    default: None,
148                },
149                // col_num1 and subsequent
150                ArgSchema {
151                    kinds: smallvec::smallvec![ArgKind::Number],
152                    required: true,
153                    by_ref: false,
154                    shape: ShapeKind::Scalar,
155                    coercion: CoercionPolicy::NumberLenientText,
156                    max: None,
157                    repeating: Some(1),
158                    default: None,
159                },
160            ]
161        });
162        &SCHEMA
163    }
164    fn eval_scalar<'a, 'b>(
165        &self,
166        args: &'a [ArgumentHandle<'a, 'b>],
167        ctx: &dyn FunctionContext,
168    ) -> Result<LiteralValue, ExcelError> {
169        if args.len() < 2 {
170            return Ok(LiteralValue::Error(ExcelError::new(ExcelErrorKind::Value)));
171        }
172        let rows = materialize_rows_2d(&args[0], ctx)?;
173        if rows.is_empty() {
174            return Ok(LiteralValue::Array(vec![]));
175        }
176        let width = rows[0].len();
177        if rows.iter().any(|r| r.len() != width) {
178            return Ok(LiteralValue::Error(ExcelError::new(ExcelErrorKind::Value)));
179        }
180        let mut indices: Vec<usize> = Vec::new();
181        for a in &args[1..] {
182            let v = a.value()?;
183            if let LiteralValue::Error(e) = v.as_ref() {
184                return Ok(LiteralValue::Error(e.clone()));
185            }
186            let raw = match v.as_ref() {
187                LiteralValue::Int(i) => *i,
188                LiteralValue::Number(n) => *n as i64,
189                _ => return Ok(LiteralValue::Error(ExcelError::new(ExcelErrorKind::Value))),
190            };
191            if raw == 0 {
192                return Ok(LiteralValue::Error(ExcelError::new(ExcelErrorKind::Value)));
193            }
194            let adj = if raw > 0 {
195                raw - 1
196            } else {
197                (width as i64) + raw
198            };
199            if adj < 0 || adj as usize >= width {
200                return Ok(LiteralValue::Error(ExcelError::new(ExcelErrorKind::Value)));
201            }
202            indices.push(adj as usize);
203        }
204        let mut out: Vec<Vec<LiteralValue>> = Vec::with_capacity(rows.len());
205        for r in &rows {
206            let mut new_row = Vec::with_capacity(indices.len());
207            for &c in &indices {
208                new_row.push(r[c].clone());
209            }
210            out.push(new_row);
211        }
212        if out.len() == 1 && out[0].len() == 1 {
213            return Ok(out[0][0].clone());
214        }
215        Ok(LiteralValue::Array(out))
216    }
217}
218
219impl Function for ChooseRowsFn {
220    func_caps!(PURE, LOOKUP);
221    fn name(&self) -> &'static str {
222        "CHOOSEROWS"
223    }
224    fn min_args(&self) -> usize {
225        2
226    }
227    fn variadic(&self) -> bool {
228        true
229    }
230    fn arg_schema(&self) -> &'static [ArgSchema] {
231        use once_cell::sync::Lazy;
232        static SCHEMA: Lazy<Vec<ArgSchema>> = Lazy::new(|| {
233            vec![
234                // array
235                ArgSchema {
236                    kinds: smallvec::smallvec![ArgKind::Range, ArgKind::Any],
237                    required: true,
238                    by_ref: false,
239                    shape: ShapeKind::Range,
240                    coercion: CoercionPolicy::None,
241                    max: None,
242                    repeating: None,
243                    default: None,
244                },
245                // row_num1...
246                ArgSchema {
247                    kinds: smallvec::smallvec![ArgKind::Number],
248                    required: true,
249                    by_ref: false,
250                    shape: ShapeKind::Scalar,
251                    coercion: CoercionPolicy::NumberLenientText,
252                    max: None,
253                    repeating: Some(1),
254                    default: None,
255                },
256            ]
257        });
258        &SCHEMA
259    }
260    fn eval_scalar<'a, 'b>(
261        &self,
262        args: &'a [ArgumentHandle<'a, 'b>],
263        ctx: &dyn FunctionContext,
264    ) -> Result<LiteralValue, ExcelError> {
265        if args.len() < 2 {
266            return Ok(LiteralValue::Error(ExcelError::new(ExcelErrorKind::Value)));
267        }
268        let rows = materialize_rows_2d(&args[0], ctx)?;
269        if rows.is_empty() {
270            return Ok(LiteralValue::Array(vec![]));
271        }
272        let height = rows.len();
273        let mut indices: Vec<usize> = Vec::new();
274        for a in &args[1..] {
275            let v = a.value()?;
276            if let LiteralValue::Error(e) = v.as_ref() {
277                return Ok(LiteralValue::Error(e.clone()));
278            }
279            let raw = match v.as_ref() {
280                LiteralValue::Int(i) => *i,
281                LiteralValue::Number(n) => *n as i64,
282                _ => return Ok(LiteralValue::Error(ExcelError::new(ExcelErrorKind::Value))),
283            };
284            if raw == 0 {
285                return Ok(LiteralValue::Error(ExcelError::new(ExcelErrorKind::Value)));
286            }
287            let adj = if raw > 0 {
288                raw - 1
289            } else {
290                (height as i64) + raw
291            };
292            if adj < 0 || adj as usize >= height {
293                return Ok(LiteralValue::Error(ExcelError::new(ExcelErrorKind::Value)));
294            }
295            indices.push(adj as usize);
296        }
297        let mut out: Vec<Vec<LiteralValue>> = Vec::with_capacity(indices.len());
298        for &r in &indices {
299            out.push(rows[r].clone());
300        }
301        if out.len() == 1 && out[0].len() == 1 {
302            return Ok(out[0][0].clone());
303        }
304        Ok(LiteralValue::Array(out))
305    }
306}
307
308#[cfg(test)]
309mod tests {
310    use super::*;
311    use crate::test_workbook::TestWorkbook;
312    use formualizer_parse::parser::{ASTNode, ASTNodeType, ReferenceType};
313    use std::sync::Arc;
314
315    fn lit(v: LiteralValue) -> ASTNode {
316        ASTNode::new(ASTNodeType::Literal(v), None)
317    }
318
319    #[test]
320    fn choose_basic() {
321        let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
322        let ctx = wb.interpreter();
323        let f = ctx.context.get_function("", "CHOOSE").unwrap();
324
325        // CHOOSE(2, "A", "B", "C") -> "B"
326        let two = lit(LiteralValue::Int(2));
327        let a = lit(LiteralValue::Text("A".into()));
328        let b = lit(LiteralValue::Text("B".into()));
329        let c = lit(LiteralValue::Text("C".into()));
330
331        let args = vec![
332            ArgumentHandle::new(&two, &ctx),
333            ArgumentHandle::new(&a, &ctx),
334            ArgumentHandle::new(&b, &ctx),
335            ArgumentHandle::new(&c, &ctx),
336        ];
337
338        let result = f.dispatch(&args, &ctx.function_context(None)).unwrap();
339        assert_eq!(result, LiteralValue::Text("B".into()));
340    }
341
342    #[test]
343    fn choose_numeric_values() {
344        let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
345        let ctx = wb.interpreter();
346        let f = ctx.context.get_function("", "CHOOSE").unwrap();
347
348        // CHOOSE(3, 10, 20, 30, 40) -> 30
349        let three = lit(LiteralValue::Int(3));
350        let ten = lit(LiteralValue::Int(10));
351        let twenty = lit(LiteralValue::Int(20));
352        let thirty = lit(LiteralValue::Int(30));
353        let forty = lit(LiteralValue::Int(40));
354
355        let args = vec![
356            ArgumentHandle::new(&three, &ctx),
357            ArgumentHandle::new(&ten, &ctx),
358            ArgumentHandle::new(&twenty, &ctx),
359            ArgumentHandle::new(&thirty, &ctx),
360            ArgumentHandle::new(&forty, &ctx),
361        ];
362
363        let result = f.dispatch(&args, &ctx.function_context(None)).unwrap();
364        assert_eq!(result, LiteralValue::Int(30));
365    }
366
367    #[test]
368    fn choose_out_of_range() {
369        let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
370        let ctx = wb.interpreter();
371        let f = ctx.context.get_function("", "CHOOSE").unwrap();
372
373        // CHOOSE(5, "A", "B", "C") -> #VALUE! (only 3 choices)
374        let five = lit(LiteralValue::Int(5));
375        let a = lit(LiteralValue::Text("A".into()));
376        let b = lit(LiteralValue::Text("B".into()));
377        let c = lit(LiteralValue::Text("C".into()));
378
379        let args = vec![
380            ArgumentHandle::new(&five, &ctx),
381            ArgumentHandle::new(&a, &ctx),
382            ArgumentHandle::new(&b, &ctx),
383            ArgumentHandle::new(&c, &ctx),
384        ];
385
386        let result = f.dispatch(&args, &ctx.function_context(None)).unwrap();
387        assert!(matches!(result, LiteralValue::Error(e) if e.kind == ExcelErrorKind::Value));
388
389        // CHOOSE(0, "A", "B") -> #VALUE! (index must be >= 1)
390        let zero = lit(LiteralValue::Int(0));
391        let args2 = vec![
392            ArgumentHandle::new(&zero, &ctx),
393            ArgumentHandle::new(&a, &ctx),
394            ArgumentHandle::new(&b, &ctx),
395        ];
396
397        let result2 = f.dispatch(&args2, &ctx.function_context(None)).unwrap();
398        assert!(matches!(result2, LiteralValue::Error(e) if e.kind == ExcelErrorKind::Value));
399    }
400
401    #[test]
402    fn choose_decimal_index() {
403        let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
404        let ctx = wb.interpreter();
405        let f = ctx.context.get_function("", "CHOOSE").unwrap();
406
407        // CHOOSE(2.7, "A", "B", "C") -> "B" (truncates to 2)
408        let two_seven = lit(LiteralValue::Number(2.7));
409        let a = lit(LiteralValue::Text("A".into()));
410        let b = lit(LiteralValue::Text("B".into()));
411        let c = lit(LiteralValue::Text("C".into()));
412
413        let args = vec![
414            ArgumentHandle::new(&two_seven, &ctx),
415            ArgumentHandle::new(&a, &ctx),
416            ArgumentHandle::new(&b, &ctx),
417            ArgumentHandle::new(&c, &ctx),
418        ];
419
420        let result = f.dispatch(&args, &ctx.function_context(None)).unwrap();
421        assert_eq!(result, LiteralValue::Text("B".into()));
422    }
423
424    #[test]
425    fn choose_text_index_numeric_string() {
426        let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
427        let ctx = wb.interpreter();
428        let f = ctx.context.get_function("", "CHOOSE").unwrap();
429        let two_txt = lit(LiteralValue::Text("2".into()));
430        let a = lit(LiteralValue::Text("A".into()));
431        let b = lit(LiteralValue::Text("B".into()));
432        let c = lit(LiteralValue::Text("C".into()));
433        let args = vec![
434            ArgumentHandle::new(&two_txt, &ctx),
435            ArgumentHandle::new(&a, &ctx),
436            ArgumentHandle::new(&b, &ctx),
437            ArgumentHandle::new(&c, &ctx),
438        ];
439        let result = f.dispatch(&args, &ctx.function_context(None)).unwrap();
440        // Current engine uses NumberStrict coercion for index: numeric text not accepted -> #VALUE!
441        assert!(matches!(result, LiteralValue::Error(e) if e.kind == ExcelErrorKind::Value));
442    }
443
444    #[test]
445    fn choose_decimal_less_than_one() {
446        let wb = TestWorkbook::new().with_function(Arc::new(ChooseFn));
447        let ctx = wb.interpreter();
448        let f = ctx.context.get_function("", "CHOOSE").unwrap();
449        let zero_nine = lit(LiteralValue::Number(0.9));
450        let a = lit(LiteralValue::Text("A".into()));
451        let b = lit(LiteralValue::Text("B".into()));
452        let args = vec![
453            ArgumentHandle::new(&zero_nine, &ctx),
454            ArgumentHandle::new(&a, &ctx),
455            ArgumentHandle::new(&b, &ctx),
456        ];
457        let result = f.dispatch(&args, &ctx.function_context(None)).unwrap();
458        assert!(matches!(result, LiteralValue::Error(e) if e.kind == ExcelErrorKind::Value));
459    }
460
461    fn range(r: &str, sr: u32, sc: u32, er: u32, ec: u32) -> ASTNode {
462        ASTNode::new(
463            ASTNodeType::Reference {
464                original: r.into(),
465                reference: ReferenceType::Range {
466                    sheet: None,
467                    start_row: Some(sr),
468                    start_col: Some(sc),
469                    end_row: Some(er),
470                    end_col: Some(ec),
471                },
472            },
473            None,
474        )
475    }
476
477    #[test]
478    fn choosecols_basic_and_negative_and_duplicates() {
479        let wb = TestWorkbook::new().with_function(Arc::new(ChooseColsFn));
480        let wb = wb
481            .with_cell_a1("Sheet1", "A1", LiteralValue::Int(1))
482            .with_cell_a1("Sheet1", "B1", LiteralValue::Int(2))
483            .with_cell_a1("Sheet1", "C1", LiteralValue::Int(3))
484            .with_cell_a1("Sheet1", "A2", LiteralValue::Int(10))
485            .with_cell_a1("Sheet1", "B2", LiteralValue::Int(20))
486            .with_cell_a1("Sheet1", "C2", LiteralValue::Int(30));
487        let ctx = wb.interpreter();
488        let arr = range("A1:C2", 1, 1, 2, 3);
489        let f = ctx.context.get_function("", "CHOOSECOLS").unwrap();
490        let one = lit(LiteralValue::Int(1));
491        let three = lit(LiteralValue::Int(3));
492        let neg_one = lit(LiteralValue::Int(-1));
493        // pick first & third (positive indices)
494        let args = vec![
495            ArgumentHandle::new(&arr, &ctx),
496            ArgumentHandle::new(&one, &ctx),
497            ArgumentHandle::new(&three, &ctx),
498        ];
499        let v = f.dispatch(&args, &ctx.function_context(None)).unwrap();
500        match v {
501            LiteralValue::Array(a) => {
502                assert_eq!(a.len(), 2);
503                assert_eq!(a[0], vec![LiteralValue::Int(1), LiteralValue::Int(3)]);
504            }
505            other => panic!("expected array got {other:?}"),
506        }
507        // negative -1 -> last col only
508        let args_neg = vec![
509            ArgumentHandle::new(&arr, &ctx),
510            ArgumentHandle::new(&neg_one, &ctx),
511        ];
512        let v2 = f.dispatch(&args_neg, &ctx.function_context(None)).unwrap();
513        match v2 {
514            LiteralValue::Array(a) => {
515                assert_eq!(a[0], vec![LiteralValue::Int(3)]);
516            }
517            other => panic!("expected array last col got {other:?}"),
518        }
519        // duplicates (1,1)
520        let args_dup = vec![
521            ArgumentHandle::new(&arr, &ctx),
522            ArgumentHandle::new(&one, &ctx),
523            ArgumentHandle::new(&one, &ctx),
524        ];
525        let v3 = f.dispatch(&args_dup, &ctx.function_context(None)).unwrap();
526        match v3 {
527            LiteralValue::Array(a) => {
528                assert_eq!(a[0], vec![LiteralValue::Int(1), LiteralValue::Int(1)]);
529            }
530            other => panic!("expected dup cols got {other:?}"),
531        }
532    }
533
534    #[test]
535    fn choosecols_out_of_range() {
536        let wb = TestWorkbook::new().with_function(Arc::new(ChooseColsFn));
537        let wb = wb
538            .with_cell_a1("Sheet1", "A1", LiteralValue::Int(1))
539            .with_cell_a1("Sheet1", "B1", LiteralValue::Int(2));
540        let ctx = wb.interpreter();
541        let arr = range("A1:B1", 1, 1, 1, 2);
542        let f = ctx.context.get_function("", "CHOOSECOLS").unwrap();
543        let three = lit(LiteralValue::Int(3));
544        let args = vec![
545            ArgumentHandle::new(&arr, &ctx),
546            ArgumentHandle::new(&three, &ctx),
547        ];
548        let v = f.dispatch(&args, &ctx.function_context(None)).unwrap();
549        match v {
550            LiteralValue::Error(e) => assert_eq!(e.kind, ExcelErrorKind::Value),
551            other => panic!("expected #VALUE! got {other:?}"),
552        }
553    }
554
555    #[test]
556    fn chooserows_basic_and_negative() {
557        let wb = TestWorkbook::new().with_function(Arc::new(ChooseRowsFn));
558        let wb = wb
559            .with_cell_a1("Sheet1", "A1", LiteralValue::Int(1))
560            .with_cell_a1("Sheet1", "B1", LiteralValue::Int(2))
561            .with_cell_a1("Sheet1", "A2", LiteralValue::Int(10))
562            .with_cell_a1("Sheet1", "B2", LiteralValue::Int(20))
563            .with_cell_a1("Sheet1", "A3", LiteralValue::Int(100))
564            .with_cell_a1("Sheet1", "B3", LiteralValue::Int(200));
565        let ctx = wb.interpreter();
566        let arr = range("A1:B3", 1, 1, 3, 2);
567        let f = ctx.context.get_function("", "CHOOSEROWS").unwrap();
568        let one = lit(LiteralValue::Int(1));
569        let neg_one = lit(LiteralValue::Int(-1));
570        let args = vec![
571            ArgumentHandle::new(&arr, &ctx),
572            ArgumentHandle::new(&one, &ctx),
573            ArgumentHandle::new(&neg_one, &ctx),
574        ];
575        let v = f.dispatch(&args, &ctx.function_context(None)).unwrap();
576        match v {
577            LiteralValue::Array(a) => {
578                assert_eq!(a.len(), 2);
579                assert_eq!(a[0][0], LiteralValue::Int(1));
580                assert_eq!(a[1][0], LiteralValue::Int(100));
581            }
582            other => panic!("expected array got {other:?}"),
583        }
584    }
585
586    #[test]
587    fn chooserows_out_of_range() {
588        let wb = TestWorkbook::new().with_function(Arc::new(ChooseRowsFn));
589        let wb = wb.with_cell_a1("Sheet1", "A1", LiteralValue::Int(1));
590        let ctx = wb.interpreter();
591        let arr = range("A1:A1", 1, 1, 1, 1);
592        let f = ctx.context.get_function("", "CHOOSEROWS").unwrap();
593        let two = lit(LiteralValue::Int(2));
594        let args = vec![
595            ArgumentHandle::new(&arr, &ctx),
596            ArgumentHandle::new(&two, &ctx),
597        ];
598        let v = f.dispatch(&args, &ctx.function_context(None)).unwrap();
599        match v {
600            LiteralValue::Error(e) => assert_eq!(e.kind, ExcelErrorKind::Value),
601            other => panic!("expected #VALUE! got {other:?}"),
602        }
603    }
604}
formualizer_eval/builtins/lookup/choose.rs

formualizer_eval/builtins/lookup/
choose.rs
