mathew/frontend/app/utils/plot-expression.ts

/* The graphing calculator's little algebra: a hand-rolled tokenizer and
   recursive-descent parser that reads an expression in x — "sin(x) + x / 4",
   "2^x", "x(x+1)" — and compiles it to a function the canvas evaluates pixel
   by pixel. No library: the grammar is small enough for one file, and the
   parser can be strict about what it promises — everything it accepts is a
   real function of x, and everything else is refused with a reason.

   Precedence is the usual one: + - bind loosest; then * / % (left, implicit
   multiplication included — "2x" is 2·x); then unary minus, so -x^2 negates
   the square; then ^, right-associative, its exponent may carry a sign so
   2^-x reads. Functions call, constants stand, parentheses group. */

export type PlotFunction = (x: number) => number

/* why a text was refused: nothing to read, a name the calculator does not
   speak, or a shape the grammar rejects */
export type PlotErrorReason = 'empty' | 'unknown' | 'syntax'

export type CompiledPlot =
  | { ok: true; fn: PlotFunction }
  | { ok: false; reason: PlotErrorReason }

type Op = '+' | '-' | '*' | '/' | '%' | '^' | '(' | ')' | ','

type Token = { t: Op } | { t: 'num'; value: number } | { t: 'id'; name: string }

type Node =
  | { kind: 'num'; value: number }
  | { kind: 'var' }
  | { kind: 'const'; value: number }
  | { kind: 'neg'; arg: Node }
  | { kind: 'bin'; op: '+' | '-' | '*' | '/' | '%' | '^'; a: Node; b: Node }
  | { kind: 'call'; name: string; args: Node[] }

/* the constants a name may stand for */
const CONSTANTS: Record<string, number> = {
  pi: Math.PI,
  tau: Math.PI * 2,
  e: Math.E,
  phi: (1 + Math.sqrt(5)) / 2,
}

/* the functions a name may call — all on the variable x, ln and log split
   the natural and the base-10 way as calculators write them */
const FUNCTIONS: Record<string, (...args: number[]) => number> = {
  sin: Math.sin,
  cos: Math.cos,
  tan: Math.tan,
  asin: Math.asin,
  acos: Math.acos,
  atan: Math.atan,
  sinh: Math.sinh,
  cosh: Math.cosh,
  tanh: Math.tanh,
  ln: Math.log,
  log: Math.log10,
  log2: Math.log2,
  log10: Math.log10,
  sqrt: Math.sqrt,
  cbrt: Math.cbrt,
  abs: Math.abs,
  exp: Math.exp,
  floor: Math.floor,
  ceil: Math.ceil,
  round: Math.round,
  sign: Math.sign,
  min: (...a) => Math.min(...a),
  max: (...a) => Math.max(...a),
  mod: (a, b) => a % b,
}

/* min and max take as many arguments as they are given; mod takes two;
   everything else wants exactly one */
function arityOk(name: string, args: number): boolean {
  if (name === 'min' || name === 'max') return args >= 1
  if (name === 'mod') return args === 2
  return args === 1
}

class PlotParseError extends Error {
  reason: Exclude<PlotErrorReason, 'empty'>
  constructor(reason: 'unknown' | 'syntax') {
    super(reason)
    this.reason = reason
  }
}

/* the input wears its "y =" beside the field; a reader who types it too is
   answered the same — the rest is what gets plotted. f(x) = reads too. */
function stripAssignment(src: string): string {
  return src.replace(/^\s*(?:y|f\s*\(\s*x\s*\))\s*=/i, '')
}

function tokenize(src: string): Token[] {
  /* the typographic shapes a reader may paste in stand for the plain ones */
  const s = src
    .replace(/[×·∗]/g, '*')
    .replace(/÷/g, '/')
    .replace(/[−–]/g, '-')
    .replace(/π/g, 'pi')
    .replace(/τ/g, 'tau')

  const toks: Token[] = []
  let i = 0
  while (i < s.length) {
    const c = s[i]!
    if (c === ' ' || c === '\t' || c === '\n') {
      i++
      continue
    }
    if ((c >= '0' && c <= '9') || c === '.') {
      const start = i
      let dot = false
      while (i < s.length) {
        const d = s[i]!
        if (d >= '0' && d <= '9') i++
        else if (d === '.' && !dot) {
          dot = true
          i++
        } else break
      }
      const value = Number.parseFloat(s.slice(start, i))
      if (!Number.isFinite(value)) throw new PlotParseError('syntax')
      toks.push({ t: 'num', value })
      continue
    }
    if (/[a-z_]/i.test(c)) {
      const start = i
      while (i < s.length && /[a-z0-9_]/i.test(s[i]!)) i++
      toks.push({ t: 'id', name: s.slice(start, i).toLowerCase() })
      continue
    }
    if (c === '*' && s[i + 1] === '*') {
      i += 2
      toks.push({ t: '^' }) // ** as an exponent, the programming shape
      continue
    }
    if (c === '+' || c === '-' || c === '*' || c === '/' || c === '%' || c === '^') {
      toks.push({ t: c as Op })
      i++
      continue
    }
    if (c === '(' || c === ')' || c === ',') {
      toks.push({ t: c as Op })
      i++
      continue
    }
    throw new PlotParseError('syntax')
  }
  return toks
}

interface Parser {
  toks: Token[]
  pos: number
}

const peek = (p: Parser): Token | null => p.toks[p.pos] ?? null

const startsAtom = (tk: Token): boolean => tk.t === 'num' || tk.t === 'id' || tk.t === '('

function parseExpr(p: Parser): Node {
  let node = parseTerm(p)
  for (;;) {
    const tk = peek(p)
    if (!tk || (tk.t !== '+' && tk.t !== '-')) break
    p.pos++
    node = { kind: 'bin', op: tk.t, a: node, b: parseTerm(p) }
  }
  return node
}

function parseTerm(p: Parser): Node {
  let node = parseUnary(p)
  for (;;) {
    const tk = peek(p)
    if (!tk) break
    if (tk.t === '*' || tk.t === '/' || tk.t === '%') {
      p.pos++
      node = { kind: 'bin', op: tk.t, a: node, b: parseUnary(p) }
      continue
    }
    /* nothing between two factors and a factor following: that's a
       multiplication — 2x, 3(x+1), x sin(x) */
    if (startsAtom(tk)) {
      node = { kind: 'bin', op: '*', a: node, b: parseUnary(p) }
      continue
    }
    break
  }
  return node
}

function parseUnary(p: Parser): Node {
  const tk = peek(p)
  if (tk && tk.t === '+') {
    p.pos++
    return parseUnary(p)
  }
  if (tk && tk.t === '-') {
    p.pos++
    return { kind: 'neg', arg: parseUnary(p) }
  }
  return parsePower(p)
}

function parsePower(p: Parser): Node {
  const base = parseAtom(p)
  const tk = peek(p)
  if (tk && tk.t === '^') {
    p.pos++
    // right-associative, and the exponent may carry its own sign
    return { kind: 'bin', op: '^', a: base, b: parseUnary(p) }
  }
  return base
}

function parseAtom(p: Parser): Node {
  const tk = peek(p)
  if (!tk) throw new PlotParseError('syntax')
  if (tk.t === 'num') {
    p.pos++
    return { kind: 'num', value: tk.value }
  }
  if (tk.t === '(') {
    p.pos++
    const inner = parseExpr(p)
    const close = peek(p)
    if (!close || close.t !== ')') throw new PlotParseError('syntax')
    p.pos++
    return inner
  }
  if (tk.t === 'id') {
    p.pos++
    if (tk.name === 'x') return { kind: 'var' }
    if (Object.hasOwn(CONSTANTS, tk.name)) return { kind: 'const', value: CONSTANTS[tk.name]! }
    if (Object.hasOwn(FUNCTIONS, tk.name)) {
      const open = peek(p)
      if (!open || open.t !== '(') throw new PlotParseError('syntax')
      p.pos++
      const args: Node[] = [parseExpr(p)]
      while (peek(p)?.t === ',') {
        p.pos++
        args.push(parseExpr(p))
      }
      const close = peek(p)
      if (!close || close.t !== ')') throw new PlotParseError('syntax')
      p.pos++
      if (!arityOk(tk.name, args.length)) throw new PlotParseError('syntax')
      return { kind: 'call', name: tk.name, args }
    }
    throw new PlotParseError('unknown')
  }
  throw new PlotParseError('syntax')
}

function evaluate(node: Node, x: number): number {
  switch (node.kind) {
    case 'num':
      return node.value
    case 'var':
      return x
    case 'const':
      return node.value
    case 'neg':
      return -evaluate(node.arg, x)
    case 'call': {
      const fn = FUNCTIONS[node.name]!
      return fn(...node.args.map((a) => evaluate(a, x)))
    }
    case 'bin': {
      const a = evaluate(node.a, x)
      const b = evaluate(node.b, x)
      switch (node.op) {
        case '+':
          return a + b
        case '-':
          return a - b
        case '*':
          return a * b
        case '/':
          return a / b
        case '%':
          return a % b
        case '^':
          return Math.pow(a, b)
      }
    }
  }
}

/* read one expression; a refusal says whether a name was not known or the
   shape was not a function — an empty text is the caller's to treat as
   "nothing asked for", not as an error */
export function compilePlotExpression(src: string): CompiledPlot {
  if (src.length > 400) return { ok: false, reason: 'syntax' }
  try {
    const toks = tokenize(stripAssignment(src))
    if (toks.length === 0) return { ok: false, reason: 'empty' }
    const p: Parser = { toks, pos: 0 }
    const ast = parseExpr(p)
    if (p.pos !== toks.length) throw new PlotParseError('syntax')
    return { ok: true, fn: (x) => evaluate(ast, x) }
  } catch (err) {
    if (err instanceof PlotParseError) return { ok: false, reason: err.reason }
    return { ok: false, reason: 'syntax' }
  }
}