25 Arithmetic expression (infix notation) parsing & evaluation & transformation into prefix notation

arithmetic.js

// arithmetic expression (in infix notation) parsing & evaluation & transformation into prefix notation
// and implementation of algebraic expression AST transformation rules (derivative & simplification)
// * for numerical values, only integer is supported

var TOKEN_TYPE = {
  'op': Symbol('op'),
  'num': Symbol('num'),
  'var': Symbol('var'),
  'paren': Symbol('paren')
}

var ASSOCIATIVITY = {
  'right': Symbol('right'),
  'left': Symbol('left')
}

var OP_TABLE = {
  '+': { prec: 2, assoc: ASSOCIATIVITY['left'], f: (a, b) => a + b },
  '-': { prec: 2, assoc: ASSOCIATIVITY['left'], f: (a, b) => a - b },
  '*': { prec: 3, assoc: ASSOCIATIVITY['left'], f: (a, b) => a * b },
  '/': { prec: 3, assoc: ASSOCIATIVITY['left'], f: (a, b) => a / b },
  '^': { prec: 4, assoc: ASSOCIATIVITY['right'], f: (a, b) => Math.pow(a, b) },
  '(': { prec: 0 }, // hack: parentheses are actually not operator
  ')': { prec: 0 },
}

class ASTNode {
  constructor(value, left, right) {
    this.value = value
    this.left = left
    this.right = right
  }
}

class Expression {
  constructor(literal) {
    this.literal = literal
    this.init()
  }
  init() {
    var tokens = this.tokenize(this.literal)
    this.ast = this.parse(tokens)
  }
  tokenize(literal) {
    var char = literal[0]
    var i = 0
    var next = (c) => {
      if (c) {
        if (c == char) {
          i++
          char = literal[i]
        } else {
          throw new Error(`Expect ${c} at ${i}, got ${char} instead`)
        }
      } else {
        i++
        char = literal[i]
      }
    }

    // single charactor operators
    var caret = () => {
      next('^')
      return {
        value: '^',
        type: TOKEN_TYPE['op']
      }
    }
    var add = () => {
      next('+')
      return {
        value: '+',
        type: TOKEN_TYPE['op']
      }
    }
    var subt = () => {
      next('-')
      return {
        value: '-',
        type: TOKEN_TYPE['op']
      }
    }
    var mult = () => {
      next('*')
      return {
        value: '*',
        type: TOKEN_TYPE['op']
      }
    }
    var div = () => {
      next('/')
      return {
        value: '/',
        type: TOKEN_TYPE['op']
      }
    }
    var leftParen = () => {
      next('(')
      return {
        value: '(',
        type: TOKEN_TYPE['paren']
      }
    }
    var rightParen = () => {
      next(')')
      return {
        value: ')',
        type: TOKEN_TYPE['paren']
      }
    }
    var whitespace = () => {
      while (char == ' ') {
        next()
      }
    }
    var variable = () => {
      if (/[a-z]/.test(char)) {
        var v = char
        next()
        return {
          value: v,
          type: TOKEN_TYPE['var']
        }
      }
    }
    // number
    var number = () => {
      var value = ''
      while (/[0-9]/.test(char)) {
        value += char
        next()
      }
      value = parseInt(value)
      return {
        value: value,
        type: TOKEN_TYPE['num']
      }
    }

    var tokens = []
    var tick = () => {
      if (char == '^') {
        return caret()
      } else if (char == '+') {
        return add()
      } else if (char == '-') {
        return subt()
      } else if (char == '*') {
        return mult()
      } else if (char == '/') {
        return div()
      } else if (char == '(') {
        return leftParen()
      } else if (char == ')') {
        return rightParen()
      } else if (/[a-z]/.test(char)) {
        return variable()
      } else if (/[0-9]/.test(char)) {
        return number()
      } else {
        throw new Error(`Token ${char} unrecognizable at ${i}`)
      }
    }
    whitespace()
    while (char) {
      tokens.push(tick())
      whitespace()
    }
    return tokens
  }
  // util function for finding precedence of operator token
  prec(token) {
    return OP_TABLE[token.value].prec
  }
  // util function for finding associativity of operator token
  assoc(token) {
    return OP_TABLE[token.value].assoc
  }
  // parse algebraic expression tokens into ast using shunting yard algorithm
  parse(tokens) {
    var output = []
    var opstack = []
    var i = 0
    var t = tokens[i]
    // reading tokens and build up AST node stack and operation stack
    while (t) {
      // if there is a token input of type variable or number
      if (t.type == TOKEN_TYPE['var'] || t.type == TOKEN_TYPE['num']) {
        output.push(t)
      // if there is a left parenthesis token input
      } else if (t.value == '(') {
        opstack.push(t)
      // if there is a right parenthesis token input
      } else if (t.value == ')') {
        var op
        // until matching parenthesis is found
        while (true) {
          op = opstack.pop()
          if (op.value == '(') {
            break
          }
          var b = output.pop()
          var a = output.pop()
          output.push(new ASTNode(op.value, a, b))
        }
      // if there is a token input of type operator
      } else {
        var op
        while(true) {
          op = opstack.pop()
          if (op) {
            if ((this.prec(op) > this.prec(t)) ||
              (this.prec(op) == this.prec(t) && this.assoc(t) == ASSOCIATIVITY['left'] && op.value != '(')
            ) {
              var b = output.pop()
              var a = output.pop()
              output.push(new ASTNode(op.value, a, b))
            } else {
              opstack.push(op)
              break
            }
          } else {
            break
          }
        }
        opstack.push(t)
      }
      i++
      t = tokens[i]
    }
    // resolve remaining operation stacks
    while (opstack.length > 0) {
      var b = output.pop()
      var a = output.pop()
      var op = opstack.pop()
      output.push(new ASTNode(op.value, a, b))
    }
    return output[0]
  }
  // get result of the expression, will return NaN if the expression contains variable
  eval() {
    var _eval = (ast) => {
      if (ast.constructor.name == 'ASTNode') {
        var op = OP_TABLE[ast.value].f
        var a = ast.left
        var b = ast.right
        return op(_eval(a), _eval(b))
      } else {
        return ast.value
      }
    }
    return _eval(this.ast)
  }
}

// print ast in parenthesized prefix notation with identation (Lisp style) into console
function prefix(ast) {
  var _prefix = (ast, level) => {
    if (ast.constructor.name == 'ASTNode') {
      var op = ast.value
      var a = ast.left
      var b = ast.right
      var ident = new Array(level * 3).fill(' ').join('')
      if (b.constructor.name != 'ASTNode') {
        return `(${op} ${_prefix(a, level + 1)} ${b.value})`
      } else {
        return `(${op} ${_prefix(a, level + 1)}\n${ident}${_prefix(b, level + 1)})`
      }
    } else {
      return ast.value
    }
  }
  console.log(_prefix(ast, 1))
}

// print ast in parenthesized infix notation with identation (Lisp style) into console
function infix(ast) {
  var _infix = (ast, level) => {
    if (ast.constructor.name == 'ASTNode') {
      var op = ast.value
      var a = ast.left
      var b = ast.right
      if (b.constructor.name != 'ASTNode') {
        return `${_infix(a, level + 1)}${op}${b.value}`
      } else {
        return `${_infix(a, level + 1)}${op}${_infix(b, level + 1)}`
      }
    } else {
      return ast.value
    }
  }
  console.log(_infix(ast, 1))
}

function deriv(ast, respect) {
  if (ast.constructor.name == 'ASTNode') {
    // sum rule
    if (ast.value == '+') {
      return new ASTNode('+', deriv(ast.left, respect), deriv(ast.right, respect))
    // product rule
    } else if (ast.value == '*') {
      var t1 = new ASTNode('*', deriv(ast.left, respect), ast.right)
      var t2 = new ASTNode('*', ast.left, deriv(ast.right, respect))
      return new ASTNode('+', t1, t2)
    // power rule
    } else if (ast.value == '^' && ast.right.type == TOKEN_TYPE['num']) {
      var t1 = { value: ast.right.value - 1, type: TOKEN_TYPE['num']}
      var t2 = { value: ast.right.value, type: TOKEN_TYPE['num']} 
      var t3 = { value: respect, type: TOKEN_TYPE['var']}
      var t4 = new ASTNode('^', t3, t1)
      var t5 = new ASTNode('*', t2, t4)
      return new ASTNode('*', t5, deriv(ast.left, respect))
    }
  // leaf node
  } else {
    if (ast.type == TOKEN_TYPE['num']) {
      return { value: 0, type: TOKEN_TYPE['num']}
    } else if (ast.type == TOKEN_TYPE['var']) {
      if (ast.value == respect) {
        return { value: 1, type: TOKEN_TYPE['num']}
      } else {
        return { value: 0, type: TOKEN_TYPE['num']}
      }
    }
  }
}

// simplify polynomial factors by some simple rules
function simplify(ast) {
  if (ast.constructor.name == 'ASTNode') {
    if (ast.value == '+') {
      if (ast.left.value === 0) {
        return simplify(ast.right)
      } else if (ast.right.value === 0) {
        return simplify(ast.left)
      } else if (ast.left.type == TOKEN_TYPE['num'] && ast.right.type == TOKEN_TYPE['num']) {
        return simplify({value: ast.left.value + ast.right.value, type: TOKEN_TYPE['num']})
      } else {
        return new ASTNode('+', simplify(ast.left), simplify(ast.right))
      }
    } else if (ast.value == '*') {
      if (ast.left.value === 1) {
        return simplify(ast.right)
      } else if (ast.right.value === 1) {
        return simplify(ast.left)
      } else if (ast.left.value === 0 || ast.right.value === 0) {
        return simplify({value: 0, type: TOKEN_TYPE['num']})
      // combination rule: ca * (cb * f(x)) = (ca * cb) * f(x)
      } else if (ast.left.type == TOKEN_TYPE['num'] && ast.right.left && ast.right.value == '*' && ast.right.left.type == TOKEN_TYPE['num']) {
        var t1 = {value: ast.left.value * ast.right.left.value, type: TOKEN_TYPE['num']}
        var t2 = ast.right.right
        var t3 = new ASTNode('*', t1, t2)
        return simplify(t3)
      } else {
        return new ASTNode('*', simplify(ast.left), simplify(ast.right))
      }
    } else if (ast.value == '^') {
      if (ast.right.value === 0) {
        return simplify({value: 1, type: TOKEN_TYPE['num']})
      } else if (ast.right.value === 1) {
        return simplify(ast.left)
      } else {
        return new ASTNode('^', simplify(ast.left), simplify(ast.right))
      }
    }
  } else {
    return ast
  }
}

// test
var a = new Expression('x^3+5*x^2+6*x+2')
infix(a.ast)
prefix(a.ast)
console.log('=========')

// find the derivative of above function
var dx = deriv(a.ast, 'x')
infix(dx)
prefix(dx)
console.log('=========')

// simplify the derivative function expression
var s = simplify
var simplified = s(s(s(dx)))
infix(simplified)
prefix(simplified)
console.log('=========')

// another example with pure arithmetic operations
var z = new Expression("3 + 4 * 2 / ( 1 - 5 ) ^ 2 ^ 3")
infix(z.ast)
prefix(z.ast)
console.log(z.eval())
console.log('=========')