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 RETURN ergebnis',
};

describe('expandUserFunctions (AST-level): kind "block" -- FOR loops', () => {
  it('a single FOR loop unrolls into one nested expression, left-to-right', () => {
    const index = buildFunctionIndex([def({ id: 'f1', name: 'PRODUCT2', ...productBlock })]);
    const expanded = expandUserFunctions(parseFormula('=PRODUCT2(A1:A3)'), index);
    // ((1 * A1) * A2) * A3
    expect(expanded).toEqual({
      kind: 'binary',
      op: '*',
      left: {
        kind: 'binary',
        op: '*',
        left: {
          kind: 'binary',
          op: '*',
          left: { kind: 'number', value: 1 },
          right: { kind: 'cell', qualifier: null, col: 0, row: 0 },
        },
        right: { kind: 'cell', qualifier: null, col: 0, row: 1 },
      },
      right: { kind: 'cell', qualifier: null, col: 0, row: 2 },
    });
  });

  it('a single-cell FOR source (no range) unrolls to exactly one iteration', () => {
    const index = buildFunctionIndex([def({ id: 'f1', name: 'PRODUCT2', ...productBlock })]);
    const expanded = expandUserFunctions(parseFormula('=PRODUCT2(A1)'), index);
    expect(expanded).toEqual({
      kind: 'binary',
      op: '*',
      left: { kind: 'number', value: 1 },
      right: { kind: 'cell', qualifier: null, col: 0, row: 0 },
    });
  });

  it('a loop variable / local stays visible (holding its last value) after its own loop', () => {
    const index = buildFunctionIndex([
      def({
        id: 'f1',
        name: 'LASTPLUSONE',
        params: ['bereich'],
        kind: 'block',
        body: 'acc = 0 FOR x IN bereich DO acc = acc + x RETURN x + 1', // returns the loop VAR, not acc
      }),
    ]);
    const expanded = expandUserFunctions(parseFormula('=LASTPLUSONE(A1:A3)'), index);
    expect(expanded).toEqual({
      kind: 'binary',
      op: '+',
      left: { kind: 'cell', qualifier: null, col: 0, row: 2 }, // A3 -- the last iterated cell
      right: { kind: 'number', value: 1 },
    });
  });

  it('a call inside a FOR body to another (expression-kind) function is expanded too', () => {
    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) RETURN total',
      }),
    ]);
    const expanded = expandUserFunctions(parseFormula('=SUMDOUBLED(A1:A2)'), index);
    // (0 + (A1*2)) + (A2*2) -- DOUBLE() fully inlined per iteration, no func nodes left
    expect(expanded).toEqual({
      kind: 'binary',
      op: '+',
      left: {
        kind: 'binary',
        op: '+',
        left: { kind: 'number', value: 0 },
        right: {
          kind: 'binary',
          op: '*',
          left: { kind: 'cell', qualifier: null, col: 0, row: 0 },
          right: { kind: 'number', value: 2 },
        },
      },
      right: {
        kind: 'binary',
        op: '*',
        left: { kind: 'cell', qualifier: null, col: 0, row: 1 },
        right: { kind: 'number', value: 2 },
      },
    });
  });

  it('multiple independent FOR loops each accumulate correctly', () => {
    const index = buildFunctionIndex([
      def({
        id: 'c1',
        name: 'COMBINED',
        params: ['sumRange', 'prodRange'],
        kind: 'block',
        body:
          'total = 0 FOR a IN sumRange DO total = total + a ' +
          'prod = 1 FOR b IN prodRange DO prod = prod * b ' +
          'RETURN total + prod',
      }),
    ]);
    const expanded = expandUserFunctions(parseFormula('=COMBINED(A1:A2; B1:B2)'), index);
    expect(expanded).toEqual({
      kind: 'binary',
      op: '+',
      left: {
        kind: 'binary',
        op: '+',
        left: { kind: 'binary', op: '+', left: { kind: 'number', value: 0 }, right: { kind: 'cell', qualifier: null, col: 0, row: 0 } },
        right: { kind: 'cell', qualifier: null, col: 0, row: 1 },
      },
      right: {
        kind: 'binary',
        op: '*',
        left: { kind: 'binary', op: '*', left: { kind: 'number', value: 1 }, right: { kind: 'cell', qualifier: null, col: 1, row: 0 } },
        right: { kind: 'cell', qualifier: null, col: 1, row: 1 },
      },
    });
  });

  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();
  });

  it('a FOR source that is not a cell/range after substitution throws', () => {
    const index = buildFunctionIndex([
      def({
        id: 'f1',
        name: 'BAD',
        params: ['x'],
        kind: 'block',
        body: 'total = 0 FOR y IN x + 1 DO total = total + y RETURN total',
      }),
    ]);
    // x substituted with a NumberNode -> the FOR source becomes "5 + 1", not a cell/range
    expect(() => expandUserFunctions(parseFormula('=BAD(5)'), index)).toThrow();
  });

  it('a FOR loop over more cells than the unroll cap throws', () => {
    const index = buildFunctionIndex([def({ id: 'f1', name: 'PRODUCT2', ...productBlock })]);
    // A1:A20000 -> 20000 cells, well over the 10,000-cell cap
    expect(() => expandUserFunctions(parseFormula('=PRODUCT2(A1:A20000)'), index)).toThrow();
  });
});

describe('Workbook integration: block-kind (FOR loop) 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('exceeding the unroll 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 RETURN total',
        }),
      ],
    );
    expect((val(wb, 's-umsatz', 'B9') as FormulaError).code).toBe('#ERR');
  });
});