(* instruction operand types, etc. *)

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 (* | NE | LE | GE *)

(* instruction types (suffix denotes number of successors) *)

type i0 =
  | Ret of arg

type i1 =
  | Mov of reg * arg
  | Opr of opr * reg * arg
  | Get of reg * loc
  | Set of loc * arg
  | Con of reg * Value.vtable
  | Loc of reg * reg * string
  | Cal of reg * loc * reg list

type i2 =
  | IfT of arg
  | IfC of cnd * reg * arg

(* bytecode graph nodes *)

type t = {
  mutable edge : edge;
  mutable label : string option;
  mutable preds : int;
}

and edge =
  | I0 of i0
  | I1 of i1 * t
  | I2 of i2 * t * t

let make edge = {
  edge;
  label = None;
  preds = 0;
}

let sucs = function
  | I0 _ -> []
  | I1 (_, t) -> [t]
  | I2 (_, t1, t2) -> [t1; t2]

let registers e =
  let arg = function #reg as r -> [r] | _ -> [] in
  let loc (r, o) = r :: arg o in
  match e with
  | I0 (Ret a) -> arg a
  | I1 (Mov (a, b), _) -> a :: arg b
  | I1 (Opr (_, a, b), _) -> a :: arg b
  | I1 (Get (a, b), _) -> a :: loc b
  | I1 (Set (a, b), _) -> loc a @ arg b
  | I1 (Con (a, _), _) -> [a]
  | I1 (Loc (a, b, _), _) -> [a; b]
  | I1 (Cal (a, f, bs), _) -> a :: loc f @ bs
  | I2 (IfT a, _, _) -> arg a
  | I2 (IfC (_, a, b), _, _) -> a :: arg b

let preorder t0 =
  let rec go t =
    t.preds <- t.preds + 1;
    if t.preds = 1 then
      t :: List.flatten (List.map go (sucs t.edge))
    else
      []
  in
  List.map (fun t -> t.preds <- 0; t) (go t0)

(* functions *)

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

type Value.mthd += 
  | Function of funct

let make_funct n_args entry =
  let frame_size =
    List.map (fun t -> registers t.edge) (preorder entry)
    |> List.flatten
    |> List.fold_left
         (fun fs (`R i) -> max fs (i + 1))
         (n_args + 1)
  in
  { n_args; frame_size; entry }

(* helper module for constructing and combining bytecode graphs *)

module B = struct
  type nonrec bcg = t
  type t = { build : bcg -> bcg } [@@unboxed]

  let empty =
    {build = Fun.id}

  let append t1 t2 =
    {build = fun b -> t1.build (t2.build b)}

  let concat ts =
    let ts_r = List.rev ts in
    {build = fun b -> List.fold_left (fun b t -> t.build b) b ts_r}

  (* let fix (f : t -> t) : t = *)
  (*   let _ = f in failwith "TODO: B.fix" *)

  module Infix = struct
    let ( +> ) = append
  end

  module Private = struct
    let i0 (i : i0) (b : t) : bcg =
      b.build (make (I0 i))

    let i1 (i : i1) : t =
      {build = fun t -> make (I1 (i, t))}

    let i2 (i : i2) (b1 : t) (b2 : t) : t =
      {build = fun b -> make (I2 (i, b1.build b, b2.build b))}
  end

  open Private

  let nil = `Cst Value.Nil
  let int64 x = `Cst (Value.Int x)
  let int x = int64 (Int64.of_int x)

  let mov dst src = i1 (Mov (dst, src))
  let opr op dst src = i1 (Opr (op, dst, src))
  let get dst loc = i1 (Get (dst, loc))
  let set loc src = i1 (Set (loc, src))
  let con dst vtb = i1 (Con (dst, vtb))
  let loc dst src nam = i1 (Loc (dst, src, nam))
  let cal dst fn args =
    (* TODO: check if fn,args well formed *)
    i1 (Cal (dst, fn, args))
  let ret v = i0 (Ret v)
  let if_ = function
    | #arg as x      -> i2 (IfT x)
    | `Cmp (c, x, y) -> i2 (IfC (c, x, y))

  let add = opr ADD
  let sub = opr SUB
  let mul = opr MUL
  let div = opr DIV
  let mod_= opr MOD
  let not_= opr NOT
  let neg = opr NEG
  let ceq = opr (Cmp EQ)
  let cgt = opr (Cmp GT)
  let clt = opr (Cmp LT)
end

(* pretty printer *)

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 v -> Value.pp ppf v

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 ppf (vtb : Value.vtable) =
  Fmt.pf ppf "(%d){" 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_cnd ~prefix = function
  | EQ -> prefix ^ "eq"
  | GT -> prefix ^ "gt"
  | LT -> prefix ^ "lt"

let string_of_opr = function
  | NOT -> "not"
  | NEG -> "neg"
  | ADD -> "add"
  | SUB -> "sub"
  | MUL -> "mul"
  | DIV -> "div"
  | MOD -> "mod"
  | Cmp c -> string_of_cnd c ~prefix:"c"

let pp_i0 ppf = function
  | Ret a -> Fmt.pf ppf "ret %a" pp_arg a

let pp_i1 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
  | Loc (a, b, nam) -> Fmt.pf ppf "loc %a, %a.<%s>" pp_reg a pp_reg b nam
  | Con (a, vtb) -> Fmt.pf ppf "con %a, %a" pp_reg a pp_vtable vtb
  | Cal (r, f, args) ->
    Fmt.pf ppf "cal %a, %a(" pp_reg r pp_loc f;
    List.iteri (fun i a -> if i > 0 then Fmt.pf ppf ","; pp_reg ppf a) args;
    Fmt.pf ppf ")"

let pp_i2 ppf = function
  | IfT v -> Fmt.pf ppf "btr %a" pp_arg v
  | IfC (c, a, b) ->
    let name = string_of_cnd c ~prefix:"b" in
    Fmt.pf ppf "%s %a, %a" name pp_reg a pp_arg b

let generate_labels ep =
  let nl = ref 0 in
  let rec go t require =
    t.preds <- t.preds + 1;
    if t.label = None && (t.preds > 1 || require) then begin
      t.label <- Some (Fmt.str "L%d" !nl);
      incr nl
    end;
    if t.preds = 1 then begin
      match t.edge with
      | I0 _ -> ()
      | I1 (_, t1) -> go t1 false
      | I2 (_, t1, t2) -> go t1 false; go t2 true
    end
  in
  ep.label <- Some "EP";
  go ep false

let dump ?(margin = 8) println main =
  let printf ?l fmt =
    let prefix = match l with
      | None -> ""
      | Some l -> l ^ ":"
    in
    Fmt.kstr println ("%-*s" ^^ fmt) margin prefix
  in

  let rec pr_code t =
    if t.preds = 0 then
      pr_jmp t
    else begin
      t.preds <- 0;
      match t.edge with
      | I0 i ->
        printf ?l:t.label "%a" pp_i0 i

      | I1 (i, t1) ->
        printf ?l:t.label "%a" pp_i1 i;
        pr_code t1

      | I2 (i, t1, t2) ->
        printf ?l:t.label "%a" pp_i2 i;
        pr_jmp t2;
        pr_code t1;
        maybe_pr_code t2
    end

  and pr_jmp t =
    printf "jmp %s" (Option.get t.label)

  and maybe_pr_code t =
    if t.preds > 0 then
      pr_code t

  in

  let pr_funct name fn =
    println (Fmt.str "# fun %s(%d)" name fn.n_args);
    generate_labels fn.entry;
    pr_code fn.entry
  in

  pr_funct "main" main