Arithmetic Expression Evaluator in OCaml

calculate.ml

open Printf

type token =
  | Digit of int
  | Operator of char
  | LeftParen
  | RightParen
;;

let token_to_string = function
  | Digit d -> string_of_int d
  | Operator c -> String.make 1 c
  | LeftParen -> "("
  | RightParen -> ")"
;;

let rec tokens_to_string = function
  | [] -> ""
  | token :: rest -> (token_to_string token) ^ " " ^ (tokens_to_string rest)
;;

let print_tokens tokens = printf "%s\n" (tokens_to_string tokens);;

let precedence = function
  | '+' | '-' -> 1
  | '*' | '/' -> 2
  | '^' -> 3
  | _ -> raise (Invalid_argument "Bad token")
;;

let rec tilt output operators pivot =
  let p = precedence pivot in
  match operators with
  | [] -> (output, [Operator pivot])
  | (Operator c as o) :: rest ->
    if precedence c >= p then
      tilt (o :: output) rest pivot
    else
      (output, (Operator pivot) :: operators)
  | LeftParen :: rest -> (output, (Operator pivot) :: operators)
  | _ -> raise (Invalid_argument "Operators contained non-operator token")
;;

let rec tilt_and_pop output operators target =
  match operators with
  | [] -> raise (Invalid_argument "Unmatched token")
  | token :: rest ->
    if token = target then
      (output, rest) (* discard the token *)
    else
      tilt_and_pop (token :: output) rest target

let rec tilt_all (output, operators) = List.rev output @ operators;;

let shunting_yard tokens =
    let rec run (output, operators) = function
    | [] -> (output, operators)
    | token :: rest -> begin match token with
      | Digit _ as d -> run ((d :: output), operators) rest
      | Operator c as o -> begin
          match c with
          | '^' -> run (output, (o :: operators)) rest
          | '+' | '-' | '*' | '/' | '%' -> run (tilt output operators c) rest
          | _ -> raise (Invalid_argument "Bad token")
        end
      | LeftParen as l -> run (output, l :: operators) rest
      | RightParen -> run (tilt_and_pop output operators LeftParen) rest
      end
    in tilt_all (run ([], []) tokens)
;;

let apply x y = function
    | '+' -> x + y
    | '-' -> x - y
    | '*' -> x * y
    | '/' -> x / y
    | '^' -> int_of_float ((float_of_int x) ** (float_of_int y))
    | _ -> raise (Invalid_argument "Bad operator")
;;

let lex_number string =
  let rec loop number s = match s with
    | [] -> (number, [])
    | c :: rest -> match c with
      | '0' | '1' | '2' | '3' | '4' | '5' | '6' | '7' | '8' | '9' ->
        loop (number ^ (String.make 1 c)) rest
      | _ -> (number, s)
  in let (number, rest) = loop "" string
  in (Digit (int_of_string number), rest);
;;


let lex string =
  let rec list_of_string = function
    | "" -> []
    | s -> s.[0] :: (list_of_string (String.sub s 1 (String.length s - 1)))
  in
  let rec loop = function
    | [] -> []
    | c :: rest -> match c with
      | '0' | '1' | '2' | '3' | '4' | '5' | '6' | '7' | '8' | '9'
        -> let (number, tail) = lex_number (c :: rest) in number :: loop tail
      | '+' | '-' | '/' | '*' | '^' -> (Operator c) :: (loop rest)
      | '(' -> LeftParen :: (loop rest)
      | ')' -> RightParen :: (loop rest)
      | ' ' | '\t' | '\n' | '\r' -> loop rest
      | _ -> raise (Invalid_argument "bad token")
  in loop (list_of_string string)
;;

let calculate arithmetic_expression =
  let infix = lex arithmetic_expression in
  let postfix = shunting_yard infix in
  let rec loop stack = function
    | [] -> List.hd stack
    | token :: rest -> match token with
        | Digit d -> loop (d :: stack) rest
        | Operator o -> begin match stack with
            | y :: x :: bottom -> loop ((apply x y o) :: bottom) rest
            | _ -> raise (Invalid_argument "missing operand")
          end
        | _ -> raise (Invalid_argument "bad token")
  in loop [] postfix
;;

let equation = "1+5^(2*2)-400*2-3-3-3-3+5^2^3-50/2";;
printf "%d\n" (calculate equation);;