open Printf
open Types
open Util

let rec strip_comments = function
  | Comment _ :: tl -> strip_comments tl
  | InlineComment (EmptyLine, _) :: tl -> strip_comments tl
  | InlineComment (instr, _) :: tl -> strip_comments (instr :: tl)
  | hd :: tl -> hd :: (strip_comments tl)
  | [] -> []

let rec peephole = function
  (* Constant load before branch becomes a jump when the branch condition
   * matches the loaded value ... *)
  | LoadImm (BoolVal b) :: Branch (cond, tgt) :: tl when cond = b ->
    InlineComment (Jump tgt, "branch -> jump") :: (peephole tl)

  (* ... otherwise, both instructions can be removed *)
  | LoadImm (BoolVal _) :: Branch (_, tgt) :: tl ->
    InlineComment (EmptyLine, "load + branch removed") :: peephole tl

  (* Transform addition/subtraction by constant to increment/decrement:
   * iload L          |   iload L
   * iloadc[_ ]C      |   iloadc_1
   * i{add,sub}       |   i{add,sub}
   * istore L         |   istore L
   *     |                    |
   *     v                    v
   * i{inc,dec} L C   |   i{inc,dec}_1 L
   *)
  | Load (Int, Current, index) ::
    LoadConst (_, i) ::
    Op (Add, Int) ::
    Store (Int, Current, store) :: tl
  | LoadConst (_, i) ::
    Load (Int, Current, index) ::
    Op (Add, Int) ::
    Store (Int, Current, store) :: tl
    when store = index ->
      InlineComment (Inc (index, i), "add -> inc") :: (peephole tl)

  | Load (Int, Current, index) ::
    LoadConst (_, i) ::
    Op (Sub, Int) ::
    Store (Int, Current, store) :: tl
  | LoadConst (_, i) ::
    Load (Int, Current, index) ::
    Op (Sub, Int) ::
    Store (Int, Current, store) :: tl
    when store = index ->
      InlineComment (Dec (index, i), "sub -> dec") :: (peephole tl)

  | Load (Int, Current, index) ::
    LoadImm (IntVal 1l) ::
    Op (Add, Int) ::
    Store (Int, Current, store) :: tl
  | LoadImm (IntVal 1l) ::
    Load (Int, Current, index) ::
    Op (Add, Int) ::
    Store (Int, Current, store) :: tl
    when store = index ->
      InlineComment (IncOne index, "add -> inc") :: (peephole tl)

  | Load (Int, Current, index) :: LoadImm (IntVal 1l) :: Op (Sub, Int) ::
    Store (Int, Current, store) :: tl
  | LoadImm (IntVal 1l) ::
    Load (Int, Current, index) ::
    Op (Sub, Int) ::
    Store (Int, Current, store) :: tl
    when store = index ->
      InlineComment (DecOne index, "sub -> dec") :: (peephole tl)

  | hd :: tl -> hd :: (peephole tl)
  | [] -> []

(* Count actual instructions, ignoring comments and labels *)
let count_instrs instrs =
  let rec trav n = function
  | [] -> n
  | (Comment _ | Label _ | EmptyLine) :: tl -> trav n tl
  | InlineComment (hd, _) :: tl -> trav (trav n [hd]) tl
  | hd :: tl -> trav (n + 1) tl
  in trav 0 instrs

let phase = function
  | Assembly instrs ->
    let oldcount = count_instrs instrs in
    let instrs = peephole (strip_comments instrs) in
    let newcount = count_instrs instrs in
    log_line 2 (sprintf
      "Optimized %d to %d instructions (%d fewer)"
      oldcount newcount (oldcount - newcount)
    );
    Assembly instrs
  | _ -> raise InvalidInput