module Ast = Spice_syntax.Ast

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

type reg = [`R of int]
type cst = [`Cst of Value.t]
type arg = [reg | cst]
type ofs = [reg | `Ofs of int]
type loc = reg * ofs

type opr = NOT | NEG | ADD | SUB | MUL | DIV | MOD | Cmp of cnd
and  cnd = EQ | LT | GT
(* and  cnd = EQ | NE | LT | GE | GT | LE *)

type ins =
  | Mov of reg * arg
  | Opr of opr * reg * arg
  | Get of reg * loc
  | Set of loc * arg
  | Con of reg * vtable
  | Loc of reg * reg * string
  | Cal of reg * loc * reg list
  | Btr of arg * block * block
  | Jmp of block
  | Ret of arg

and block =
  { mutable ins_list_rev : ins list }

let arg_regs = function
  | #reg as r -> [r]
  | #cst -> []

let loc_regs = function
  | (r1, (#reg as r2)) -> [r1; r2]
  | (r1, #ofs)         -> [r1]

let registers = function
  | Ret v
  | Btr (v, _, _)   -> arg_regs v
  | Mov (r, v)
  | Opr (_, r, v)   -> r :: arg_regs v
  | Get (r, l)      -> r :: loc_regs l
  | Set (l, v)      -> arg_regs v @ loc_regs l
  | Con (r, _)      -> [r]
  | Loc (r1, r2, _) -> [r1; r2]
  | Cal (r, l, rs)  -> loc_regs l @ r :: rs
  | Jmp _           -> []

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
          | Btr (_, 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_arg ppf = function
  | #reg as r -> pp_reg ppf r
  | `Cst c -> Value.pp ppf c

let pp_loc ppf = function
  | (r, (#reg as i)) -> Fmt.pf ppf "%a[%a]" pp_reg r pp_reg i
  | (r, `Ofs ofs)    -> Fmt.pf ppf "%a[%d]" pp_reg r ofs

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 string_of_opr = function
  | NOT -> "not"
  | NEG -> "neg"
  | ADD -> "add"
  | SUB -> "sub"
  | MUL -> "mul"
  | DIV -> "div"
  | MOD -> "mod"
  | Cmp EQ -> "ceq"
  | Cmp LT -> "clt"
  | Cmp GT -> "cgt"
  (* | Cmp NE -> "cne" *)
  (* | Cmp GE -> "cge" *)
  (* | Cmp LE -> "cle" *)

let pp_ins ~tbname ~label ppf = function
  | Mov (a, b) -> Fmt.pf ppf "mov %a, %a" pp_reg a pp_arg b
  | Opr (o, a, b) -> Fmt.pf ppf "%s %a, %a" (string_of_opr o) pp_reg a pp_arg b
  | Get (a, b) -> Fmt.pf ppf "mov %a, %a" pp_reg a pp_loc b
  | Set (a, b) -> Fmt.pf ppf "mov %a, %a" pp_loc a pp_arg b
  | Con (a, vtb) -> Fmt.pf ppf "con %a, %a" pp_reg a (pp_vtable ~tbname) vtb
  | Loc (a, b, nam) -> Fmt.pf ppf "loc %a, %a.<%s>" pp_reg a pp_reg b nam
  | Jmp b -> Fmt.pf ppf "jmp %s" (label b)
  | Ret a -> Fmt.pf ppf "ret %a" pp_arg a
  | Btr (a, b1, b2) ->
    let l1 = label b1 in
    let l2 = label b2 in
    Fmt.pf ppf "btr %a, %s, %s" pp_arg a l1 l2
  | Cal (a, f, args) ->
    Fmt.pf ppf "cal %a, %a(" pp_reg a pp_loc f;
    List.iteri (fun i d -> if i > 0 then Fmt.pf ppf ","; pp_reg ppf d) args;
    Fmt.pf ppf ")"

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 ()