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 registers = function
  | JMP _ -> []
  | RET r
  | LDI (r, _)
  | CON (r, _)
  | CBR (r, _, _) -> [r]
  | LDR (r1, r2)
  | NOT (r1, r2)
  | LOC (r1, r2, _) -> [r1; r2]
  | 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) -> [r1; r2; r3]
  | CAL (r1, r2, r3, rs) -> r1::r2::r3::rs

let make_block () =
  { ins_list_rev = [] }

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

let instructions b =
  List.rev b.ins_list_rev

let iter_blocks_df f b0 =
  let queue = ref [ b0 ] in
  let visited = ref !queue in
  let enqueue b =
    if not (List.memq b !visited) then (
      queue := !queue @ [b];
      visited := b :: !visited)
  in
  let rec loop () =
    match !queue with
    | [] -> ()
    | b :: rest ->
      queue := rest;
      f b;
      (* NOTE: only [List.hd b.ins_list_rev] should be a branching instruction, so iterating
         the whole list is pointless. but just to be safe ... *)
      List.iter
        (function
          | JMP b1 -> enqueue b1
          | CBR (_, b1, b2) -> enqueue b1; enqueue b2
          | _ -> ())
        b.ins_list_rev;
      loop ()
  in
  loop ()


type funct =
  { n_args : int;
    frame_size : int;
    entry : block }

type Value.mthd +=
  | Function of funct

let make_funct n_args entry =
  let frame_size =
    let fsize = ref (n_args + 1) in
    iter_blocks_df
      (fun b ->
         fsize :=
           List.rev_map registers b.ins_list_rev
           |> List.flatten
           |> List.fold_left (fun fs (R i) -> max fs (i + 1))
                !fsize)
      entry;
    !fsize
  in
  { n_args; frame_size; entry }


(* pretty printing *)

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

let pp_vtable ~tbname ppf (vtb : vtable) =
  Fmt.pf ppf "%s(%d){" (tbname vtb) vtb.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 := ",")
    vtb.elems;
  Fmt.pf ppf "}"

let pp_ins ~tbname ~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 "loc %a, %a.%s" pp_reg a pp_reg b el
  | CON (a, vt) -> Fmt.pf ppf "con %a, %a" pp_reg a (pp_vtable ~tbname) 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 dump ?(recursive = true) println main_fn =
  let tbqueue = ref [] in
  let tbnames = ref [] in
  let tbname vtb =
    try List.assq vtb !tbnames
    with Not_found ->
      if recursive then tbqueue := !tbqueue @ [vtb];
      let n = List.length !tbnames in
      let l = Fmt.str "$tbl%d" n in
      tbnames := (vtb, l) :: !tbnames; l
  in

  let dump_fn fn =
    let labels = ref [ fn.entry, "ENTRY" ] in
    let label b =
      try List.assq b !labels
      with Not_found ->
        let n = List.length !labels - 1 in
        let l = Fmt.str "B%d" n in
        labels := (b, l) :: !labels; l
    in
    let pp_ins = pp_ins ~tbname ~label in
    iter_blocks_df
      (fun b ->
         List.fold_left
           (fun pfx ins ->
              println (Fmt.str "%-8s%a" pfx pp_ins ins);
              "")
           (label b ^ ":")
           (instructions b)
         |> ignore)
      fn.entry
  in

  println "# fun main(0)";
  dump_fn main_fn;

  let rec loop () =
    match !tbqueue with
    | [] -> ()
    | vtb :: rest ->
      tbqueue := rest;
      Hashtbl.iter
        (fun fname -> function
           | Value.Field _ -> ()
           | Value.Method i ->
             match vtb.mthds.(i) with
             | Function fn ->
               println "";
               println (Fmt.str "# fun %s.%s(%d)" (tbname vtb) fname fn.n_args);
               dump_fn fn
             | _ -> ())
        vtb.elems;
      loop ()
  in
  loop ()