Math expression evaluator and expression solver.

mathy.lua

--------------------------------
-- StringBuffer (for printing out Trees)
--------------------------------
local stringBuffer
do
  local function insert(buffer, item)
    local len = buffer.length + 1
    buffer.length = len
    buffer[len] = item
  end
  
  local concat = table.concat
  
  stringBuffer = function()
    return { length = 0, insert = insert, concat = concat }
  end
end

--------------------------------
-- Tree Class
--------------------------------
local function Tree(value, ...)
  local tree = { ... }
  tree.value = value
  return tree
end

local function TreeToStringStep(tree, buffer)
  if tree.value == nil then return end
  
  local len = #tree
  
  if len == 0 then
    buffer:insert(tree.value)
  elseif len == 2 then
    buffer:insert("(")
    TreeToStringStep(tree[1], buffer)
    buffer:insert(" ")
    buffer:insert(tree.value)
    buffer:insert(" ")
    TreeToStringStep(tree[2], buffer)
    buffer:insert(")")
  else
    buffer:insert(tree.value)
    buffer:insert("(")
    for i = 1,len do
      TreeToStringStep(tree[i], buffer)
      buffer:insert(",")
    end
    buffer:insert(")")
  end
end

local function TreeToString(tree)
  local buffer = stringBuffer()
  TreeToStringStep(tree, buffer)
  return buffer:concat()
end

local function findNode(tree, value, path)
  if tree == nil then return nil end
  if not path then path = "" end
  
  if tree.value == value then
    return path
  end
  
  for i = 1,#tree do
    local loc = findNode(tree[i], value, path .. i)
    if loc then return loc end
  end
end

local function findNodes(tree, value, acc, path)
  if tree == nil then return nil end
  if not path then path = "" end
  
  if tree.value == value then
    acc[#acc+1] = path
  end
  
  for i = 1,#tree do
    findNodes(tree[i], value, acc, path .. i)
  end
end

--------------------------------
-- Parsing functions
--------------------------------

-- Count Parentheses
local function balancedParens(str)
  local stack, size = {}, 0
  for i = 1,#str do
    local c = str:sub(i,i)
    if c == "(" then
      size = size + 1
      stack[size] = true
    elseif c == ")" then
      if stack[size] then
        stack[size] = nil
        size = size - 1
      else
        return false
      end
    end
  end
  return size == 0
end

-- Operator List
local opList = { {"="}, {"-", "+"}, {"/", "*"}, {"^"} }

-- Main parsing function
local function parse(expr)

  do
    -- Remove outer parens if they exist
    local tmpExpr = expr:match "^%s*%(%s*(.*)%s*%)%s*$"
    if tmpExpr then return parse(tmpExpr) end
    -- Return a number if expression is just a number
    local n = expr:match "^%s*([%w%.]+)%s*$"
    if n then return Tree(n) end
  end

  local exprLen = #expr
  for i = 1,#opList do
    local ops = opList[i]
    local expL, op, expR
    local opPos = 0

    for j = 1,#ops do
      local searchFrom = 1

      while true do
        local opInd, opEnd = expr:find(ops[j], searchFrom, true)

        if opInd then
          local _expL = expr:sub(1, opInd - 1)

          if balancedParens(_expL) and opInd >= opPos then
            local _op, _expR =
            expr:sub(opInd, opEnd),
            expr:sub(opEnd + 1, exprLen)
            
            expL, op, expR = _expL, _op, _expR
            opPos = opInd
          end

          searchFrom = opInd + 1
        else break end
      end

    end

    if op then
      return Tree(op, parse(expL), parse(expR))
    end
  end

  error("Error parsing \"" .. expr .. "\".")
end

--------------------------------
-- Evaluate Expression Tree
--------------------------------
local function evalTree(tree)
  local value = tree.value
  
  if value:match("%d+%.?%d*") then
    return value
  elseif value == "+" then
    return evalTree(tree[1]) + evalTree(tree[2])
  elseif value == "-" then
    return evalTree(tree[1]) - evalTree(tree[2])
  elseif value == "*" then
    return evalTree(tree[1]) * evalTree(tree[2])
  elseif value == "/" then
    return evalTree(tree[1]) / evalTree(tree[2])
  elseif value == "^" then
    return evalTree(tree[1]) ^ evalTree(tree[2])
  elseif value == "log" then
    return math.log(evalTree(tree[1]))
  end
end

--------------------------------
-- Solve Expression Tree
--------------------------------
local function solveTreeStep(tree, value)
  local opNode = tree[1]
  local location = findNode(opNode, value)
  
  if location == "" then return end
  
  local sideA = location:sub(1,1) == "1"
  local rhs = tree[2]
  local opA, opB = opNode[1], opNode[2]
  local operator = opNode.value
  
  if sideA then
    tree[1] = opA
  else
    tree[1] = opB
  end
  
  if operator == "+" then
    if sideA then
      tree[2] = Tree("-", rhs, opB)
    else
      tree[2] = Tree("-", rhs, opA)
    end
  elseif operator == "-" then
    if sideA then
      tree[2] = Tree("+", rhs, opB)
    else
      tree[2] = Tree("-", opA, rhs)
    end
  elseif operator == "*" then
    if sideA then
      tree[2] = Tree("/", rhs, opB)
    else
      tree[2] = Tree("/", rhs, opA)
    end
  elseif operator == "/" then
    if sideA then
      tree[2] = Tree("*", rhs, opB)
    else
      tree[2] = Tree("/", opA, rhs)
    end
  elseif operator == "^" then
    if sideA then
      tree[2] = Tree("^", rhs, Tree("/", Tree("1"), opB))
    else
      tree[2] = Tree("/", Tree("log", rhs), Tree("log", opA))
    end
  else
    error("Unknown operator: " .. operator)
  end
  
  return solveTreeStep(tree, value)
end

local function solveTree(tree, value)
  local location = findNode(tree, value)

  if location:sub(1,1) == "2" then
    tree[1], tree[2] = tree[2], tree[1]
  end
  
  solveTreeStep(tree, value)
end

--------------------------------
-- Main
--------------------------------
function main(expression, variable)
  if expression then
    local banana = parse(expression)

    print(TreeToString(banana))

    if variable and banana.value == "=" then
      
      local varLoc = {}
      findNodes(banana, variable, varLoc)
      local numVars = #varLoc

      if numVars < 1 then
        print("Variable " .. variable .. " does not appear in expression.")
      elseif numVars > 1 then
        print("Variable " .. variable .. " appears multiple times in expression.")
      else
        solveTree(banana, variable, varLoc)
        print(TreeToString(banana))
        print(variable .. " = " .. evalTree(banana[2]))
      end
    else
      print("ans = " .. evalTree(banana))
    end
  else
    print(string.format("Usage: %s <expression>", arg[0]))
    print("\tExamples:")
    print(string.format("\t%s \"(1 + 5^0.5) / 2\"", arg[0]))
    print(string.format("\t%s \" 1 + 5^0.5 = 2*x\" x", arg[0]))
  end
end

main(arg[1], arg[2])