open Analysis

module Variable = struct
  type t = string
  let pp (ppf: out_channel) (v: t) : unit =
    Printf.fprintf ppf "%s" v

  let pp_list (ppf: out_channel) (vv: t list) : unit =
    List.iter (Printf.fprintf ppf "%s, ") vv

  let compare a b =
    String.compare a b
end

module RISCCfg = CfgRISC.RISCCfg

module DVCfg = Dataflow.Make (CfgRISC.RISCSimpleStatements) (Variable)
module DVCeltSet = Set.Make(Variable)


let variables (instr : DVCfg.elt) : DVCfg.internal list =
  let aux (acc: DVCeltSet.t) (instr: DVCfg.elt) =
    match instr with
    | Nop ->
      acc
    | BRegOp (_, r1, r2, r3) ->
      DVCeltSet.add r1.index acc |>
      DVCeltSet.add r2.index |>
      DVCeltSet.add r3.index
    | BImmOp (_, r1, _, r3)
    | URegOp (_, r1, r3)
    | Load (r1, r3)
    | Store (r1, r3) ->
      DVCeltSet.add r1.index acc |>
      DVCeltSet.add r3.index
    | LoadI (_, r3) ->
      DVCeltSet.add r3.index acc
  in

  aux DVCeltSet.empty instr |> DVCeltSet.to_list

let variables_all (instructions : DVCfg.elt list) : DVCfg.internal list =
  List.fold_left (fun (acc: DVCeltSet.t) (instr: DVCfg.elt) ->
      DVCeltSet.union acc (variables instr |> DVCeltSet.of_list)
    ) DVCeltSet.empty instructions |> DVCeltSet.to_list

let variables_used (instr : DVCfg.elt) : DVCfg.internal list =
  let aux (acc: DVCeltSet.t) (instr: DVCfg.elt) =
    match instr with
    | Nop
    | LoadI (_, _) ->
      acc
    | Store (r1, r2)
    | BRegOp (_, r1, r2, _) ->
      DVCeltSet.add r1.index acc |>
      DVCeltSet.add r2.index
    | BImmOp (_, r1, _, _)
    | URegOp (_, r1, _)
    | Load (r1, _) ->
      DVCeltSet.add r1.index acc
  in

  aux DVCeltSet.empty instr |> DVCeltSet.to_list

let variables_used_all (instructions : DVCfg.elt list) : DVCfg.internal list =
  List.fold_left (fun (acc: DVCeltSet.t) (instr: DVCfg.elt) ->
      DVCeltSet.union acc (variables_used instr |> DVCeltSet.of_list)
    ) DVCeltSet.empty instructions |> DVCeltSet.to_list


let variables_defined (instructions : DVCfg.elt) : DVCfg.internal list =
  let aux (acc: DVCeltSet.t) (instr: DVCfg.elt) =
    match instr with
    | Nop
    | Store (_, _) -> acc
    | BRegOp (_, _, _, r3)
    | BImmOp (_, _, _, r3)
    | URegOp (_, _, r3)
    | Load (_, r3)
    | LoadI (_, r3) ->
      DVCeltSet.add r3.index acc
  in

  aux DVCeltSet.empty instructions |> DVCeltSet.to_list


(* init function, assign the bottom to everything *)
let _init_bottom : (DVCfg.elt list -> DVCfg.internal_node) =
  (fun l -> {internal_in = [];
             internal_out = [];
             internal_between = (List.init (List.length l) (fun _ -> ([], [])))}
  )

(* init function, assign the top to everything *)
let init_top (all_variables) : (DVCfg.elt list -> DVCfg.internal_node) =
  (fun l -> {internal_in = all_variables;
             internal_out = all_variables;
             internal_between = (List.init (List.length l)
                                  (fun _ -> (all_variables, all_variables)))})


let lub (t: DVCfg.t) (node: Cfg.Node.t) : DVCfg.internal_node =
  let prev_internal_var = Cfg.NodeMap.find node t.internal_var in
  let code = match Cfg.NodeMap.find_opt node t.t.content with
      None -> []
    | Some c -> c
  in

  let new_internal_between = (
    List.map
      (fun (code, (i, _o)) ->
         (i, Utility.unique_union i (variables_defined code)))
      (List.combine code prev_internal_var.internal_between)
  ) in

  let new_internal_out =
    match new_internal_between with
    | [] ->
      prev_internal_var.internal_in
    | _ ->
      let _, newinternalout = (Utility.last_list new_internal_between) in
      newinternalout
  in

  { prev_internal_var with
    internal_between = new_internal_between;
    internal_out = new_internal_out }

let lucf (t: DVCfg.t) (node: Cfg.Node.t) : DVCfg.internal_node =
  let prev_internal_var = Cfg.NodeMap.find node t.internal_var in

  if Option.equal (=) (Some node) t.t.initial then
    (* if L is initial set dvin to the "in" register *)
    let new_internal_in = (
      match t.t.input_output_var with
        Some (i, _) -> [i]
      | None -> []
    ) in

    let new_internal_between = (
      (* set the dvin of each to the previous dvout *)
      match prev_internal_var.internal_between with
        [] -> []
      | [(_i, o)] -> [(new_internal_in, o)]
      | (_i, o) :: btwrest ->
        (new_internal_in, o) :: (
          List.map (fun ((_i, o), (_previ, prevo)) -> (prevo, o))
            (Utility.combine_twice btwrest prev_internal_var.internal_between)
        )
    ) in
    { prev_internal_var with
      internal_in = new_internal_in;
      internal_between = new_internal_between }
  else
    (* if L is not initial set dvin to the intersection of the previous node's
       dvouts *)
    let prev_nodes = Cfg.NodeMap.find node t.t.reverse_edges in
    let new_internal_in = (
      match prev_nodes with
      | [] ->
        []
      | [prevnode] ->
        (Cfg.NodeMap.find prevnode t.internal_var).internal_out
      | [prevnode1; prevnode2] ->
        Utility.unique_intersection
          (Cfg.NodeMap.find prevnode1 t.internal_var).internal_out
          (Cfg.NodeMap.find prevnode2 t.internal_var).internal_out
      | prevnode :: restnodes ->
        List.fold_left (* intersection of all previous nodes' dvout *)
          (fun acc prevnode ->
             Utility.unique_intersection
               acc
               (Cfg.NodeMap.find prevnode t.internal_var).internal_out)
          (Cfg.NodeMap.find prevnode t.internal_var).internal_out
          restnodes
    ) in

    let new_internal_between =
      match prev_internal_var.internal_between with
        [] -> []
      | [(_i, o)] -> [(new_internal_in, o)]
      | (_i, o) :: btwrest ->
        (new_internal_in, o) :: (
          List.map (fun ((_i, o), (_previ, prevo)) -> (prevo, o))
            (Utility.combine_twice btwrest prev_internal_var.internal_between)
        )
    in
    { prev_internal_var with
      internal_in = new_internal_in;
      internal_between = new_internal_between }


let update (t: DVCfg.t) (node: Cfg.Node.t) : DVCfg.internal_node =
  let newt =
    {t with internal_var = (Cfg.NodeMap.add node (lucf t node) t.internal_var)}
  in
  lub newt node


let compute_defined_variables (cfg: RISCCfg.t) : DVCfg.t =
  (* creates the DVCfg structure and finds the fixed point *)
  let all_variables = List.fold_left
      (fun acc (_, code) ->
         Utility.unique_union acc (variables_all code))
      []
      (Cfg.NodeMap.to_list cfg.content)
  in
  let all_variables =
    match cfg.input_output_var with
    | None -> all_variables
    | Some (i, o) -> Utility.unique_union all_variables [i;o]
  in
  DVCfg.from_cfg cfg
  |> DVCfg.fixed_point ~init:(init_top all_variables) ~update:update



let check_undefined_variables (dvcfg: DVCfg.t) : Variable.t list option =
  (* returns all undefined variables previously computed *)
  let aux (node: Cfg.Node.t) (dvcfg: DVCfg.t) : Variable.t list option =
    let code = match Cfg.NodeMap.find_opt node dvcfg.t.content with
        None -> []
      | Some c -> c
    in
    let internal_var = Cfg.NodeMap.find node dvcfg.internal_var in
    let vua = variables_used_all code in

    let outvar =
      match (Option.equal (=) (Some node) dvcfg.t.terminal,
             dvcfg.t.input_output_var,
             internal_var.internal_out) with
      | (true, Some (_, outvar), vout) ->
        if List.mem outvar vout
        then None
        else Some outvar
      | (_, _, _) ->
        None
    in

    if Utility.inclusion vua (internal_var.internal_in) then
      match outvar with None -> None
                      | Some outvar -> Some [outvar]
    else
      (* the variable might be defined inside the block, so check all vin and
         return true only if all variables are properly defined *)
      let vua_between = List.map variables_used code in
      let undef_vars = List.fold_left
          (fun acc (codevars, (vin, _vout)) ->
             (Utility.subtraction codevars vin) @ acc)
          []
          (List.combine vua_between internal_var.internal_between)
      in
      match outvar, undef_vars with
        None, [] -> None
      | None, undef_vars -> Some undef_vars
      | Some outvar, [] -> Some [outvar]
      | Some outvar, undef_vars -> Some (outvar :: undef_vars)
  in
  Cfg.NodeSet.fold (fun node acc ->
      match acc, (aux node dvcfg) with
        None, None -> None
      | None, Some x -> Some x
      | Some acc, None -> Some acc
      | Some acc, Some x -> Some (acc @ x)
    ) dvcfg.t.nodes None


let compute_cfg (dvcfg: DVCfg.t) : RISCCfg.t =
  (* no change to the cfg so returned as is *)
  DVCfg.to_cfg dvcfg