Merge pull request #252 from nihilist-narwhal/fix-choose-array-index

CHOOSE: pick element by element when the index is an array

crates/calc/src/eval.rs

@@ -554,7 +554,20 @@ impl<'h> Evaluator<'h> {
                 if args.len() < 2 {
                     return err(CellError::Value);
                 }
-                let i = match self.arg_num(args, 0, 0.0) {
+                let index = self.value(&args[0]);
+                if let Value::Array(_) = index {
+                    // An array of indexes picks element by element: `CHOOSE({1,2},A1:A3,B1:B3)`
+                    // is a two-column array. The library CHOOSE lifts over its arguments.
+                    let Some(spec) = gridcraft_functions::lookup(name) else { return err(CellError::Name) };
+                    let mut evaluated = vec![Arg { value: index, from_ref: false }];
+                    for a in &args[1..] {
+                        let value = self.value(a);
+                        evaluated.push(Arg { value, from_ref: false });
+                    }
+                    let mut ctx = FnCtx { ev: self };
+                    return gridcraft_functions::call(spec, &evaluated, &mut ctx).pipe(Ev::V);
+                }
+                let i = match index.to_number() {
                     Ok(n) => n.trunc() as i64,
                     Err(e) => return err(e),
                 };

crates/calc/src/tests.rs

@@ -219,6 +219,25 @@ fn ifs_and_lazy() {
     assert_eq!(t.get("B4"), Value::from("five"));
 }
 
+#[test]
+fn choose_with_an_array_index() {
+    let mut t = T::new();
+    for (i, (k, v)) in [(1, 10), (2, 20), (3, 30)].iter().enumerate() {
+        t.set(&format!("A{}", i + 1), &k.to_string());
+        t.set(&format!("B{}", i + 1), &v.to_string());
+    }
+    t.set("D1", "=CHOOSE({1,2},\"a\",\"b\")");
+    assert_eq!(t.get("D1"), Value::from("a"));
+    assert_eq!(t.get("E1"), Value::from("b"));
+    // The usual "look left" idiom: a two-column array from two ranges.
+    t.set("D3", "=VLOOKUP(20,CHOOSE({1,2},B1:B3,A1:A3),2,FALSE)");
+    assert_eq!(t.num("D3"), 2.0);
+    t.set("D4", "=SUM(CHOOSE({1,2},A1:A3,B1:B3))");
+    assert_eq!(t.num("D4"), 66.0);
+    t.set("D5", "=CHOOSE(2,A1:A3,B1:B3)");
+    assert_eq!(t.num("D7"), 30.0);
+}
+
 #[test]
 fn references_functions() {
     let mut t = T::new();