Momento/memento-note/lib/interactive-page/sim-eval.ts

/**
 * Safe math-expression evaluator for the generic-formula simulator.
 * Tokenizer + recursive-descent parser — NO eval/Function, no object access,
 * no loops. Identifiers resolve only against an explicit environment;
 * functions come from a fixed allowlist.
 */

type Token =
  | { t: 'num'; v: number }
  | { t: 'id'; v: string }
  | { t: 'op'; v: string }
  | { t: 'lparen' }
  | { t: 'rparen' }
  | { t: 'comma' }

const CONSTANTS: Record<string, number> = {
  pi: Math.PI,
  e: Math.E,
}

type Fn = (...args: number[]) => number
const FUNCTIONS: Record<string, { fn: Fn; minArgs: number; maxArgs: number }> = {
  sqrt: { fn: Math.sqrt, minArgs: 1, maxArgs: 1 },
  abs: { fn: Math.abs, minArgs: 1, maxArgs: 1 },
  exp: { fn: Math.exp, minArgs: 1, maxArgs: 1 },
  ln: { fn: Math.log, minArgs: 1, maxArgs: 1 },
  log: { fn: Math.log10, minArgs: 1, maxArgs: 1 },
  round: { fn: Math.round, minArgs: 1, maxArgs: 1 },
  floor: { fn: Math.floor, minArgs: 1, maxArgs: 1 },
  ceil: { fn: Math.ceil, minArgs: 1, maxArgs: 1 },
  min: { fn: Math.min, minArgs: 1, maxArgs: 8 },
  max: { fn: Math.max, minArgs: 1, maxArgs: 8 },
}

const ID_RE = /^[A-Za-z_][A-Za-z0-9_]*$/

export type SimExprError = { message: string }

type Ast =
  | { k: 'num'; v: number }
  | { k: 'id'; v: string }
  | { k: 'call'; name: string; args: Ast[] }
  | { k: 'un'; op: '-'; a: Ast }
  | { k: 'bin'; op: string; a: Ast; b: Ast }

function tokenize(src: string): Token[] | SimExprError {
  const tokens: Token[] = []
  let i = 0
  while (i < src.length) {
    const ch = src[i]
    if (ch === ' ' || ch === '\t' || ch === '\n') {
      i++
      continue
    }
    if (ch >= '0' && ch <= '9') {
      let j = i
      while (j < src.length && /[0-9.]/.test(src[j])) j++
      const raw = src.slice(i, j)
      const v = Number(raw)
      if (!Number.isFinite(v)) return { message: `invalid number "${raw}"` }
      tokens.push({ t: 'num', v })
      i = j
      continue
    }
    if (ch === '.' && src[i + 1] >= '0' && src[i + 1] <= '9') {
      let j = i + 1
      while (j < src.length && /[0-9]/.test(src[j])) j++
      tokens.push({ t: 'num', v: Number(src.slice(i, j)) })
      i = j
      continue
    }
    if (/[A-Za-z_]/.test(ch)) {
      let j = i
      while (j < src.length && /[A-Za-z0-9_]/.test(src[j])) j++
      tokens.push({ t: 'id', v: src.slice(i, j) })
      i = j
      continue
    }
    if ('+-*/^%'.includes(ch)) {
      tokens.push({ t: 'op', v: ch })
      i++
      continue
    }
    if (ch === '(') {
      tokens.push({ t: 'lparen' })
      i++
      continue
    }
    if (ch === ')') {
      tokens.push({ t: 'rparen' })
      i++
      continue
    }
    if (ch === ',') {
      tokens.push({ t: 'comma' })
      i++
      continue
    }
    return { message: `unexpected character "${ch}"` }
  }
  return tokens
}

class Parser {
  private pos = 0
  constructor(private tokens: Token[]) {}

  private peek(): Token | undefined {
    return this.tokens[this.pos]
  }
  private next(): Token | undefined {
    return this.tokens[this.pos++]
  }
  private expectOp(): string | null {
    const t = this.peek()
    return t?.t === 'op' ? t.v : null
  }

  parseExpr(): Ast | SimExprError {
    let left = this.parseTerm()
    if ('message' in left) return left
    for (;;) {
      const op = this.expectOp()
      if (op !== '+' && op !== '-') break
      this.next()
      const right = this.parseTerm()
      if ('message' in right) return right
      left = { k: 'bin', op, a: left, b: right }
    }
    return left
  }

  private parseTerm(): Ast | SimExprError {
    let left = this.parseUnary()
    if ('message' in left) return left
    for (;;) {
      const op = this.expectOp()
      if (op !== '*' && op !== '/' && op !== '%') break
      this.next()
      const right = this.parseUnary()
      if ('message' in right) return right
      left = { k: 'bin', op, a: left, b: right }
    }
    return left
  }

  private parseUnary(): Ast | SimExprError {
    if (this.expectOp() === '-') {
      this.next()
      const a = this.parseUnary()
      if ('message' in a) return a
      return { k: 'un', op: '-', a }
    }
    if (this.expectOp() === '+') {
      this.next()
      return this.parseUnary()
    }
    return this.parsePower()
  }

  private parsePower(): Ast | SimExprError {
    const base = this.parseAtom()
    if ('message' in base) return base
    if (this.expectOp() === '^') {
      this.next()
      const exp = this.parseUnary() // right-assoc
      if ('message' in exp) return exp
      return { k: 'bin', op: '^', a: base, b: exp }
    }
    return base
  }

  private parseAtom(): Ast | SimExprError {
    const t = this.next()
    if (!t) return { message: 'unexpected end of expression' }
    if (t.t === 'num') return { k: 'num', v: t.v }
    if (t.t === 'lparen') {
      const inner = this.parseExpr()
      if ('message' in inner) return inner
      const close = this.next()
      if (close?.t !== 'rparen') return { message: 'missing closing parenthesis' }
      return inner
    }
    if (t.t === 'id') {
      if (this.peek()?.t === 'lparen') {
        this.next() // consume (
        const args: Ast[] = []
        if (this.peek()?.t !== 'rparen') {
          for (;;) {
            const arg = this.parseExpr()
            if ('message' in arg) return arg
            args.push(arg)
            if (this.peek()?.t === 'comma') {
              this.next()
              continue
            }
            break
          }
        }
        const close = this.next()
        if (close?.t !== 'rparen') return { message: 'missing closing parenthesis' }
        return { k: 'call', name: t.v, args }
      }
      return { k: 'id', v: t.v }
    }
    return { message: `unexpected token "${'v' in t ? t.v : t.t}"` }
  }

  parseTop(): Ast | SimExprError {
    const ast = this.parseExpr()
    if ('message' in ast) return ast
    if (this.pos < this.tokens.length) {
      return { message: 'trailing tokens after expression' }
    }
    return ast
  }
}

function evalAst(ast: Ast, env: Record<string, number>): number {
  switch (ast.k) {
    case 'num':
      return ast.v
    case 'id': {
      if (ast.v in CONSTANTS) return CONSTANTS[ast.v]
      const v = env[ast.v]
      return typeof v === 'number' ? v : NaN
    }
    case 'un':
      return -evalAst(ast.a, env)
    case 'bin': {
      const a = evalAst(ast.a, env)
      const b = evalAst(ast.b, env)
      switch (ast.op) {
        case '+':
          return a + b
        case '-':
          return a - b
        case '*':
          return a * b
        case '/':
          return b === 0 ? NaN : a / b
        case '%':
          return b === 0 ? NaN : a % b
        case '^':
          return Math.pow(a, b)
        default:
          return NaN
      }
    }
    case 'call': {
      const def = FUNCTIONS[ast.name]
      if (!def) return NaN
      const args = ast.args.map((a) => evalAst(a, env))
      if (args.some((x) => Number.isNaN(x))) return NaN
      return def.fn(...args)
    }
  }
}

function collectIds(ast: Ast, out: Set<string>): void {
  switch (ast.k) {
    case 'num':
      return
    case 'id':
      out.add(ast.v)
      return
    case 'un':
      collectIds(ast.a, out)
      return
    case 'bin':
      collectIds(ast.a, out)
      collectIds(ast.b, out)
      return
    case 'call':
      ast.args.forEach((a) => collectIds(a, out))
      return
  }
}

export type ParsedSimExpr = {
  /** Evaluate against an environment; NaN when uncomputable. */
  evaluate(env: Record<string, number>): number
  /** Identifiers used (params/computed refs), excluding constants. */
  identifiers: string[]
}

/**
 * Parse a safe math expression. Returns error message on syntax problems.
 * Unknown function names are rejected at parse time.
 */
export function parseSimExpr(src: string): ParsedSimExpr | SimExprError {
  const trimmed = src.trim()
  if (!trimmed || trimmed.length > 200) {
    return { message: 'expression empty or too long (max 200 chars)' }
  }
  const tokens = tokenize(trimmed)
  if ('message' in tokens) return tokens
  const ast = new Parser(tokens).parseTop()
  if ('message' in ast) return ast

  const ids = new Set<string>()
  collectIds(ast, ids)
  for (const id of ids) {
    if (id in FUNCTIONS) {
      return { message: `"${id}" is a function name — call it with (...)` }
    }
  }
  // Validate function calls (unknown names, arity)
  const checkCalls = (node: Ast): SimExprError | null => {
    if (node.k === 'call') {
      const def = FUNCTIONS[node.name]
      if (!def) return { message: `unknown function "${node.name}"` }
      if (node.args.length < def.minArgs || node.args.length > def.maxArgs) {
        return { message: `function "${node.name}" expects ${def.minArgs}–${def.maxArgs} args` }
      }
      for (const a of node.args) {
        const err = checkCalls(a)
        if (err) return err
      }
    } else if (node.k === 'un') {
      return checkCalls(node.a)
    } else if (node.k === 'bin') {
      return checkCalls(node.a) ?? checkCalls(node.b)
    }
    return null
  }
  const callErr = checkCalls(ast)
  if (callErr) return callErr

  const identifiers = [...ids].filter((id) => !(id in CONSTANTS))
  return {
    evaluate: (env) => evalAst(ast, env),
    identifiers,
  }
}

/**
 * Validation helper: parse + require every identifier ∈ allowedIds.
 * Returns a list of issue messages (empty = OK).
 */
export function validateSimExprRefs(
  src: string,
  allowedIds: Set<string>
): string[] {
  const parsed = parseSimExpr(src)
  if ('message' in parsed) return [parsed.message]
  return parsed.identifiers
    .filter((id) => !allowedIds.has(id))
    .map((id) => `unknown identifier "${id}"`)
}

/** Type guard helper for zod refinements. */
export function isValidSimParamId(id: string): boolean {
  return ID_RE.test(id) && !(id in FUNCTIONS) && !(id in CONSTANTS)
}