open Maps
open RTL
open BTL
open Registers
open DebugPrint

(* TODO gourdinl duplicated parts from RTLpathLivegenaux.ml
        This should be accessible everywhere. *)
let get_some = function
| None -> failwith "Got None instead of Some _"
| Some thing -> thing

let successors_inst = function
| Inop n | Iop (_,_,_,n) | Iload (_,_,_,_,_,n) | Istore (_,_,_,_,n) | Icall (_,_,_,_,n) | Ibuiltin (_,_,_,n) -> [n]
| Icond (_,_,n1,n2,_) -> [n1; n2]
| Ijumptable (_,l) -> l
| Itailcall _ | Ireturn _ -> []

let predicted_successor = function
| Inop n | Iop (_,_,_,n) | Iload (_,_,_,_,_,n) | Istore (_,_,_,_,n) -> Some n
| Icall (_,_,_,_,n) | Ibuiltin (_,_,_,n) -> None
| Icond (_,_,n1,n2,p) -> (
    match p with
    | Some true -> Some n1
    | Some false -> Some n2
    | None -> None )
| Ijumptable _ | Itailcall _ | Ireturn _ -> None

(* TODO gourdinl Adapted from RTLpathLivegenaux to avoid redundant computations ... *)
let non_predicted_successors i = function
  | None -> successors_inst i
  | Some ps -> List.filter (fun s -> s != ps) (successors_inst i)

(* adapted from Linearizeaux.get_join_points *)
let get_join_points code entry =
  let reached = ref (PTree.map (fun n i -> false) code) in
  let reached_twice = ref (PTree.map (fun n i -> false) code) in
  let rec traverse pc =
    if get_some @@ PTree.get pc !reached then begin
      if not (get_some @@ PTree.get pc !reached_twice) then
        reached_twice := PTree.set pc true !reached_twice
    end else begin
      reached := PTree.set pc true !reached;
      traverse_succs (successors_inst @@ get_some @@ PTree.get pc code)
    end
  and traverse_succs = function
    | [] -> ()
    | [pc] -> traverse pc
    | pc :: l -> traverse pc; traverse_succs l
  in traverse entry; !reached_twice

let translate_inst = function
    | Inop (_) -> Bnop
    | Iop (op,lr,rd,_) -> Bop (op,lr,rd)
    | Iload (trap,chk,addr,lr,rd,_) -> Bload (trap,chk,addr,lr,rd)
    | Istore (chk,addr,lr,rd,_) -> Bstore (chk,addr,lr,rd)
    | Icall (sign,fn,lr,rd,s) -> BF (Bcall (sign,fn,lr,rd,s))
    | Itailcall (sign,fn,lr) -> BF (Btailcall (sign,fn,lr))
    | Ibuiltin (ef,lr,rd,s) -> BF (Bbuiltin (ef,lr,rd,s))
    | Icond (cond,lr,ifso,ifnot,info) -> (
        (* TODO gourdinl remove this *)
        assert (info = None);
        Bcond (cond,lr,BF (Bgoto(ifso)),BF (Bgoto(ifnot)),info))
    | Ijumptable (arg, tbl) -> BF (Bjumptable (arg,tbl))
    | Ireturn (oreg) -> BF (Breturn (oreg))

let translate_function code entry join_points =
  let reached = ref (PTree.map (fun n i -> false) code) in
  let btl_code = ref PTree.empty in
  (* SEPARATE IN A BETTER WAY!! *)
  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 rec build_btl_block n =
        let inst = get_some @@ PTree.get n code in
        let psucc = predicted_successor inst 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) -> (
              (* TODO gourdinl remove this *)
              assert (s = ifso || s = ifnot);
              next_nodes := !next_nodes @ non_predicted_successors inst psucc;
              if s = ifso then
                Bcond (cond,lr,build_btl_block s,BF (Bgoto (ifnot)),info)
              else
                Bcond (cond,lr,BF (Bgoto (ifso)),build_btl_block s,info))
          | _ -> Bseq (translate_inst inst, build_btl_block s))
        | None -> (
            next_nodes := !next_nodes @ successors_inst inst;
            translate_inst inst)
      in
      let ib = build_btl_block e in
      let succs = !next_nodes in
      let ibf = { entry = ib; input_regs = Regset.empty } 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 btl = translate_function code entry join_points in
  debug_flag := true;
  debug"Code: ";
  print_code code;
  debug "Join points: ";
  print_true_nodes join_points;
  debug "\n";
  debug_flag := false;
  ((btl, entry), PTree.empty)