open Maps
open RTL
open BTL
open PrintBTL
open RTLcommonaux
open DebugPrint
open BTLtypes
open BTLcommonaux

let encaps_final inst osucc =
  match inst with
  | BF _ -> inst
  | _ -> Bseq (inst, BF (Bgoto (get_some @@ osucc), def_iinfo ()))

let translate_inst (iinfo : BTL.inst_info) inst is_final =
  let osucc = ref None in
  let btli =
    match inst with
    | Inop s ->
        osucc := Some s;
        Bnop (Some iinfo)
    | Iop (op, lr, rd, s) ->
        osucc := Some s;
        Bop (op, lr, rd, iinfo)
    | Iload (trap, chk, addr, lr, rd, s) ->
        osucc := Some s;
        Bload (trap, chk, addr, lr, rd, iinfo)
    | Istore (chk, addr, lr, src, s) ->
        osucc := Some s;
        Bstore (chk, addr, lr, src, iinfo)
    | Icall (sign, fn, lr, rd, s) -> BF (Bcall (sign, fn, lr, rd, s), iinfo)
    | Itailcall (sign, fn, lr) -> BF (Btailcall (sign, fn, lr), iinfo)
    | Ibuiltin (ef, lr, rd, s) -> BF (Bbuiltin (ef, lr, rd, s), iinfo)
    | Icond (cond, lr, ifso, ifnot, info) ->
        osucc := Some ifnot;
        iinfo.opt_info <- info;
        Bcond
          ( cond,
            lr,
            BF (Bgoto ifso, def_iinfo ()),
            Bnop None,
            iinfo )
    | Ijumptable (arg, tbl) -> BF (Bjumptable (arg, tbl), iinfo)
    | Ireturn oreg -> BF (Breturn oreg, iinfo)
  in
  if is_final then encaps_final btli !osucc else btli

let translate_function code entry join_points liveness (typing : RTLtyping.regenv) =
  let reached = ref (PTree.map (fun n i -> false) code) in
  let btl_code = ref PTree.empty in
  let rec build_btl_tree e =
    if get_some @@ PTree.get e !reached then ()
    else (
      reached := PTree.set e true !reached;
      let next_nodes = ref [] in
      let last = ref None in
      let rec build_btl_block n =
        let inst = get_some @@ PTree.get n code in
        let psucc = predicted_successor inst in
        let iinfo = mk_iinfo (p2i n) None in
        let succ =
          match psucc with
          | Some ps ->
              if get_some @@ PTree.get ps join_points then None else Some ps
          | None -> None
        in
        match succ with
        | Some s -> (
            match inst with
            | Icond (cond, lr, ifso, ifnot, info) ->
                assert (s = ifnot);
                next_nodes := !next_nodes @ non_predicted_successors inst psucc;
                iinfo.opt_info <- info;
                Bseq
                  ( Bcond
                      (cond, lr, BF (Bgoto ifso, def_iinfo ()), Bnop None, iinfo),
                    build_btl_block s )
            | _ -> Bseq (translate_inst iinfo inst false, build_btl_block s))
        | None ->
            debug "BLOCK END.\n";
            next_nodes := !next_nodes @ successors_inst inst;
            last := Some inst;
            translate_inst iinfo inst true
      in
      let ib = build_btl_block e in
      let succs = !next_nodes in

      let inputs = get_some @@ PTree.get e liveness in
      let soutput = get_outputs liveness e (get_some @@ !last) in

      let bi = mk_binfo (p2i e) soutput typing in
      let ibf = { entry = ib; input_regs = inputs; binfo = bi } in
      btl_code := PTree.set e ibf !btl_code;
      List.iter build_btl_tree succs)
  in
  build_btl_tree entry;
  !btl_code

let rtl2btl (f : RTL.coq_function) =
  let code = f.fn_code in
  let entry = f.fn_entrypoint in
  let join_points = get_join_points code entry in
  let liveness = analyze f in
  let typing = get_ok @@ RTLtyping.type_function f in
  let btl = translate_function code entry join_points liveness typing in
  let dm = PTree.map (fun n i -> n) btl in
  (*debug_flag := true;*)
  debug "Entry %d\n" (p2i entry);
  debug "RTL Code: ";
  print_code code;
  debug "BTL Code:\n";
  print_btl_code stderr btl;
  (*debug_flag := false;*)
  debug "Dupmap:\n";
  print_ptree print_pint dm;
  debug "Join points: ";
  print_true_nodes join_points;
  debug "\n";
  ((btl, entry), dm)