numolus/src/engine/userFunctions.ts

/**
 * User-defined functions: a document-level library of named, parameterised
 * formulas (`MARGE(revenue; cost) = (revenue - cost) / revenue`) callable
 * from any cell like a built-in (`=MARGE(Umsatz!B2; Umsatz!C2)`).
 *
 * The evaluator (`evaluator.ts`) evaluates function arguments *eagerly* --
 * scalar `CellValue`s in, `CellValue` out (`FunctionFn` in `functions.ts`) --
 * which has no room for a "body formula with its own cell references and a
 * parameter environment". Rather than rework the evaluator, a `kind:
 * 'expression'` call (a single-formula body) is *inlined* into the caller's
 * AST at parse time, before the graph/evaluator ever see it: the `FuncNode`
 * is replaced by a (recursively expanded) copy of the function body, with
 * every `ParamNode` swapped for the caller's own argument subtree. After
 * that, the stored AST looks exactly as if the user had written the
 * expanded formula by hand -- `collectRefs`/`precedentsOf`/`evaluate` all
 * work completely unmodified, and correctly pick up cell references that
 * live *only* inside the function body, not just the ones passed as
 * arguments.
 *
 * A `kind: 'block'` call (a FOR/IF/RETURN "code" body, see parser.ts's
 * `parseFunctionBlock`) is different: this module only resolves *nested*
 * function calls inside it (`expandBlock`, below) -- still at parse time,
 * still the same "inside-out" inlining -- and wraps the result in a
 * `BlockCallNode`. The block's *own* control flow (FOR iteration, IF
 * branching, an early RETURN) is only resolved at evaluation time, by a
 * small interpreter in `evaluator.ts`, since it genuinely depends on real
 * cell values. See src/engine/README.md's "Schleifen" section for the full
 * design and why that split is necessary.
 */

import {
  parseFormula,
  parseFunctionBlock,
  type BlockNode,
  type ForNode,
  type Node,
} from './parser';

export interface FunctionDef {
  id: string;
  /** document-unique, matched case-insensitively (like built-in names). */
  name: string;
  /** parameter names, in call order. */
  params: string[];
  /** raw formula text, no leading "=" -- a single expression for
   *  `kind: 'expression'`, a statement sequence (assignments, FOR/DO/END,
   *  IF/THEN/ELSE/END, RETURN anywhere) for `kind: 'block'`. */
  body: string;
  /** 'expression' (the original, Phase-1 form) or 'block' (FOR/IF/RETURN
   *  "code"). Every function predating this field is 'expression'. */
  kind: 'expression' | 'block';
}

/** Recursion guard for `expandUserFunctions` -- a cycle should always be
 *  rejected at definition time (see the store's cycle check), this is only
 *  a last-resort safety net so a bug there degrades to a cell error rather
 *  than a stack overflow. */
const MAX_EXPANSION_DEPTH = 32;

export function buildFunctionIndex(defs: readonly FunctionDef[]): Map<string, FunctionDef> {
  return new Map(defs.map((d) => [d.name.toUpperCase(), d]));
}

/**
 * Replace every call to a user-defined function in `ast` with either a
 * (recursively expanded, parameter-substituted) copy of its body
 * (`kind: 'expression'`), or a `BlockCallNode` wrapping its (nested-calls-
 * resolved) block for evaluator.ts's interpreter to run (`kind: 'block'`).
 * Calls to built-in or genuinely unknown names are left untouched --
 * `evaluate()`'s existing `#NAME?` path already covers "unknown function",
 * built-in or not.
 *
 * Throws a plain `Error` (not `ParseError` -- there's nothing to parse here,
 * just structural problems: wrong argument count, or a definition cycle
 * `MAX_EXPANSION_DEPTH` didn't catch at save time) on failure; the caller
 * (`Workbook.parseCell`) wraps this the same way it already wraps
 * `parseFormula` errors, turning it into a `#ERR` cell value.
 */
export function expandUserFunctions(ast: Node, index: ReadonlyMap<string, FunctionDef>): Node {
  if (index.size === 0) return ast; // fast path: no user functions defined at all
  return expand(ast, index, new Map(), new Map(), 0);
}

function expand(
  node: Node,
  index: ReadonlyMap<string, FunctionDef>,
  bodyCache: Map<string, Node>,
  blockCache: Map<string, BlockNode>,
  depth: number,
): Node {
  switch (node.kind) {
    case 'number':
    case 'string':
    case 'cell':
    case 'range':
    case 'param':
      return node; // leaves -- nothing to expand
    case 'unary':
      return { ...node, operand: expand(node.operand, index, bodyCache, blockCache, depth) };
    case 'binary':
      return {
        ...node,
        left: expand(node.left, index, bodyCache, blockCache, depth),
        right: expand(node.right, index, bodyCache, blockCache, depth),
      };
    case 'blockcall':
      // never present in freshly-parsed input -- only expand() itself ever
      // produces one (see the 'func' case below) -- but the type checker
      // doesn't know that, so: a defensive leaf, nothing to do.
      return node;
    case 'func': {
      // expand the arguments regardless of whether this call is itself a
      // user function -- a built-in like SUM(...) may still take a
      // user-function call as one of its arguments.
      const args = node.args.map((a) => expand(a, index, bodyCache, blockCache, depth));
      const def = index.get(node.name.toUpperCase());
      if (!def) return { ...node, args };

      if (depth >= MAX_EXPANSION_DEPTH) {
        throw new Error(`'${def.name}' expansion too deep (likely a definition cycle)`);
      }
      if (args.length !== def.params.length) {
        throw new Error(
          `'${def.name}' expects ${def.params.length} argument(s), got ${args.length}`,
        );
      }

      if (def.kind === 'block') {
        let block = blockCache.get(def.id);
        if (!block) {
          block = parseFunctionBlock(def.body, def.params); // ParseError propagates -- caller wraps it
          blockCache.set(def.id, block);
        }
        // resolve any *other* function calls the block contains (built-ins
        // untouched, other user functions inlined/turned into their own
        // BlockCallNode) -- same "inside-out" principle as the
        // expression-kind path below. This function's *own* parameters
        // stay as ParamNode inside `block`; evaluator.ts's interpreter
        // binds them from `args` at evaluation time, since a FOR loop's
        // iteration count and an IF's branch both depend on real cell
        // values that aren't known yet here.
        const expandedBlock = expandBlock(block, index, bodyCache, blockCache, depth + 1);
        return { kind: 'blockcall', block: expandedBlock, args, paramNames: def.params, label: def.name };
      }

      let body = bodyCache.get(def.id);
      if (!body) {
        body = parseFormula(def.body, def.params); // ParseError propagates -- caller wraps it
        bodyCache.set(def.id, body);
      }
      // expand any calls *inside* the body first (it may call other user
      // functions -- possibly reusing the same parameter names as this one;
      // resolving inside-out like this keeps each function's own parameters
      // scoped correctly no matter what the caller's happen to be named),
      // then substitute this function's own parameters last.
      const expandedBody = expand(body, index, bodyCache, blockCache, depth + 1);
      const bound = new Map(def.params.map((p, i) => [p.toUpperCase(), args[i]]));
      return substitute(expandedBody, bound);
    }
  }
}

/**
 * Replace every `ParamNode` in `node` matching a name in `bound` with its
 * bound subtree. Exported for evaluator.ts's block interpreter, which reuses
 * this exact mechanism to resolve names (parameters *and* local variables --
 * it doesn't distinguish the two, both are just entries in its environment
 * map) against the current execution state before evaluating an expression.
 */
export function substitute(node: Node, bound: ReadonlyMap<string, Node>): Node {
  switch (node.kind) {
    case 'param':
      // always found in practice (arity is checked before this runs, and
      // parameter names are exactly what the body was parsed against) --
      // the fallback just leaves an unresolved param for evaluate()'s own
      // defensive #ERR case rather than assuming it can't happen.
      return bound.get(node.name.toUpperCase()) ?? node;
    case 'number':
    case 'string':
    case 'cell':
    case 'range':
      return node;
    case 'unary':
      return { ...node, operand: substitute(node.operand, bound) };
    case 'binary':
      return { ...node, left: substitute(node.left, bound), right: substitute(node.right, bound) };
    case 'func':
      return { ...node, args: node.args.map((a) => substitute(a, bound)) };
    case 'blockcall':
      // only `.args` can still reference an *outer* scope's parameter (e.g.
      // this call sits inside another function's own body, referencing
      // *that* function's parameter) -- `.block` itself only refers to this
      // call's *own* parameters, a separate scope bound later, at
      // evaluation time, so it's left untouched here.
      return { ...node, args: node.args.map((a) => substitute(a, bound)) };
  }
}

/** {@link expand}, run over every leaf expression a `BlockNode` contains
 *  (each statement's value/condition/source, recursing into nested FOR/IF
 *  bodies) instead of over a single `Node` -- resolves nested function
 *  calls throughout a block-kind function's body before it's wrapped in a
 *  `BlockCallNode`. */
function expandBlock(
  block: BlockNode,
  index: ReadonlyMap<string, FunctionDef>,
  bodyCache: Map<string, Node>,
  blockCache: Map<string, BlockNode>,
  depth: number,
): BlockNode {
  return {
    kind: 'block',
    statements: block.statements.map((stmt) => {
      switch (stmt.kind) {
        case 'assign':
          return { ...stmt, value: expand(stmt.value, index, bodyCache, blockCache, depth) };
        case 'return':
          return { ...stmt, value: expand(stmt.value, index, bodyCache, blockCache, depth) };
        case 'for':
          return {
            ...stmt,
            source: expand(stmt.source, index, bodyCache, blockCache, depth),
            body: expandBlock(stmt.body, index, bodyCache, blockCache, depth),
          };
        case 'ifstmt':
          return {
            ...stmt,
            cond: expand(stmt.cond, index, bodyCache, blockCache, depth),
            then: expandBlock(stmt.then, index, bodyCache, blockCache, depth),
            else: stmt.else ? expandBlock(stmt.else, index, bodyCache, blockCache, depth) : null,
          };
      }
    }),
  };
}

/**
 * Every distinct function name called anywhere in `ast` (built-in or not --
 * callers filter against their own registry). Used by the store to build a
 * definition-time call graph between user functions for cycle detection,
 * without a second AST-walking implementation alongside `collectRefs`.
 */
export function collectCalledFunctionNames(ast: Node, out: Set<string> = new Set()): Set<string> {
  switch (ast.kind) {
    case 'func':
      out.add(ast.name.toUpperCase());
      ast.args.forEach((a) => collectCalledFunctionNames(a, out));
      break;
    case 'binary':
      collectCalledFunctionNames(ast.left, out);
      collectCalledFunctionNames(ast.right, out);
      break;
    case 'unary':
      collectCalledFunctionNames(ast.operand, out);
      break;
    default:
      break; // number/string/cell/range/param/blockcall -- no nested calls
  }
  return out;
}

/** Every name a block declares -- an assignment target or a loop variable --
 *  in the order first declared, recursing into nested FOR/IF bodies. Used
 *  by the store to validate these local names the same way it validates
 *  parameters (not cell-shaped, no collision with a parameter, not a
 *  reserved word); the engine itself has no opinion on naming policy or
 *  German error text, so it only exposes the raw names here. The *same*
 *  name legitimately reappears (an accumulator re-assigned inside a loop,
 *  or across two separate loops) -- that's not the store's concern either,
 *  it's intentional under the flat-scope design (see parser.ts). */
export function declaredBlockNames(block: BlockNode, out: string[] = []): string[] {
  for (const stmt of block.statements) {
    switch (stmt.kind) {
      case 'assign':
        out.push(stmt.name);
        break;
      case 'for':
        out.push(stmt.varName);
        declaredBlockNames(stmt.body, out);
        break;
      case 'ifstmt':
        declaredBlockNames(stmt.then, out);
        if (stmt.else) declaredBlockNames(stmt.else, out);
        break;
      case 'return':
        break;
    }
  }
  return out;
}

/** The first FOR loop, anywhere in the block (including nested inside
 *  another FOR/IF), whose source isn't (syntactically) a parameter, a
 *  single cell, or a range -- i.e. can never resolve to a concrete list of
 *  cells no matter what the caller passes. Checked at save time so this is
 *  caught immediately, not just on the first real call; the evaluator's
 *  interpreter re-checks the same thing on the *substituted* source at call
 *  time regardless (needed either way, since a parameter's shape isn't
 *  known until then) as the actual enforcement. */
export function findInvalidForSource(block: BlockNode): ForNode | null {
  for (const stmt of block.statements) {
    if (stmt.kind === 'for') {
      if (stmt.source.kind !== 'param' && stmt.source.kind !== 'cell' && stmt.source.kind !== 'range') {
        return stmt;
      }
      const nested = findInvalidForSource(stmt.body);
      if (nested) return nested;
    } else if (stmt.kind === 'ifstmt') {
      const nested = findInvalidForSource(stmt.then) ?? (stmt.else && findInvalidForSource(stmt.else));
      if (nested) return nested;
    }
  }
  return null;
}

/** {@link collectCalledFunctionNames}, for a `kind: 'block'` body: runs it
 *  over every leaf expression a `BlockNode` contains (each statement's
 *  value/condition/source, recursing into nested FOR/IF bodies) and merges
 *  the results -- `collectCalledFunctionNames` itself needs no changes,
 *  since a block is never part of the `Node` union it switches on. */
export function collectCalledFunctionNamesInBlock(
  block: BlockNode,
  out: Set<string> = new Set(),
): Set<string> {
  for (const stmt of block.statements) {
    switch (stmt.kind) {
      case 'assign':
      case 'return':
        collectCalledFunctionNames(stmt.value, out);
        break;
      case 'for':
        collectCalledFunctionNames(stmt.source, out);
        collectCalledFunctionNamesInBlock(stmt.body, out);
        break;
      case 'ifstmt':
        collectCalledFunctionNames(stmt.cond, out);
        collectCalledFunctionNamesInBlock(stmt.then, out);
        if (stmt.else) collectCalledFunctionNamesInBlock(stmt.else, out);
        break;
    }
  }
  return out;
}

/**
 * Conservative static check: does *every* execution path through `block`
 * reach a `return`? True if the block contains a `return` statement
 * (statements after it are unreachable but harmless), or an `IF` with an
 * `ELSE` where *both* branches always return. A `FOR` loop is never enough
 * on its own -- even if its body always returns when it runs, the loop
 * might cover zero cells at call time, which this static check can't see
 * (the interpreter's own "ended without RETURN" -> `#ERR` is the actual,
 * always-correct enforcement; this is only an early, save-time warning).
 */
export function blockAlwaysReturns(block: BlockNode): boolean {
  return block.statements.some((stmt) => {
    if (stmt.kind === 'return') return true;
    if (stmt.kind === 'ifstmt' && stmt.else) {
      return blockAlwaysReturns(stmt.then) && blockAlwaysReturns(stmt.else);
    }
    return false;
  });
}