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_used (instr : DVCfg.elt) : DVCfg.internal list =
  let helper (acc: DVCeltSet.t) (instr: DVCfg.elt) =
    match instr with
    | Nop
    | LoadI (_, _) ->
      acc
    | BRegOp (_, r1, r2, _) ->
      DVCeltSet.add r1.index acc |>
      DVCeltSet.add r2.index
    | BImmOp (_, r1, _, _)
    | URegOp (_, r1, _)
    | Load (r1, _)
    | Store (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 -> acc
    | BRegOp (_, _, _, r3)
    | BImmOp (_, _, _, r3)
    | URegOp (_, _, r3)
    | Load (_, r3)
    | LoadI (_, r3)
    | Store (_, r3) ->
      DVCeltSet.add r3.index acc
  in

  helper DVCeltSet.empty instructions |> DVCeltSet.to_list

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

let _variables_defined_nth (instructions : DVCfg.elt list) (i: int) : DVCfg.internal list =
  variables_defined (List.nth instructions i)

let _variables_defined_last (instructions : DVCfg.elt list) : DVCfg.internal list =
  variables_defined (List.nth instructions ((List.length instructions) - 1))




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


(* piece of code that computes vout for the whole block, not used,
   use lub below *)
let _dumb_lub (t: DVCfg.t) (node: Cfg.Node.t) : DVCfg.internalnode =
  let previnternalvar = Cfg.NodeMap.find node t.internalvar in
  let code = Cfg.NodeMap.find node t.t.content in
  { previnternalvar with
    internalout =
      Utility.unique_union (variables_defined_all code) (previnternalvar.internalin)
  }


(* We consider only the propagation in the middle elements during the lub.
   This incurs in a performance penality, but it is simpler to implement.
   Each node is connected to one previus node.
*)
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
  { previnternalvar with
    internalbetween =
      List.mapi (* we don't NEED the index but i = 0 is easier to write than
                   to check if vinout is None *)
        (fun i (ithcode, vinout, ithcodeprev) ->
           if i = 0 then
             let dvin = previnternalvar.internalin in
             (dvin, Utility.unique_union dvin (variables_defined ithcode))
           else (
             let ithcodeprev = match ithcodeprev with
                 None -> ([], [])
               | Some x -> x
             in
             match vinout with
               None ->
               ([], variables_defined ithcode)
             | Some prevdvbtw ->
               (snd prevdvbtw,
                Utility.unique_union
                  (variables_defined ithcode)
                  (ithcodeprev |> fst)
               ))
        )
        (* ugly code that zips the three lists that we need to compute each vin
           and vout for the middle of the code *)
        (Utility.combine_thrice
           code
           (Utility.pad_opt
              (Utility.prev previnternalvar.internalbetween None) None (List.length code))
           (Utility.pad previnternalvar.internalbetween None (List.length code))
        );
    internalout =
      match previnternalvar.internalbetween with
        [] -> previnternalvar.internalin
      | _ -> (snd (Utility.last_list previnternalvar.internalbetween))
  }

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
    { previnternalvar with
      internalin =
        match t.t.inputOutputVar with
          Some (i, _) -> [i]
        | None -> []
    }
  else
    let prevnodes = Cfg.NodeMap.find node t.t.reverseEdges in
    { previnternalvar with
      internalin =
        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
        | _ ->
          List.fold_left (* intersection of all previous nodes' dvout *)
            (fun acc prevnode ->
               Utility.unique_intersection acc (Cfg.NodeMap.find prevnode t.internalvar).internalout)
            []
            prevnodes
    }


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 =
  DVCfg.from_cfg cfg
  |> DVCfg.fixed_point ~init:init ~update:update



let check_defined_variables (dvcfg: DVCfg.t) : bool =
  let helper node (dvcfg: DVCfg.t) =
    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 = (* is true if we are in the last node and the out variable is
                    not in vout, so its true if the out variable is not
                    defined *)
      match (Option.equal (=) (Some node) dvcfg.t.terminal,
             dvcfg.t.inputOutputVar,
             internalvar.internalout) with
      | (true, Some (_, outvar), vout) ->
        not (List.mem outvar vout)
      | (_, _, _) ->
        false
    in

    if Utility.inclusion vua (internalvar.internalin) then
      not 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 check = List.fold_left
          (fun acc (codevars, (vin, _vout)) ->
             acc && (Utility.inclusion codevars vin))
          true
          (List.combine vuabetween internalvar.internalbetween)
      in
      check && (not outvar)
  in
  Cfg.NodeSet.fold (fun node acc -> acc && (helper node dvcfg)) dvcfg.t.nodes true


let undefined_variables (dvcfg: DVCfg.t) : Variable.t list =
  let helper (node: Cfg.Node.t) (dvcfg: DVCfg.t) =
    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 -> [] | Some outvar -> [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 with None -> undef_vars | Some outvar -> outvar :: undef_vars
  in
  Cfg.NodeSet.fold (fun node acc -> acc @ (helper node dvcfg)) dvcfg.t.nodes []


let compute_cfg (dvcfg: DVCfg.t) : RISCCfg.t =
  DVCfg.to_cfg dvcfg