open RTL
open RTLpath
open Registers
open Maps
open Camlcoq
open Datatypes
open DebugPrint
open RTLcommonaux

let rec list_to_regset = function
  | [] -> Regset.empty
  | r::l -> Regset.add r (list_to_regset l)

let get_input_regs i = 
  let empty = Regset.empty in
  match i with
  | Inop _ -> empty
  | Iop (_,lr,_,_) | Iload (_,_,_,lr,_,_) | Icond (_,lr,_,_,_) -> list_to_regset lr
  | Istore (_,_,lr,r,_) -> Regset.add r (list_to_regset lr)
  | Icall (_, ri, lr, _, _) | Itailcall (_, ri, lr) -> begin
      let rs = list_to_regset lr in
      match ri with
      | Coq_inr _ -> rs
      | Coq_inl r -> Regset.add r rs
    end
  | Ibuiltin (_, lbr, _, _) -> list_to_regset @@ AST.params_of_builtin_args lbr
  | Ijumptable (r, _) -> Regset.add r empty
  | Ireturn opr -> (match opr with Some r -> Regset.add r empty | None -> empty)

let get_output_reg i =
  match i with
  | Inop _ | Istore _ | Icond _ | Itailcall _ | Ijumptable _ | Ireturn _ -> None
  | Iop (_, _, r, _) | Iload (_, _, _, _, r, _) | Icall (_, _, _, r, _) -> Some r
  | Ibuiltin (_, _, brr, _) -> (match brr with AST.BR r -> Some r | _ -> None)

(* Does not set the input_regs and liveouts field *)
let get_path_map code entry join_points =
  let visited = ref (PTree.map (fun n i -> false) code) in
  let path_map = ref PTree.empty in
  let rec dig_path e =
    if (get_some @@ PTree.get e !visited) then
      ()
    else begin
      visited := PTree.set e true !visited;
      let psize = ref (-1) in
      let path_successors = ref [] in
      let rec dig_path_rec n : (path_info * node list) option =
        let inst = get_some @@ PTree.get n code in
        begin
          psize := !psize + 1;
          let psucc = predicted_successor inst in
          let successor = match psucc with
          | None -> None
          | Some n' -> if get_some @@ PTree.get n' join_points then None else Some n'
          in match successor with
          | Some n' -> begin
                path_successors := !path_successors @ non_predicted_successors inst psucc;
                dig_path_rec n'
              end
          | None -> Some ({ psize = (Camlcoq.Nat.of_int !psize);
                            input_regs = Regset.empty; pre_output_regs = Regset.empty; output_regs = Regset.empty },
                        !path_successors @ successors_inst inst)
        end
      in match dig_path_rec e with
      | None -> ()
      | Some ret ->
          let (path_info, succs) = ret in
          begin
            path_map := PTree.set e path_info !path_map;
            List.iter dig_path succs
          end
    end
  in begin
    dig_path entry;
    !path_map
  end

(** OLD CODE - If needed to have our own kildall

let transfer after = let open Liveness in function
  | Inop _ -> after
  | Iop (_, args, res, _) ->
      reg_list_live args (Regset.remove res after)
  | Iload (_, _, _, args, dst, _) ->
      reg_list_live args (Regset.remove dst after)
  | Istore (_, _, args, src, _) ->
      reg_list_live args (Regset.add src after)
  | Icall (_, ros, args, res, _) ->
      reg_list_live args (reg_sum_live ros (Regset.remove res after))
  | Itailcall (_, ros, args) ->
      reg_list_live args (reg_sum_live ros Regset.empty)
  | Ibuiltin (_, args, res, _) ->
      reg_list_live (AST.params_of_builtin_args args)
        (reg_list_dead (AST.params_of_builtin_res res) after)
  | Icond (_, args, _, _, _) ->
      reg_list_live args after
  | Ijumptable (arg, _) ->
      Regset.add arg after
  | Ireturn optarg ->
      reg_option_live optarg Regset.empty

let get_last_nodes f =
  let visited = ref (PTree.map (fun n i -> false) f.fn_code) in
  let rec step n =
    let inst = get_some @@ PTree.get n f.fn_code in
    let successors = successors_inst inst in
    if get_some @@ PTree.get n !visited then []
    else begin

let analyze f =
  let liveness = ref (PTree.map (fun n i -> None) f.fn_code) in
  let predecessors = Duplicateaux.get_predecessors_rtl f.fn_code in
  let last_nodes = get_last_nodes f in
  let rec step liveout n = (* liveout is the input_regs from the successor *)
    let inst = get_some @@ PTree.get n f.fn_code in
    let continue = ref true in
    let alive = match get_some @@ PTree.get n !liveness with
    | None -> transfer liveout inst
    | Some pre_alive -> begin
        let union = Regset.union pre_alive liveout in
        let new_alive = transfer union inst in
        (if Regset.equal pre_alive new_alive then continue := false);
        new_alive
      end
    in begin
      liveness := PTree.set n (Some alive) !liveness;
      if !continue then
        let preds = get_some @@ PTree.get n predecessors in
        List.iter (step alive) preds
    end
  in begin
    List.iter (step Regset.empty) last_nodes;
    let liveness_noopt = PTree.map (fun n i -> get_some i) !liveness in
    begin
      debug_flag := true;
      dprintf "Liveness: "; print_ptree_regset liveness_noopt; dprintf "\n";
      debug_flag := false;
      liveness_noopt
    end
  end
*)

let rec traverse code n size =
  let inst = get_some @@ PTree.get n code in
  if (size == 0) then (inst, n)
  else 
    let n' = get_some @@ predicted_successor inst in
    traverse code n' (size-1)

let get_outputs liveness f n pi =
  let (last_instruction, pc_last) = traverse f.fn_code n (Camlcoq.Nat.to_int pi.psize) in
  let path_last_successors = successors_inst last_instruction in
  let list_input_regs = List.map (
      fun n -> get_some @@ PTree.get n liveness
    ) path_last_successors in
  let outputs = List.fold_left Regset.union Regset.empty list_input_regs in
  let por = match last_instruction with (* see RTLpathLivegen.final_inst_checker *)
    | Icall (_, _, _, res, _) -> Regset.remove res outputs
    | Ibuiltin (_, _, res, _) ->  Liveness.reg_list_dead (AST.params_of_builtin_res res) outputs
    | Itailcall (_, _, _) | Ireturn _ ->
       assert (outputs = Regset.empty); (* defensive check for performance *)
       outputs
    | _ -> outputs
  in (por, outputs)

let set_pathmap_liveness f pm =
  let liveness = analyze f in
  let new_pm = ref PTree.empty in
  begin
    debug "Liveness: "; print_ptree_regset liveness; debug "\n";
    List.iter (fun (n, pi) ->
      let inputs = get_some @@ PTree.get n liveness in
      let (por, outputs) = get_outputs liveness f n pi in
      new_pm := PTree.set n
                  {psize=pi.psize; input_regs=inputs; pre_output_regs=por; output_regs=outputs} !new_pm
    ) (PTree.elements pm);
    !new_pm
  end

let print_path_info pi = begin
  debug "(psize=%d; " (Camlcoq.Nat.to_int pi.psize);
  debug "\ninput_regs=";
  print_regset pi.input_regs;
  debug "\n; pre_output_regs=";
  print_regset pi.pre_output_regs;
  debug "\n; output_regs=";
  print_regset pi.output_regs;
  debug ")\n"
end

let print_path_map path_map = begin
  debug "[";
  List.iter (fun (n,pi) ->
    debug "\n\t";
    debug "%d: " (P.to_int n);
    print_path_info pi
  ) (PTree.elements path_map);
  debug "]"
end

let build_path_map f = 
  let code = f.fn_code in
  let entry = f.fn_entrypoint in
  let join_points = get_join_points code entry in
  let path_map = set_pathmap_liveness f @@ get_path_map code entry join_points in
  begin
    debug "Join points: ";
    print_true_nodes join_points;
    debug "\nPath map: ";
    print_path_map path_map;
    debug "\n";
    path_map
  end