open RTL
open Maps
open Camlcoq
open Registers
open Kildall
open Lattice

let p2i r = P.to_int r

let i2p i = P.of_int i

let get_some = function
  | None -> failwith "Got None instead of Some _"
  | Some thing -> thing

let get_ok r = match r with Errors.OK x -> x | _ -> failwith "Did not get OK"

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

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 (
      if not (get_some @@ PTree.get pc !reached_twice) then
        reached_twice := PTree.set pc true !reached_twice)
    else (
      reached := PTree.set pc true !reached;
      traverse_succs (successors_inst @@ get_some @@ PTree.get pc code))
  and traverse_succs = function
    | [] -> ()
    | [ pc ] -> traverse pc
    | pc :: l ->
        traverse pc;
        traverse_succs l
  in
  traverse entry;
  !reached_twice

let transfer f pc after =
  let open Liveness in
  match PTree.get pc f.fn_code with
  | Some i -> (
      match i with
      | 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)
  | None -> Regset.empty

module RegsetLat = LFSet (Regset)
module DS = Backward_Dataflow_Solver (RegsetLat) (NodeSetBackward)

let analyze f =
  let liveouts =
    get_some @@ DS.fixpoint f.fn_code successors_instr (transfer f)
  in
  PTree.map
    (fun n _ ->
      let lo = PMap.get n liveouts in
      transfer f n lo)
    f.fn_code

let get_outputs liveness n last =
  let path_last_successors = successors_inst last in
  let list_input_regs =
    List.map (fun n -> get_some @@ PTree.get n liveness) path_last_successors
  in
  List.fold_left Regset.union Regset.empty list_input_regs