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. A whole equation answers too: "y^2 + x^2 = 1" reads as a
   relation F(x, y) = left − right, and the curve is where F is zero. So a
   text with an "=" is one relation of x and y; without it, an expression
   mentioning y is the relation that expression equals zero, and one in x
   alone is the function y = f(x). 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 or a real relation in the
   plane, 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

/* a relation of the plane: the curve is its zero set */
export type PlaneFunction = (x: number, y: 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; kind: 'function'; fn: PlotFunction }
  | { ok: true; kind: 'relation'; fn: PlaneFunction }
  // "(x, y)" or "P = (x, y)": a single dot at fixed coordinates, optionally named
  | { ok: true; kind: 'point'; x: number; y: number; label: string | null }
  | { 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'; name: 'x' | 'y' }
  | { 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 may wear its "f(x) =" beside the field; a reader who types it
   too is answered the same — the rest is what gets plotted. A leading "y ="
   is read by the equation rule instead, so a right side that also carries y
   still means the same relation, not a mystery name. */
function stripAssignment(src: string): string {
  return src.replace(/^\s*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 === '=') {
      toks.push({ t: '=' })
      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', name: 'x' }
    if (tk.name === 'y') return { kind: 'var', name: 'y' }
    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, y: number): number {
  switch (node.kind) {
    case 'num':
      return node.value
    case 'var':
      return node.name === 'x' ? x : y
    case 'const':
      return node.value
    case 'neg':
      return -evaluate(node.arg, x, y)
    case 'call': {
      const fn = FUNCTIONS[node.name]!
      return fn(...node.args.map((a) => evaluate(a, x, y)))
    }
    case 'bin': {
      const a = evaluate(node.a, x, y)
      const b = evaluate(node.b, x, y)
      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)
      }
    }
  }
}

/* whether y stands anywhere in a tree — the difference between a function of
   x, which the canvas can trace column by column, and a relation, which it
   must search cell by cell */
function mentionsY(node: Node): boolean {
  switch (node.kind) {
    case 'var':
      return node.name === 'y'
    case 'neg':
      return mentionsY(node.arg)
    case 'bin':
      return mentionsY(node.a) || mentionsY(node.b)
    case 'call':
      return node.args.some(mentionsY)
    default:
      return false
  }
}

/* a point's coordinates hold no variable — a lone x or y has no single value
   to place a dot at, so only constant expressions (numbers, pi, sqrt(2)…)
   make a coordinate */
function isConstant(node: Node): boolean {
  switch (node.kind) {
    case 'var':
      return false
    case 'neg':
      return isConstant(node.arg)
    case 'bin':
      return isConstant(node.a) && isConstant(node.b)
    case 'call':
      return node.args.every(isConstant)
    default:
      return true
  }
}

/* the inside of a "(…, …)" split into its two coordinate token runs, or null
   unless there is exactly one comma at this depth */
function splitPair(inner: Token[]): [Token[], Token[]] | null {
  let depth = 0
  let comma = -1
  for (let i = 0; i < inner.length; i++) {
    const t = inner[i]!
    if (t.t === '(') depth++
    else if (t.t === ')') depth--
    else if (t.t === ',' && depth === 0) {
      if (comma !== -1) return null // more than one separator: not a plane point
      comma = i
    }
  }
  if (comma === -1) return null
  return [inner.slice(0, comma), inner.slice(comma + 1)]
}

/* tokens shaped "(x, y)": one pair of parens wrapping the whole run, split by
   a single top-level comma — the point "(2, 3)". A "(x+1)" has no comma and a
   "(1)+(2)" closes its parens early, so neither is read as a point. */
function pointPair(toks: Token[]): [Token[], Token[]] | null {
  if (toks.length < 4 || toks[0]!.t !== '(' || toks[toks.length - 1]!.t !== ')') return null
  let depth = 0
  for (let i = 0; i < toks.length; i++) {
    if (toks[i]!.t === '(') depth++
    else if (toks[i]!.t === ')') depth--
    if (depth === 0 && i < toks.length - 1) return null // parens closed before the end
  }
  return splitPair(toks.slice(1, toks.length - 1))
}

/* one run of tokens, the whole of it — a leftover is a shape the grammar
   does not know */
function parseAll(toks: Token[]): Node {
  const p: Parser = { toks, pos: 0 }
  const ast = parseExpr(p)
  if (p.pos !== toks.length) throw new PlotParseError('syntax')
  return ast
}

/* read one expression, equation, or point; a refusal says whether a name was
   not known or the shape was wrong — an empty text is the caller's to treat as
   "nothing asked for", not as an error. A "(x, y)" — or "P = (x, y)" with a
   name out front — is a point at fixed coordinates. With an "=" otherwise the
   text is the relation left − right = 0 (a lone "y =" ahead of an expression
   in x is still the fast function); without it, an expression in x is y = f(x)
   and one carrying y is the relation that expression equals zero. */
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 at = toks.findIndex((tk) => tk.t === '=')

    // A point: "(x, y)" on its own, or "P = (x, y)" with a lone name ahead of
    // the "=". Coordinates must be constant — a variable has no single value
    // to place a dot at — and finite once read.
    const named = at === 1 && toks[0]!.t === 'id'
    const pair = pointPair(named ? toks.slice(2) : toks)
    if (pair) {
      const cx = parseAll(pair[0])
      const cy = parseAll(pair[1])
      if (!isConstant(cx) || !isConstant(cy)) return { ok: false, reason: 'syntax' }
      const x = evaluate(cx, 0, 0)
      const y = evaluate(cy, 0, 0)
      if (!Number.isFinite(x) || !Number.isFinite(y)) return { ok: false, reason: 'syntax' }
      const label = named ? src.slice(0, src.indexOf('=')).trim() : null
      return { ok: true, kind: 'point', x, y, label: label || null }
    }

    if (at === -1) {
      const ast = parseAll(toks)
      return mentionsY(ast)
        ? { ok: true, kind: 'relation', fn: (x, y) => evaluate(ast, x, y) }
        : { ok: true, kind: 'function', fn: (x) => evaluate(ast, x, 0) }
    }
    const left = toks.slice(0, at)
    const right = toks.slice(at + 1)
    if (left.length === 0 || right.length === 0 || right.some((tk) => tk.t === '='))
      throw new PlotParseError('syntax')
    const rhs = parseAll(right)
    // "y = <only x>" is the explicit curve it always was: x alone decides it
    if (left.length === 1 && left[0]!.t === 'id' && left[0]!.name === 'y' && !mentionsY(rhs))
      return { ok: true, kind: 'function', fn: (x) => evaluate(rhs, x, 0) }
    const lhs = parseAll(left)
    return { ok: true, kind: 'relation', fn: (x, y) => evaluate(lhs, x, y) - evaluate(rhs, x, y) }
  } catch (err) {
    if (err instanceof PlotParseError) return { ok: false, reason: err.reason }
    return { ok: false, reason: 'syntax' }
  }
}