type regidx = int

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

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
  | JMP of basic_block
  | BRT 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 }

let make_program entrypoint = { entrypoint }

(* 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_ins ~get_bb_name 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
  | RET -> Fmt.pf ppf "ret"
  | JMP l -> Fmt.pf ppf "jmp %s" (get_bb_name l)
  | _ -> failwith "..."

let pp_program ppf pr =
  let ep = pr.entrypoint in
  let basic_blocks = ref [ ep, "START" ] in
  let work_list = ref [ ep ] in
  let get_bb_name bb =
    match List.find (fun (bb', _) -> bb == bb') !basic_blocks with
    | _, name -> name
    | exception Not_found ->
        let name = Fmt.str "L%d" (List.length !basic_blocks - 1) in
        basic_blocks := (bb, name) :: !basic_blocks;
        work_list := bb :: !work_list;
        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:[" (get_bb_name bb);
        List.iteri
          (fun i is ->
            if i > 0 then Fmt.pf ppf ",";
            let str = Fmt.str "%a" (pp_ins ~get_bb_name) is in
            Fmt.pf ppf "%S" str)
          (instructions bb);
        Fmt.pf ppf "]";
        loop (i + 1)
  in
  Fmt.pf ppf "{";
  loop 0;
  Fmt.pf ppf "}"