open Analysis

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

  let pplist (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 helper (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

  helper 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 helper (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

  helper 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 helper (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

  helper DVCeltSet.empty instructions |> DVCeltSet.to_list


(* init function, assign the bottom to everything *)
let _init_bottom : (DVCfg.elt list -> DVCfg.internalnode) =
  (fun l -> {internalin = [];
             internalout = [];
             internalbetween = (List.init (List.length l) (fun _ -> ([], [])))})

(* init function, assign the top to everything *)
let init_top (all_variables) : (DVCfg.elt list -> DVCfg.internalnode) =
  (fun l -> {internalin = all_variables;
             internalout = all_variables;
             internalbetween = (List.init (List.length l)
                                  (fun _ -> (all_variables, all_variables)))})


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

  let newinternalbetween = (
    List.map
      (fun (code, (i, _o)) ->
         (i, Utility.unique_union i (variables_defined code)))
      (List.combine code previnternalvar.internalbetween)
  ) in

  let newinternalout =
    match newinternalbetween with
    | [] -> previnternalvar.internalin
    | _ -> (snd (Utility.last_list newinternalbetween))
  in

  { previnternalvar with
    internalbetween = newinternalbetween;
    internalout = newinternalout }

let lucf (t: DVCfg.t) (node: Cfg.Node.t) : DVCfg.internalnode =
  let previnternalvar = Cfg.NodeMap.find node t.internalvar in

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

    let newinternalbetween = ( (* set the dvin of each to the previous dvout *)
      match previnternalvar.internalbetween with
        [] -> []
      | [(_i, o)] -> [(newinternalin, o)]
      | (_i, o) :: btwrest ->
        (newinternalin, o) :: (
          List.map (fun ((_i, o), (_previ, prevo)) -> (prevo, o))
            (Utility.combine_twice btwrest previnternalvar.internalbetween)
        )
    ) in
    { previnternalvar with
      internalin = newinternalin;
      internalbetween = newinternalbetween }
  else
    (* if L is not initial set dvin to the intersection of the previous node's
       dvouts *)
    let prevnodes = Cfg.NodeMap.find node t.t.reverseEdges in
    let newinternalin = (
      match prevnodes with
      | [] ->
        []
      | [prevnode] ->
        (Cfg.NodeMap.find prevnode t.internalvar).internalout
      | [prevnode1; prevnode2] ->
        Utility.unique_intersection
          (Cfg.NodeMap.find prevnode1 t.internalvar).internalout
          (Cfg.NodeMap.find prevnode2 t.internalvar).internalout
      | prevnode :: restnodes ->
        List.fold_left (* intersection of all previous nodes' dvout *)
          (fun acc prevnode ->
             Utility.unique_intersection
               acc
               (Cfg.NodeMap.find prevnode t.internalvar).internalout)
          (Cfg.NodeMap.find prevnode t.internalvar).internalout
          restnodes
    ) in

    let newinternalbetween =
      match previnternalvar.internalbetween with
        [] -> []
      | [(_i, o)] -> [(newinternalin, o)]
      | (_i, o) :: btwrest ->
        (newinternalin, o) :: (
          List.map (fun ((_i, o), (_previ, prevo)) -> (prevo, o))
            (Utility.combine_twice btwrest previnternalvar.internalbetween)
        )
    in
    { previnternalvar with
      internalin = newinternalin;
      internalbetween = newinternalbetween }


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


let compute_defined_variables (cfg: RISCCfg.t) : DVCfg.t =
  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.inputOutputVar 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 =
  let helper (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 internalvar = Cfg.NodeMap.find node dvcfg.internalvar in
    let vua = variables_used_all code in

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

    if Utility.inclusion vua (internalvar.internalin) 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 vuabetween = List.map variables_used code in
      let undef_vars = List.fold_left
          (fun acc (codevars, (vin, _vout)) ->
             (Utility.subtraction codevars vin) @ acc)
          []
          (List.combine vuabetween internalvar.internalbetween)
      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, (helper 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 =
  DVCfg.to_cfg dvcfg