open! Ast
open! Base

let counter = ref 0

let count =
  fun () ->
    counter := (!counter) + 1;
    !counter


let arithmetic =
  Map.of_alist_exn (module String)
    [ "add", [ "@SP"; "AM=M-1"; "D=M"; "A=A-1"; "M=D+M" ]
    ; "sub", [ "@SP"; "AM=M-1"; "D=M"; "A=A-1"; "M=M-D" ]
    ; "and", [ "@SP"; "AM=M-1"; "D=M"; "A=A-1"; "M=D&M" ]
    ; "or" , [ "@SP"; "AM=M-1"; "D=M"; "A=A-1"; "M=D|M" ]
    ; "neg", [ "@SP"; "A=M-1"; "M=-M" ]
    ; "not", [ "@SP"; "A=M-1"; "M=!M" ]
    ]

let eq count =
  [ "@SP"
    ; "AM=M-1"
    ; "D=M"
    ; "A=A-1"
    ; "D=M-D"
    ; "M=-1"
    ; "@END_EQ_" ^ count
    ; "D;JEQ"
    ; "@SP"
    ; "A=M-1"
    ; "M=0"
    ; "(END_EQ_" ^ count ^ ")"
  ]

let gt count =
  [ "@SP"
    ; "AM=M-1"
    ; "D=M"
    ; "A=A-1"
    ; "D=M-D"
    ; "M=-1"
    ; "@END_GT_" ^ count
    ; "D;JGT"
    ; "@SP"
    ; "A=M-1"
    ; "M=0"
    ; "(END_GT_" ^ count ^ ")"
  ]

let lt count =
  [ "@SP"
    ; "AM=M-1"
    ; "D=M"
    ; "A=A-1"
    ; "D=M-D"
    ; "M=-1"
    ; "@END_LT_" ^ count
    ; "D;JLT"
    ; "@SP"
    ; "A=M-1"
    ; "M=0"
    ; "(END_LT_" ^ count ^ ")"
  ]

let translate command =
  let i = Int.to_string (count ()) in
  match command with
  | Add -> Map.find_exn arithmetic "add"
  | Sub -> Map.find_exn arithmetic "sub"
  | Neg -> Map.find_exn arithmetic "neg"
  | And -> Map.find_exn arithmetic "and"
  | Not -> Map.find_exn arithmetic "not"
  | Or  -> Map.find_exn arithmetic "or"
  | Eq  -> eq i
  | Gt  -> gt i
  | Lt  -> lt i