Skip to main content

formualizer_eval/builtins/
reference_fns.rs

1use crate::args::{ArgSchema, CoercionPolicy, ShapeKind};
2use crate::function::{FnCaps, Function};
3use crate::traits::{ArgumentHandle, FunctionContext};
4use formualizer_common::{ArgKind, ExcelError, ExcelErrorKind, LiteralValue};
5use formualizer_parse::parser::ReferenceType;
6
7fn number_strict_scalar() -> ArgSchema {
8    ArgSchema {
9        kinds: smallvec::smallvec![ArgKind::Number],
10        required: true,
11        by_ref: false,
12        shape: ShapeKind::Scalar,
13        coercion: CoercionPolicy::NumberStrict,
14        max: None,
15        repeating: None,
16        default: None,
17    }
18}
19
20fn arg_byref_array() -> Vec<ArgSchema> {
21    vec![
22        // Accept both references and array literals
23        ArgSchema {
24            kinds: smallvec::smallvec![ArgKind::Any],
25            required: true,
26            by_ref: false,
27            shape: ShapeKind::Range,
28            coercion: CoercionPolicy::None,
29            max: None,
30            repeating: None,
31            default: None,
32        },
33        number_strict_scalar(),
34        // Column is optional for 1D arrays
35        ArgSchema {
36            kinds: smallvec::smallvec![ArgKind::Number],
37            required: false,
38            by_ref: false,
39            shape: ShapeKind::Scalar,
40            coercion: CoercionPolicy::NumberStrict,
41            max: None,
42            repeating: None,
43            default: None,
44        },
45    ]
46}
47
48fn arg_byref_reference() -> Vec<ArgSchema> {
49    vec![
50        ArgSchema {
51            kinds: smallvec::smallvec![ArgKind::Range],
52            required: true,
53            by_ref: true,
54            shape: ShapeKind::Range,
55            coercion: CoercionPolicy::None,
56            max: None,
57            repeating: None,
58            default: None,
59        },
60        number_strict_scalar(),
61        number_strict_scalar(),
62        ArgSchema {
63            // height optional
64            kinds: smallvec::smallvec![ArgKind::Number],
65            required: false,
66            by_ref: false,
67            shape: ShapeKind::Scalar,
68            coercion: CoercionPolicy::NumberStrict,
69            max: None,
70            repeating: None,
71            default: None,
72        },
73        ArgSchema {
74            // width optional
75            kinds: smallvec::smallvec![ArgKind::Number],
76            required: false,
77            by_ref: false,
78            shape: ShapeKind::Scalar,
79            coercion: CoercionPolicy::NumberStrict,
80            max: None,
81            repeating: None,
82            default: None,
83        },
84    ]
85}
86
87#[derive(Debug)]
88pub struct IndexFn;
89impl Function for IndexFn {
90    fn caps(&self) -> FnCaps {
91        FnCaps::PURE | FnCaps::RETURNS_REFERENCE
92    }
93    fn name(&self) -> &'static str {
94        "INDEX"
95    }
96    fn min_args(&self) -> usize {
97        2
98    }
99    fn arg_schema(&self) -> &'static [ArgSchema] {
100        use once_cell::sync::Lazy;
101        static SCHEMA: Lazy<Vec<ArgSchema>> = Lazy::new(arg_byref_array);
102        &SCHEMA
103    }
104
105    fn eval_reference<'a, 'b, 'c>(
106        &self,
107        args: &'c [ArgumentHandle<'a, 'b>],
108        _ctx: &dyn FunctionContext<'b>,
109    ) -> Option<Result<ReferenceType, ExcelError>> {
110        // args: array(by_ref), row, col (col optional for 1D)
111        if args.len() < 2 {
112            return Some(Err(ExcelError::new(ExcelErrorKind::Value)));
113        }
114        // Return None for array literals so eval() handles them
115        let base = match args[0].as_reference_or_eval() {
116            Ok(r) => r,
117            Err(_) => return None,
118        };
119        let row = match args[1].value() {
120            Ok(cv) => match cv.into_literal() {
121                LiteralValue::Number(n) => n as i64,
122                LiteralValue::Int(i) => i,
123                _ => return Some(Err(ExcelError::new(ExcelErrorKind::Value))),
124            },
125            Err(e) => return Some(Err(e)),
126        };
127        let col = if args.len() >= 3 {
128            match args[2].value() {
129                Ok(cv) => match cv.into_literal() {
130                    LiteralValue::Number(n) => n as i64,
131                    LiteralValue::Int(i) => i,
132                    _ => return Some(Err(ExcelError::new(ExcelErrorKind::Value))),
133                },
134                Err(e) => return Some(Err(e)),
135            }
136        } else {
137            // TODO(phase6): Document INDEX 1D behavior when col omitted.
138            1
139        };
140
141        // Only Range supported for now
142        let (sheet, sr, sc, er, ec) = match base {
143            ReferenceType::Range {
144                sheet,
145                start_row,
146                start_col,
147                end_row,
148                end_col,
149                ..
150            } => match (start_row, start_col, end_row, end_col) {
151                (Some(sr), Some(sc), Some(er), Some(ec)) => (sheet, sr, sc, er, ec),
152                _ => return Some(Err(ExcelError::new(ExcelErrorKind::Ref))),
153            },
154            ReferenceType::Cell {
155                sheet, row, col, ..
156            } => (sheet, row, col, row, col),
157            _ => return Some(Err(ExcelError::new(ExcelErrorKind::Ref))),
158        };
159
160        // 1-based indexing per Excel
161        if row <= 0 || col <= 0 {
162            return Some(Err(ExcelError::new(ExcelErrorKind::Ref)));
163        }
164        let r = sr + (row as u32) - 1;
165        let c = sc + (col as u32) - 1;
166        if r > er || c > ec {
167            return Some(Err(ExcelError::new(ExcelErrorKind::Ref)));
168        }
169
170        Some(Ok(ReferenceType::cell(sheet, r, c)))
171    }
172
173    fn eval<'a, 'b, 'c>(
174        &self,
175        args: &'c [ArgumentHandle<'a, 'b>],
176        ctx: &dyn FunctionContext<'b>,
177    ) -> Result<crate::traits::CalcValue<'b>, ExcelError> {
178        // First try to handle as a reference
179        if let Some(result) = self.eval_reference(args, ctx) {
180            match result {
181                Ok(r) => {
182                    // Materialize to value
183                    let current_sheet = ctx.current_sheet();
184                    match ctx.resolve_range_view(&r, current_sheet) {
185                        Ok(rv) => {
186                            let (rows, cols) = rv.dims();
187                            if rows == 1 && cols == 1 {
188                                Ok(crate::traits::CalcValue::Scalar(
189                                    rv.as_1x1().unwrap_or(LiteralValue::Empty),
190                                ))
191                            } else {
192                                Ok(crate::traits::CalcValue::Range(rv))
193                            }
194                        }
195                        Err(e) => Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(e))),
196                    }
197                }
198                Err(e) => Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(e))),
199            }
200        } else {
201            // Handle array literal
202            if args.len() < 2 {
203                return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
204                    ExcelError::new(ExcelErrorKind::Value),
205                )));
206            }
207            let v = args[0].value()?.into_literal();
208            let table: Vec<Vec<LiteralValue>> = match v {
209                LiteralValue::Array(rows) => rows,
210                other => vec![vec![other]],
211            };
212            let index = match args[1].value()?.into_literal() {
213                LiteralValue::Number(n) => n as i64,
214                LiteralValue::Int(i) => i,
215                _ => {
216                    return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
217                        ExcelError::new(ExcelErrorKind::Value),
218                    )));
219                }
220            };
221
222            // Determine if this is a 1D array (single row or single column)
223            let is_single_row = table.len() == 1;
224            let is_single_col = table.iter().all(|r| r.len() == 1);
225
226            // For 1D arrays with 2 args, index is position in the array
227            if args.len() == 2 && (is_single_row || is_single_col) {
228                // TODO(phase6): Document INDEX 1D behavior when col omitted.
229                if index <= 0 {
230                    return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
231                        ExcelError::new(ExcelErrorKind::Ref),
232                    )));
233                }
234                let idx = (index - 1) as usize;
235                let val = if is_single_row {
236                    table[0].get(idx).cloned()
237                } else {
238                    table.get(idx).and_then(|r| r.first()).cloned()
239                };
240                return Ok(crate::traits::CalcValue::Scalar(val.unwrap_or_else(|| {
241                    LiteralValue::Error(ExcelError::new(ExcelErrorKind::Ref))
242                })));
243            }
244
245            // 2D array or 3 arguments: use row and col indexing
246            let row = index as usize;
247            let col = if args.len() >= 3 {
248                match args[2].value()?.into_literal() {
249                    LiteralValue::Number(n) => n as usize,
250                    LiteralValue::Int(i) => i as usize,
251                    _ => {
252                        return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
253                            ExcelError::new(ExcelErrorKind::Value),
254                        )));
255                    }
256                }
257            } else {
258                1
259            };
260
261            // 1-based indexing
262            if row == 0 || col == 0 {
263                return Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
264                    ExcelError::new(ExcelErrorKind::Ref),
265                )));
266            }
267            let val = table
268                .get(row - 1)
269                .and_then(|r| r.get(col - 1))
270                .cloned()
271                .unwrap_or_else(|| LiteralValue::Error(ExcelError::new(ExcelErrorKind::Ref)));
272            Ok(crate::traits::CalcValue::Scalar(val))
273        }
274    }
275}
276
277#[derive(Debug)]
278pub struct OffsetFn;
279impl Function for OffsetFn {
280    fn caps(&self) -> FnCaps {
281        // OFFSET is volatile in Excel semantics
282        FnCaps::PURE | FnCaps::RETURNS_REFERENCE | FnCaps::VOLATILE
283    }
284    fn name(&self) -> &'static str {
285        "OFFSET"
286    }
287    fn min_args(&self) -> usize {
288        3
289    }
290    fn arg_schema(&self) -> &'static [ArgSchema] {
291        use once_cell::sync::Lazy;
292        static SCHEMA: Lazy<Vec<ArgSchema>> = Lazy::new(arg_byref_reference);
293        &SCHEMA
294    }
295
296    fn eval_reference<'a, 'b, 'c>(
297        &self,
298        args: &'c [ArgumentHandle<'a, 'b>],
299        _ctx: &dyn FunctionContext<'b>,
300    ) -> Option<Result<ReferenceType, ExcelError>> {
301        if args.len() < 3 {
302            return Some(Err(ExcelError::new(ExcelErrorKind::Value)));
303        }
304        let base = match args[0].as_reference_or_eval() {
305            Ok(r) => r,
306            Err(e) => return Some(Err(e)),
307        };
308        let dr = match args[1].value() {
309            Ok(cv) => match cv.into_literal() {
310                LiteralValue::Number(n) => n as i64,
311                LiteralValue::Int(i) => i,
312                _ => return Some(Err(ExcelError::new(ExcelErrorKind::Value))),
313            },
314            Err(e) => return Some(Err(e)),
315        };
316        let dc = match args[2].value() {
317            Ok(cv) => match cv.into_literal() {
318                LiteralValue::Number(n) => n as i64,
319                LiteralValue::Int(i) => i,
320                _ => return Some(Err(ExcelError::new(ExcelErrorKind::Value))),
321            },
322            Err(e) => return Some(Err(e)),
323        };
324
325        let (sheet, sr, sc, er, ec) = match base {
326            ReferenceType::Range {
327                sheet,
328                start_row,
329                start_col,
330                end_row,
331                end_col,
332                ..
333            } => match (start_row, start_col, end_row, end_col) {
334                (Some(sr), Some(sc), Some(er), Some(ec)) => (sheet, sr, sc, er, ec),
335                _ => return Some(Err(ExcelError::new(ExcelErrorKind::Ref))),
336            },
337            ReferenceType::Cell {
338                sheet, row, col, ..
339            } => (sheet, row, col, row, col),
340            _ => return Some(Err(ExcelError::new(ExcelErrorKind::Ref))),
341        };
342
343        let nsr = (sr as i64) + dr;
344        let nsc = (sc as i64) + dc;
345        let height = if args.len() >= 4 {
346            match args[3].value() {
347                Ok(cv) => match cv.into_literal() {
348                    LiteralValue::Number(n) => n as i64,
349                    LiteralValue::Int(i) => i,
350                    _ => return Some(Err(ExcelError::new(ExcelErrorKind::Value))),
351                },
352                Err(e) => return Some(Err(e)),
353            }
354        } else {
355            (er as i64) - (sr as i64) + 1
356        };
357        let width = if args.len() >= 5 {
358            match args[4].value() {
359                Ok(cv) => match cv.into_literal() {
360                    LiteralValue::Number(n) => n as i64,
361                    LiteralValue::Int(i) => i,
362                    _ => return Some(Err(ExcelError::new(ExcelErrorKind::Value))),
363                },
364                Err(e) => return Some(Err(e)),
365            }
366        } else {
367            (ec as i64) - (sc as i64) + 1
368        };
369
370        if nsr <= 0 || nsc <= 0 || height <= 0 || width <= 0 {
371            return Some(Err(ExcelError::new(ExcelErrorKind::Ref)));
372        }
373        let ner = nsr + height - 1;
374        let nec = nsc + width - 1;
375
376        if height == 1 && width == 1 {
377            Some(Ok(ReferenceType::cell(sheet, nsr as u32, nsc as u32)))
378        } else {
379            Some(Ok(ReferenceType::range(
380                sheet,
381                Some(nsr as u32),
382                Some(nsc as u32),
383                Some(ner as u32),
384                Some(nec as u32),
385            )))
386        }
387    }
388
389    fn eval<'a, 'b, 'c>(
390        &self,
391        args: &'c [ArgumentHandle<'a, 'b>],
392        ctx: &dyn FunctionContext<'b>,
393    ) -> Result<crate::traits::CalcValue<'b>, ExcelError> {
394        if let Some(Ok(r)) = self.eval_reference(args, ctx) {
395            let current_sheet = ctx.current_sheet();
396            match ctx.resolve_range_view(&r, current_sheet) {
397                Ok(rv) => {
398                    let (rows, cols) = rv.dims();
399                    if rows == 1 && cols == 1 {
400                        Ok(crate::traits::CalcValue::Scalar(
401                            rv.as_1x1().unwrap_or(LiteralValue::Empty),
402                        ))
403                    } else {
404                        Ok(crate::traits::CalcValue::Range(rv))
405                    }
406                }
407                Err(e) => Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(e))),
408            }
409        } else {
410            Ok(crate::traits::CalcValue::Scalar(LiteralValue::Error(
411                ExcelError::new(ExcelErrorKind::Ref),
412            )))
413        }
414    }
415}
416
417pub fn register_builtins() {
418    crate::function_registry::register_function(std::sync::Arc::new(IndexFn));
419    crate::function_registry::register_function(std::sync::Arc::new(OffsetFn));
420}
421
422#[cfg(test)]
423mod tests {
424    use super::*;
425    use crate::test_workbook::TestWorkbook;
426    use crate::traits::ArgumentHandle;
427    use formualizer_parse::parser::{ASTNode, ASTNodeType};
428
429    fn interp(wb: &TestWorkbook) -> crate::interpreter::Interpreter<'_> {
430        wb.interpreter()
431    }
432
433    #[test]
434    fn index_returns_reference_and_materializes_in_value_context() {
435        let wb = TestWorkbook::new()
436            .with_cell_a1("Sheet1", "B2", LiteralValue::Int(42))
437            .with_function(std::sync::Arc::new(IndexFn));
438        let ctx = interp(&wb);
439
440        // Build INDEX(A1:C3,2,2) expecting B2
441        let array_ref = ASTNode::new(
442            ASTNodeType::Reference {
443                original: "A1:C3".into(),
444                reference: ReferenceType::Range {
445                    sheet: None,
446                    start_row: Some(1),
447                    start_col: Some(1),
448                    end_row: Some(3),
449                    end_col: Some(3),
450                    start_row_abs: false,
451                    start_col_abs: false,
452                    end_row_abs: false,
453                    end_col_abs: false,
454                },
455            },
456            None,
457        );
458        let row = ASTNode::new(ASTNodeType::Literal(LiteralValue::Int(2)), None);
459        let col = ASTNode::new(ASTNodeType::Literal(LiteralValue::Int(2)), None);
460        let call = ASTNode::new(
461            ASTNodeType::Function {
462                name: "INDEX".into(),
463                args: vec![array_ref.clone(), row.clone(), col.clone()],
464            },
465            None,
466        );
467
468        // Reference context
469        let r = ctx.evaluate_ast_as_reference(&call).expect("ref ok");
470        match r {
471            ReferenceType::Cell { row, col, .. } => {
472                assert_eq!((row, col), (2, 2));
473            }
474            _ => panic!(),
475        }
476
477        // Value context (scalar materialization)
478        let args = vec![
479            ArgumentHandle::new(&array_ref, &ctx),
480            ArgumentHandle::new(&row, &ctx),
481            ArgumentHandle::new(&col, &ctx),
482        ];
483        let f = ctx.context.get_function("", "INDEX").unwrap();
484        let v = f
485            .dispatch(&args, &ctx.function_context(None))
486            .unwrap()
487            .into_literal();
488        assert_eq!(v, LiteralValue::Number(42.0));
489    }
490
491    #[test]
492    fn offset_returns_reference_and_materializes() {
493        let wb = TestWorkbook::new()
494            .with_cell_a1("Sheet1", "A1", LiteralValue::Int(1))
495            .with_cell_a1("Sheet1", "B2", LiteralValue::Int(5))
496            .with_function(std::sync::Arc::new(OffsetFn));
497        let ctx = interp(&wb);
498
499        let base = ASTNode::new(
500            ASTNodeType::Reference {
501                original: "A1".into(),
502                reference: ReferenceType::Cell {
503                    sheet: None,
504                    row: 1,
505                    col: 1,
506                    row_abs: false,
507                    col_abs: false,
508                },
509            },
510            None,
511        );
512        let dr = ASTNode::new(ASTNodeType::Literal(LiteralValue::Int(1)), None);
513        let dc = ASTNode::new(ASTNodeType::Literal(LiteralValue::Int(1)), None);
514        let call = ASTNode::new(
515            ASTNodeType::Function {
516                name: "OFFSET".into(),
517                args: vec![base.clone(), dr.clone(), dc.clone()],
518            },
519            None,
520        );
521
522        let r = ctx.evaluate_ast_as_reference(&call).expect("ref ok");
523        match r {
524            ReferenceType::Cell { row, col, .. } => assert_eq!((row, col), (2, 2)),
525            _ => panic!(),
526        }
527
528        let args = vec![
529            ArgumentHandle::new(&base, &ctx),
530            ArgumentHandle::new(&dr, &ctx),
531            ArgumentHandle::new(&dc, &ctx),
532        ];
533        let f = ctx.context.get_function("", "OFFSET").unwrap();
534        let v = f
535            .dispatch(&args, &ctx.function_context(None))
536            .unwrap()
537            .into_literal();
538        assert_eq!(v, LiteralValue::Number(5.0));
539    }
540}