type regidx = int

type operand =
  | Cst_nil
  | Cst_true
  | Cst_false
  | Cst_int of int64
  | Reg of regidx

let cst_of_int i = Cst_int (Int64.of_int i)

type basic_block = {
  mutable ins_builder : ins list;
  mutable ins_list : ins list;
      (* bc_pc : int *)
      (* bc_len : int *)
}

and ins =
  | MOV of regidx * operand
  | ADD of regidx * operand
  | SUB of regidx * operand
  | MUL of regidx * operand
  | EQL of regidx * operand
  | LST of regidx * operand
  | GRT of regidx * operand
  | NOT of regidx
  | CON of regidx * Value.vtable
  | SLT of regidx * regidx * string
  | GET of regidx * regidx
  | SET of regidx * regidx
  | CLL of regidx * regidx * int
  | JMP of basic_block
  | CBR of operand * basic_block * basic_block
  | RET

let make_basic_block ins_list = { ins_builder = List.rev ins_list; ins_list }

let instructions bb =
  (* memoize computing "rev ins_builder" by storing result in ins_list *)
  if bb.ins_list = [] then bb.ins_list <- List.rev bb.ins_builder;
  bb.ins_list

let add_ins bb is =
  (* "append" instruction by prepending to ins_builder list *)
  bb.ins_builder <- is :: bb.ins_builder;
  (* invalidate the cache *)
  bb.ins_list <- []

type program = { entrypoint : basic_block }

type Value.mthd += Method of program

let make_program entrypoint = { entrypoint }

let frame_size prog =
  let visited = ref [] in
  let work_list = ref [ prog.entrypoint ] in
  let enqueue bb = if not (List.memq bb !visited) then work_list := bb :: !work_list in
  let reg acc i = max acc (i + 1) in
  let op acc = function
    | Reg i -> reg acc i
    | _ -> acc
  in
  let ins acc = function
    | MOV (r, v)
    | ADD (r, v)
    | SUB (r, v)
    | MUL (r, v)
    | EQL (r, v)
    | LST (r, v)
    | GRT (r, v) -> op (reg acc r) v
    | CON (r, _) | NOT r -> reg acc r
    | GET (o, s) | SLT (o, s, _) -> reg (reg acc o) s
    | SET (o, s) -> reg (reg acc o) (s + 1)
    | CBR (v, b1, b2) ->
        enqueue b1;
        enqueue b2;
        op acc v
    | CLL (o, m, k) -> reg (reg acc o) (m + k + 1)
    | JMP b ->
        enqueue b;
        acc
    | RET -> acc
  in
  let rec loop acc =
    match !work_list with
    | [] -> acc
    | bb :: rest ->
        visited := bb :: !visited;
        work_list := rest;
        List.fold_left ins acc (instructions bb) |> loop
  in
  loop 1

(* pretty printing *)

let pp_reg ppf r = Fmt.pf ppf "$%d" r

let pp_operand ppf = function
  | Cst_nil -> Fmt.pf ppf "nil"
  | Cst_true -> Fmt.pf ppf "true"
  | Cst_false -> Fmt.pf ppf "false"
  | Cst_int n -> Fmt.pf ppf "#%s" (Int64.to_string n)
  | Reg r -> pp_reg ppf r

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
  | MOV (l, r) -> Fmt.pf ppf "mov %a, %a" pp_reg l pp_operand r
  | ADD (l, r) -> Fmt.pf ppf "add %a, %a" pp_reg l pp_operand r
  | SUB (l, r) -> Fmt.pf ppf "sub %a, %a" pp_reg l pp_operand r
  | MUL (l, r) -> Fmt.pf ppf "mul %a, %a" pp_reg l pp_operand r
  | EQL (l, r) -> Fmt.pf ppf "eql %a, %a" pp_reg l pp_operand r
  | LST (l, r) -> Fmt.pf ppf "lst %a, %a" pp_reg l pp_operand r
  | GRT (l, r) -> Fmt.pf ppf "grt %a, %a" pp_reg l pp_operand r
  | NOT l -> Fmt.pf ppf "not %a" pp_reg l
  | CON (l, vt) -> Fmt.pf ppf "con %a, %a" pp_reg l pp_vtable vt
  | SLT (o, s, n) -> Fmt.pf ppf "mov %a, @%a.%s" pp_reg s pp_reg o n
  | GET (o, s) -> Fmt.pf ppf "mov %a, %a[%a]" pp_reg s pp_reg o pp_reg s
  | SET (o, s) -> Fmt.pf ppf "mov %a[%a], %a" pp_reg o pp_reg s pp_reg (s + 1)
  | CLL (o, m, k) ->
      Fmt.pf ppf "cll %a[%a](" pp_reg o pp_reg m;
      for i = 1 to k do
        if i > 1 then Fmt.pf ppf ",";
        Fmt.pf ppf "%a" pp_reg (m + i)
      done;
      Fmt.pf ppf ")"
  | CBR (v, b1, b2) ->
      let l1 = label b1 in
      let l2 = label b2 in
      Fmt.pf ppf "cbr %a, %s, %s" pp_operand v l1 l2
  | RET -> Fmt.pf ppf "ret"
  | JMP l -> Fmt.pf ppf "jmp %s" (label l)

let pp_program ppf pr =
  let ep = pr.entrypoint in
  let basic_blocks = ref [ ep, "START" ] in
  let work_list = ref [ ep ] 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 "}"