module Ast = Spice_syntax.Ast

type imm = Value.t
type vtable = Value.vtable

type reg = R of int [@@unboxed]

type ins =
  (* registers *)
  | LDI of reg * imm
  | LDR of reg * reg
  (* arithmetic *)
  | ADD of reg * reg * reg
  | SUB of reg * reg * reg
  | MUL of reg * reg * reg
  (* comparison *)
  | LST of reg * reg * reg
  | GRT of reg * reg * reg
  | EQL of reg * reg * reg
  | NOT of reg * reg
  (* objects *)
  | GET of reg * reg * reg
  | SET of reg * reg * reg
  | LOC of reg * reg * string
  | CON of reg * vtable
  | CAL of reg * reg * reg * reg list
  (* control flow *)
  | RET of reg
  | JMP of block
  | CBR of reg * block * block

and block =
  { mutable ins_list_rev : ins list }

let make_block () =
  { ins_list_rev = [] }

let extend t ins =
  t.ins_list_rev <- ins :: t.ins_list_rev

let instructions t =
  List.rev t.ins_list_rev


type prog =
  { entry : block }

type Value.mthd +=
  | Method of { n_args : int; body : prog }

let frame_size t =
  let queue = ref [ t.entry ] in
  let visited = ref !queue in
  let enqueue b =
    if not (List.memq b !visited) then (
      queue := b :: !queue;
      visited := b :: !visited)
  in
  let meas (R i) fs = max fs (i + 1) in
  let meas_ins fs = function
    | RET r
    | LDI (r, _)
    | CON (r, _)
      -> meas r fs
    | LDR (r1, r2)
    | NOT (r1, r2)
    | LOC (r1, r2, _)
      -> meas r1 (meas r2 fs)
    | ADD (r1, r2, r3)
    | SUB (r1, r2, r3)
    | MUL (r1, r2, r3)
    | LST (r1, r2, r3)
    | GRT (r1, r2, r3)
    | EQL (r1, r2, r3)
    | GET (r1, r2, r3)
    | SET (r1, r2, r3)
      -> meas r1 (meas r2 (meas r3 fs))
    | CAL (r1, r2, r3, rs) ->
      List.fold_right meas (r1::r2::r3::rs) fs
    | JMP b ->
      enqueue b;
      fs
    | CBR (r, b1, b2) ->
      enqueue b1;
      enqueue b2;
      meas r fs
  in
  let rec loop fs =
    match !queue with
    | [] -> fs
    | bl :: rest ->
      queue := rest;
      loop (List.fold_left meas_ins fs bl.ins_list_rev)
  in
  loop 1

(* pretty printing *)

let pp_reg ppf (R i) = Fmt.pf ppf "R%d" i

let pp_vtable ppf vt =
  Fmt.pf ppf "(%d){" vt.Value.n_slots;
  let sep = ref "" in
  Hashtbl.iter
    (fun name -> function
       | Value.Method _ -> ()
       | Value.Field idx ->
         Fmt.pf ppf "%s%s@%d" !sep name idx;
         sep := ";")
    vt.elems;
  Fmt.pf ppf "}"

let pp_ins ~label ppf = function
  | LDI (a, b) -> Fmt.pf ppf "mov %a, %s" pp_reg a (Value.to_string b)
  | LDR (a, b) -> Fmt.pf ppf "mov %a, %a" pp_reg a pp_reg b 
  | ADD (a, b, c) -> Fmt.pf ppf "add %a, %a, %a" pp_reg a pp_reg b pp_reg c
  | SUB (a, b, c) -> Fmt.pf ppf "sub %a, %a, %a" pp_reg a pp_reg b pp_reg c
  | MUL (a, b, c) -> Fmt.pf ppf "mul %a, %a, %a" pp_reg a pp_reg b pp_reg c
  | LST (a, b, c) -> Fmt.pf ppf "lst %a, %a, %a" pp_reg a pp_reg b pp_reg c
  | GRT (a, b, c) -> Fmt.pf ppf "grt %a, %a, %a" pp_reg a pp_reg b pp_reg c
  | EQL (a, b, c) -> Fmt.pf ppf "eql %a, %a, %a" pp_reg a pp_reg b pp_reg c
  | NOT (a, b) -> Fmt.pf ppf "not %a, %a" pp_reg a pp_reg b
  | GET (a, b, c) -> Fmt.pf ppf "mov %a, %a[%a]" pp_reg a pp_reg b pp_reg c
  | SET (a, b, c) -> Fmt.pf ppf "mov %a[%a], %a" pp_reg b pp_reg c pp_reg a
  | LOC (a, b, el) -> Fmt.pf ppf "mov %a, &%a.%s" pp_reg a pp_reg b el
  | CON (a, vt) -> Fmt.pf ppf "con %a, %a" pp_reg a pp_vtable vt
  | CAL (a, b, c, ds) ->
    Fmt.pf ppf "cal %a, %a[%a](" pp_reg a pp_reg b pp_reg c;
    List.iteri (fun i d -> if i > 0 then Fmt.pf ppf ","; pp_reg ppf d) ds;
    Fmt.pf ppf ")"
  | RET a -> Fmt.pf ppf "ret %a" pp_reg a
  | JMP b -> Fmt.pf ppf "jmp %s" (label b)
  | CBR (a, b1, b2) ->
    let l1 = label b1 in
    let l2 = label b2 in
    Fmt.pf ppf "cbr %a, %s, %s" pp_reg a l1 l2

let pp_program ppf prog =
  let basic_blocks = ref [ prog.entry, "START" ] in
  let work_list = ref [ prog.entry ] in
  let label bb =
    try List.assq bb !basic_blocks
    with Not_found ->
      let name = Fmt.str "L%d" (List.length !basic_blocks - 1) in
      basic_blocks := (bb, name) :: !basic_blocks;
      work_list := !work_list @ [ bb ];
      name
  in
  let rec loop i =
    match !work_list with
    | [] -> ()
    | bb :: rest ->
      work_list := rest;
      if i > 0 then Fmt.pf ppf ",";
      Fmt.pf ppf "%S:[" (label bb);
      List.iteri
        (fun i is ->
           if i > 0 then Fmt.pf ppf ",";
           let str = Fmt.str "%a" (pp_ins ~label) is in
           Fmt.pf ppf "%S" str)
        (instructions bb);
      Fmt.pf ppf "]";
      loop (i + 1)
  in
  Fmt.pf ppf "{";
  loop 0;
  Fmt.pf ppf "}"