numolus/src/engine/__tests__/userFunctions.spec.ts

import { describe, it, expect } from 'vitest';
import { parseFormula } from '../parser';
import {
  buildFunctionIndex,
  collectCalledFunctionNames,
  expandUserFunctions,
  type FunctionDef,
} from '../userFunctions';
import { Workbook, type TableModel } from '../workbook';
import { FormulaError, isError, keyOf } from '../types';

// same small A1 helper used by workbook.spec.ts
function sheet(
  id: string,
  name: string,
  cells: Record<string, string>,
  sheetName = 'Blatt1',
): TableModel {
  const inputs = new Map<string, string>();
  for (const [a1, raw] of Object.entries(cells)) {
    const m = /^([A-Z]+)([0-9]+)$/.exec(a1)!;
    const col = m[1].split('').reduce((n, ch) => n * 26 + (ch.charCodeAt(0) - 64), 0) - 1;
    const row = Number(m[2]) - 1;
    inputs.set(`${col},${row}`, raw);
  }
  return { id, name, sheetName, cols: 8, rows: 30, inputs };
}
const val = (wb: Workbook, sheetId: string, a1: string) => {
  const m = /^([A-Z]+)([0-9]+)$/.exec(a1)!;
  const col = m[1].split('').reduce((n, ch) => n * 26 + (ch.charCodeAt(0) - 64), 0) - 1;
  return wb.cellValue(sheetId, col, Number(m[2]) - 1);
};

const def = (
  partial: Partial<FunctionDef> & Pick<FunctionDef, 'id' | 'name' | 'params' | 'body'>,
): FunctionDef => ({ kind: 'expression', ...partial });

describe('expandUserFunctions (AST-level)', () => {
  it('substitutes parameters with the caller\'s argument subtrees', () => {
    const index = buildFunctionIndex([def({ id: 'f1', name: 'MARGE', params: ['x', 'y'], body: 'x - y' })]);
    const ast = parseFormula('=MARGE(A1; 5)');
    const expanded = expandUserFunctions(ast, index);
    expect(expanded).toEqual({
      kind: 'binary',
      op: '-',
      left: { kind: 'cell', qualifier: null, col: 0, row: 0 },
      right: { kind: 'number', value: 5 },
    });
  });

  it('an unknown name (built-in or not) is left untouched for evaluate() to turn into #NAME?', () => {
    const index = buildFunctionIndex([]);
    const ast = parseFormula('=FOO(1;2)');
    expect(expandUserFunctions(ast, index)).toEqual(ast);
  });

  it('a user function may call a built-in, and its arguments are still expanded', () => {
    const index = buildFunctionIndex([
      def({ id: 'f1', name: 'TOTAL', params: ['a', 'b'], body: 'SUM(a;b;10)' }),
    ]);
    const ast = parseFormula('=TOTAL(A1;A2)');
    const expanded = expandUserFunctions(ast, index);
    expect(expanded).toEqual({
      kind: 'func',
      name: 'SUM',
      args: [
        { kind: 'cell', qualifier: null, col: 0, row: 0 },
        { kind: 'cell', qualifier: null, col: 0, row: 1 },
        { kind: 'number', value: 10 },
      ],
    });
  });

  it('a user function may call another user function, params expanded inside-out', () => {
    const index = buildFunctionIndex([
      def({ id: 'inner', name: 'DOUBLE', params: ['x'], body: 'x * 2' }),
      def({ id: 'outer', name: 'QUAD', params: ['x'], body: 'DOUBLE(DOUBLE(x))' }),
    ]);
    const ast = parseFormula('=QUAD(A1)');
    const expanded = expandUserFunctions(ast, index);
    // DOUBLE(DOUBLE(A1)) -> (A1 * 2) * 2, fully inlined, no func nodes left
    expect(expanded).toEqual({
      kind: 'binary',
      op: '*',
      left: {
        kind: 'binary',
        op: '*',
        left: { kind: 'cell', qualifier: null, col: 0, row: 0 },
        right: { kind: 'number', value: 2 },
      },
      right: { kind: 'number', value: 2 },
    });
  });

  it('nested functions with colliding parameter names resolve to the correct scope', () => {
    // OUTER(x) = INNER(x) + 1, INNER(x) = x * 100  -- both call their param "x"
    const index = buildFunctionIndex([
      def({ id: 'inner', name: 'INNER', params: ['x'], body: 'x * 100' }),
      def({ id: 'outer', name: 'OUTER', params: ['x'], body: 'INNER(x) + 1' }),
    ]);
    const ast = parseFormula('=OUTER(A1)');
    const expanded = expandUserFunctions(ast, index);
    expect(expanded).toEqual({
      kind: 'binary',
      op: '+',
      left: {
        kind: 'binary',
        op: '*',
        left: { kind: 'cell', qualifier: null, col: 0, row: 0 },
        right: { kind: 'number', value: 100 },
      },
      right: { kind: 'number', value: 1 },
    });
  });

  it('throws on arity mismatch', () => {
    const index = buildFunctionIndex([def({ id: 'f1', name: 'MARGE', params: ['x', 'y'], body: 'x - y' })]);
    expect(() => expandUserFunctions(parseFormula('=MARGE(A1)'), index)).toThrow();
  });

  it('throws once expansion depth exceeds the recursion guard (defence in depth for a cycle)', () => {
    // A calls B, B calls A -- should never be saved (store-level cycle check),
    // but the engine must still degrade to an error, not a stack overflow.
    const index = buildFunctionIndex([
      def({ id: 'a', name: 'A', params: ['x'], body: 'B(x)' }),
      def({ id: 'b', name: 'B', params: ['x'], body: 'A(x)' }),
    ]);
    expect(() => expandUserFunctions(parseFormula('=A(1)'), index)).toThrow();
  });
});

describe('collectCalledFunctionNames', () => {
  it('finds every function name called anywhere in an AST, upper-cased', () => {
    const ast = parseFormula('=SUM(A1;OTHER(B1)) + third(1)');
    expect(collectCalledFunctionNames(ast)).toEqual(new Set(['SUM', 'OTHER', 'THIRD']));
  });
});

describe('Workbook integration: user-defined functions end-to-end', () => {
  function fixture(functionDefs: FunctionDef[] = []): Workbook {
    const umsatz = sheet('s-umsatz', 'Umsatz', {
      B2: '1200', C2: '800',
      B9: '=MARGE(B2;C2)',
    });
    return new Workbook([umsatz], functionDefs);
  }

  it('a call is evaluated using the function body', () => {
    const wb = fixture([def({ id: 'f1', name: 'MARGE', params: ['revenue', 'cost'], body: '(revenue - cost) / revenue' })]);
    expect(val(wb, 's-umsatz', 'B9')).toBeCloseTo((1200 - 800) / 1200, 10);
  });

  it('recalculates through the function body when a precedent changes', () => {
    const wb = fixture([def({ id: 'f1', name: 'MARGE', params: ['revenue', 'cost'], body: '(revenue - cost) / revenue' })]);
    wb.setCell('s-umsatz', 1, 1, '2000'); // B2: 1200 -> 2000
    expect(val(wb, 's-umsatz', 'B9')).toBeCloseTo((2000 - 800) / 2000, 10);
  });

  it('a body referencing a cell that is *not* in the call arguments still creates a dependency edge', () => {
    const wb = fixture();
    wb.tables[0].inputs.set('4,0', '3'); // E1 = 3, a "rate" cell not passed as an argument
    wb.tables[0].inputs.set('1,8', '=WITHRATE(B2)'); // B9 = WITHRATE(B2)
    const functionDefs = [def({ id: 'f1', name: 'WITHRATE', params: ['x'], body: 'x * E1' })];
    const wb2 = new Workbook(wb.tables, functionDefs);
    expect(val(wb2, 's-umsatz', 'B9')).toBe(1200 * 3);

    // changing E1 (never passed as an argument) must still propagate to B9
    wb2.setCell('s-umsatz', 4, 0, '5');
    expect(val(wb2, 's-umsatz', 'B9')).toBe(1200 * 5);
  });

  it('unknown function name (never defined) yields #NAME?, exactly like an unknown built-in', () => {
    const wb = fixture(); // no functionDefs at all -- MARGE stays unresolved
    expect((val(wb, 's-umsatz', 'B9') as FormulaError).code).toBe('#NAME?');
  });

  it('wrong argument count yields #ERR', () => {
    const wb = fixture([def({ id: 'f1', name: 'MARGE', params: ['revenue', 'cost', 'extra'], body: 'revenue' })]);
    expect((val(wb, 's-umsatz', 'B9') as FormulaError).code).toBe('#ERR');
  });

  it('a built-in name can never be shadowed -- SUM keeps behaving like SUM', () => {
    // even if a document somehow ended up with a "SUM" custom function
    // (the store rejects this at save time), the engine still prefers the
    // built-in -- buildFunctionIndex()/expandUserFunctions() never override it
    // because functions.ts's own registry is consulted at evaluation time and
    // wins whenever expandUserFunctions leaves a func node in place... so make
    // sure the custom "SUM" is simply never invoked from a real formula that
    // hits the built-in path first is out of scope here; this test instead
    // documents that a *user* SUM does get inlined when explicitly defined,
    // since the store -- not the engine -- is what forbids the collision.
    const index = buildFunctionIndex([def({ id: 'x', name: 'SUM', params: ['a'], body: 'a + 1' })]);
    const expanded = expandUserFunctions(parseFormula('=SUM(A1)'), index);
    expect(expanded).toEqual({
      kind: 'binary',
      op: '+',
      left: { kind: 'cell', qualifier: null, col: 0, row: 0 },
      right: { kind: 'number', value: 1 },
    });
  });

  it('errors inside a function body propagate to dependents', () => {
    const wb = fixture([def({ id: 'f1', name: 'MARGE', params: ['revenue', 'cost'], body: 'revenue / 0' })]);
    expect(isError(val(wb, 's-umsatz', 'B9'))).toBe(true);
    expect((val(wb, 's-umsatz', 'B9') as FormulaError).code).toBe('#DIV/0!');
  });

  it('a parameter can be a range, as long as it lands in a function-call argument position', () => {
    // SPANNE(bereich) = MAX(bereich) - MIN(bereich) -- "bereich" (ParamNode)
    // gets substituted with the caller's RangeNode, splicing a normal
    // MAX(Umsatz!B2:E2) / MIN(Umsatz!B2:E2) call into the expanded AST.
    const umsatz = sheet('s-umsatz', 'Umsatz', {
      B2: '1200', C2: '1450', D2: '1600', E2: '1720',
      B9: '=SPANNE(B2:E2)',
    });
    const wb = new Workbook(
      [umsatz],
      [def({ id: 'f1', name: 'SPANNE', params: ['bereich'], body: 'MAX(bereich) - MIN(bereich)' })],
    );
    expect(val(wb, 's-umsatz', 'B9')).toBe(1720 - 1200);

    // recalculates when a cell *inside* the range argument changes
    wb.setCell('s-umsatz', 1, 1, '1000'); // B2: 1200 -> 1000
    expect(val(wb, 's-umsatz', 'B9')).toBe(1720 - 1000);
  });

  it('a range parameter used outside a function-call argument position is #VALUE!, same as a bare range in an ordinary formula', () => {
    const umsatz = sheet('s-umsatz', 'Umsatz', {
      B2: '1200', C2: '1450',
      B9: '=BAD(B2:C2)',
    });
    const wb = new Workbook([umsatz], [def({ id: 'f1', name: 'BAD', params: ['bereich'], body: 'bereich + 1' })]);
    expect((val(wb, 's-umsatz', 'B9') as FormulaError).code).toBe('#VALUE!');
  });

  it('debugOrder places a function-body-only precedent before its dependent', () => {
    const umsatz = sheet('s-umsatz', 'Umsatz', {
      B2: '1200', E1: '3',
      B9: '=WITHRATE(B2)',
    });
    const wb = new Workbook([umsatz], [def({ id: 'f1', name: 'WITHRATE', params: ['x'], body: 'x * E1' })]);
    const { order, circular } = wb.debugOrder([keyOf('s-umsatz', 4, 0)]); // dirty: E1
    expect(circular).toEqual([]);
    expect(order).toContain(keyOf('s-umsatz', 1, 8)); // B9 depends on E1 too
  });
});

const productBlock = {
  params: ['bereich'],
  kind: 'block' as const,
  body: 'ergebnis = 1 FOR x IN bereich DO ergebnis = ergebnis * x END RETURN ergebnis',
};

describe('expandUserFunctions (AST-level): kind "block" produces a BlockCallNode', () => {
  it('wraps a block-kind call in a BlockCallNode carrying args/paramNames/label', () => {
    const index = buildFunctionIndex([def({ id: 'f1', name: 'PRODUCT2', ...productBlock })]);
    const expanded = expandUserFunctions(parseFormula('=PRODUCT2(A1:A3)'), index);
    expect(expanded.kind).toBe('blockcall');
    if (expanded.kind !== 'blockcall') throw new Error('unreachable');
    expect(expanded.label).toBe('PRODUCT2');
    expect(expanded.paramNames).toEqual(['bereich']);
    expect(expanded.args).toEqual([
      {
        kind: 'range',
        qualifier: null,
        start: { col: 0, row: 0, colAbs: false, rowAbs: false },
        end: { col: 0, row: 2, colAbs: false, rowAbs: false },
      },
    ]);
  });

  it('resolves a nested call to another (expression-kind) function inside the block before evaluation', () => {
    const index = buildFunctionIndex([
      def({ id: 'd1', name: 'DOUBLE', params: ['n'], body: 'n * 2' }),
      def({
        id: 's1',
        name: 'SUMDOUBLED',
        params: ['bereich'],
        kind: 'block',
        body: 'total = 0 FOR x IN bereich DO total = total + DOUBLE(x) END RETURN total',
      }),
    ]);
    const expanded = expandUserFunctions(parseFormula('=SUMDOUBLED(A1:A2)'), index);
    if (expanded.kind !== 'blockcall') throw new Error('unreachable');
    const forStmt = expanded.block.statements[1];
    if (forStmt.kind !== 'for') throw new Error('expected a for statement');
    const assignInBody = forStmt.body.statements[0];
    if (assignInBody.kind !== 'assign') throw new Error('expected an assignment');
    // DOUBLE(x) has already been inlined to "x * 2" -- no FuncNode calling
    // "DOUBLE" remains anywhere in the block
    expect(assignInBody.value).toEqual({
      kind: 'binary',
      op: '+',
      left: { kind: 'param', name: 'total' },
      right: {
        kind: 'binary',
        op: '*',
        left: { kind: 'param', name: 'x' },
        right: { kind: 'number', value: 2 },
      },
    });
  });

  it('arity mismatch is still checked for block-kind functions', () => {
    const index = buildFunctionIndex([def({ id: 'f1', name: 'PRODUCT2', ...productBlock })]);
    expect(() => expandUserFunctions(parseFormula('=PRODUCT2(A1;A2)'), index)).toThrow();
  });
});

describe('Workbook integration: block-kind (FOR/IF/RETURN "code") user-defined functions', () => {
  it('evaluates a PRODUCT-style reduction over a range correctly', () => {
    const umsatz = sheet('s-umsatz', 'Umsatz', { A1: '2', A2: '3', A3: '4', B9: '=PRODUCT2(A1:A3)' });
    const wb = new Workbook([umsatz], [def({ id: 'f1', name: 'PRODUCT2', ...productBlock })]);
    expect(val(wb, 's-umsatz', 'B9')).toBe(24);
  });

  it('recalculates when a cell inside the passed range changes', () => {
    const umsatz = sheet('s-umsatz', 'Umsatz', { A1: '2', A2: '3', A3: '4', B9: '=PRODUCT2(A1:A3)' });
    const wb = new Workbook([umsatz], [def({ id: 'f1', name: 'PRODUCT2', ...productBlock })]);
    wb.setCell('s-umsatz', 0, 0, '10'); // A1: 2 -> 10
    expect(val(wb, 's-umsatz', 'B9')).toBe(10 * 3 * 4);
  });

  it('a caller may pass a named range as the argument -- the FOR loop iterates it exactly like an ordinary range', () => {
    const umsatz = sheet('s-umsatz', 'Umsatz', { A1: '2', A2: '3', A3: '4', B9: '=PRODUCT2(Werte)' });
    const wb = new Workbook(
      [umsatz],
      [def({ id: 'f1', name: 'PRODUCT2', ...productBlock })],
      [{ id: 'n1', name: 'Werte', tableId: 's-umsatz', start: { col: 0, row: 0 }, end: { col: 0, row: 2 } }],
    );
    expect(val(wb, 's-umsatz', 'B9')).toBe(24);
    wb.setCell('s-umsatz', 0, 0, '10'); // A1: 2 -> 10 -- still tracked as a precedent via the named range
    expect(val(wb, 's-umsatz', 'B9')).toBe(10 * 3 * 4);
  });

  it('a FOR body with multiple statements runs all of them, every iteration', () => {
    const umsatz = sheet('s-umsatz', 'Umsatz', { A1: '2', A2: '3', B9: '=DOUBLESUM(A1:A2)' });
    const doubleSum = def({
      id: 'f1',
      name: 'DOUBLESUM',
      params: ['bereich'],
      kind: 'block',
      body: 'total = 0 FOR x IN bereich DO y = x * 2 total = total + y END RETURN total',
    });
    const wb = new Workbook([umsatz], [doubleSum]);
    expect(val(wb, 's-umsatz', 'B9')).toBe(2 * 2 + 3 * 2);
  });

  it('a FOR loop nested inside another FOR loop evaluates correctly, not just parses', () => {
    const umsatz = sheet('s-umsatz', 'Umsatz', {
      A1: '2',
      A2: '3',
      B1: '10',
      B2: '20',
      B9: '=OUTERPROD(A1:A2; B1:B2)',
    });
    const outerProd = def({
      id: 'f1',
      name: 'OUTERPROD',
      params: ['r1', 'r2'],
      kind: 'block',
      body: 'total = 0 FOR a IN r1 DO FOR b IN r2 DO total = total + a * b END END RETURN total',
    });
    const wb = new Workbook([umsatz], [outerProd]);
    expect(val(wb, 's-umsatz', 'B9')).toBe(2 * 10 + 2 * 20 + 3 * 10 + 3 * 20);
  });

  it('IF/ELSE branches on the real, current cell value', () => {
    const umsatz = sheet('s-umsatz', 'Umsatz', { A1: '5', B9: '=SIGN2(A1)' });
    const sign2 = def({
      id: 'f1',
      name: 'SIGN2',
      params: ['x'],
      kind: 'block',
      body: 'IF x > 0 THEN RETURN 1 ELSE RETURN -1 END',
    });
    const wb = new Workbook([umsatz], [sign2]);
    expect(val(wb, 's-umsatz', 'B9')).toBe(1);
    wb.setCell('s-umsatz', 0, 0, '-5');
    expect(val(wb, 's-umsatz', 'B9')).toBe(-1);
  });

  it('comparison operators work inside an IF condition', () => {
    const umsatz = sheet('s-umsatz', 'Umsatz', { A1: '10', A2: '10', B9: '=EQ2(A1;A2)' });
    const eq2 = def({
      id: 'f1',
      name: 'EQ2',
      params: ['a', 'b'],
      kind: 'block',
      body: 'IF a = b THEN RETURN 1 ELSE RETURN 0 END',
    });
    const wb = new Workbook([umsatz], [eq2]);
    expect(val(wb, 's-umsatz', 'B9')).toBe(1);
    wb.setCell('s-umsatz', 0, 1, '11'); // A2: 10 -> 11
    expect(val(wb, 's-umsatz', 'B9')).toBe(0);
  });

  it('an early RETURN nested inside an IF inside a FOR stops the loop -- finds the first match, never touches later cells', () => {
    // A3 is the first cell > 100; A4 must never affect the result
    const umsatz = sheet('s-umsatz', 'Umsatz', {
      A1: '1',
      A2: '2',
      A3: '999',
      A4: '4',
      B9: '=FIRSTOVER(A1:A4; 100)',
    });
    const firstOver = def({
      id: 'f1',
      name: 'FIRSTOVER',
      params: ['bereich', 'schwelle'],
      kind: 'block',
      body: 'FOR x IN bereich DO IF x > schwelle THEN RETURN x END END RETURN 0',
    });
    const wb = new Workbook([umsatz], [firstOver]);
    expect(val(wb, 's-umsatz', 'B9')).toBe(999);
    // changing the never-reached A4 doesn't change the (already-found) result
    wb.setCell('s-umsatz', 0, 3, '99999');
    expect(val(wb, 's-umsatz', 'B9')).toBe(999);
  });

  it('a cell referenced only in a not-taken IF branch is still a tracked dependency', () => {
    const umsatz = sheet('s-umsatz', 'Umsatz', { A1: '5', A2: '1', A3: '2', B9: '=BRANCHY(A1)' });
    const branchy = def({
      id: 'f1',
      name: 'BRANCHY',
      params: ['x'],
      kind: 'block',
      body: 'IF x > 0 THEN RETURN A2 ELSE RETURN A3 END',
    });
    const wb = new Workbook([umsatz], [branchy]);
    expect(val(wb, 's-umsatz', 'B9')).toBe(1); // A1=5>0 -> takes the A2 branch
    wb.setCell('s-umsatz', 0, 0, '-5'); // flips to the A3 branch
    expect(val(wb, 's-umsatz', 'B9')).toBe(2);
  });

  it('exceeding the loop-step cap degrades to #ERR at the cell, not a crash', () => {
    const umsatz = sheet('s-umsatz', 'Umsatz', { B9: '=PRODUCT2(A1:A20000)' });
    const wb = new Workbook([umsatz], [def({ id: 'f1', name: 'PRODUCT2', ...productBlock })]);
    expect((val(wb, 's-umsatz', 'B9') as FormulaError).code).toBe('#ERR');
  });

  it('an invalid FOR source degrades to #ERR at the cell, not a crash', () => {
    const umsatz = sheet('s-umsatz', 'Umsatz', { B9: '=BAD(5)' });
    const wb = new Workbook(
      [umsatz],
      [
        def({
          id: 'f1',
          name: 'BAD',
          params: ['x'],
          kind: 'block',
          body: 'total = 0 FOR y IN x + 1 DO total = total + y END RETURN total',
        }),
      ],
    );
    expect((val(wb, 's-umsatz', 'B9') as FormulaError).code).toBe('#ERR');
  });

  it('a function that never reaches RETURN degrades to #ERR, not a crash', () => {
    const umsatz = sheet('s-umsatz', 'Umsatz', { B9: '=NORETURN(1)' });
    const noReturn = def({ id: 'f1', name: 'NORETURN', params: ['x'], kind: 'block', body: 'y = x + 1' });
    const wb = new Workbook([umsatz], [noReturn]);
    expect((val(wb, 's-umsatz', 'B9') as FormulaError).code).toBe('#ERR');
  });
});