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 1) :: Op (Add, Int) ::
            Store (Int, Current, store) :: tl
    |  LoadImm (IntVal 1) :: 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 1) :: Op (Sub, Int) ::
            Store (Int, Current, store) :: tl
    |  LoadImm (IntVal 1) :: 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 input =
    match input with
    | Assembly instrs ->
        if args.optimize then (
            log_line 1 "- Peephole optimization";
            let oldcount = count_instrs instrs in
            let instrs = peephole (strip_comments instrs) in
            let newcount = count_instrs instrs in
            log_line 1 (sprintf
                "  Optimized %d to %d instructions (%d difference)"
                oldcount newcount (newcount - oldcount)
            );
            Assembly instrs
        ) else
            input
    | _ -> raise (InvalidInput "peephole")