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, _)
    | Store (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_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

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

  helper DVCeltSet.empty instruction |> 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

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

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)) ->
         (Utility.unique_union
            (variables_used code)
            (Utility.subtraction o (variables_defined code)), o))
      (Utility.combine_twice code previnternalvar.internalbetween)
  ) in

  let newinternalin =
    match newinternalbetween with
    | [] -> previnternalvar.internalout
    | (i, _)::_ -> i
  in

  { previnternalvar with
    internalbetween = newinternalbetween;
    internalin = newinternalin; }

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

  let newinternalout = (
    if Option.equal (=) (Some node) t.t.terminal then (
      match t.t.inputOutputVar with
        Some (_, o) -> [o]
      | None -> []
    ) else (
      let nextnodes = Cfg.NodeMap.find_opt node t.t.edges in
      match nextnodes with
      | None -> []
      | Some (node, None) ->
        (Cfg.NodeMap.find node t.internalvar).internalin
      | Some (node1, Some node2) ->
        Utility.unique_union
          (Cfg.NodeMap.find node1 t.internalvar).internalin
          (Cfg.NodeMap.find node2 t.internalvar).internalin
    )
  ) in

  let newinternalbetween = (
    match List.rev previnternalvar.internalbetween with
    | [] -> []
    | (i, _o) :: btwrest ->
      let btwrest = List.rev btwrest in
      let newbtwrest = List.map2
          (fun (i, _o) (nexti, _nexto) -> (i, nexti))
          btwrest
          (Utility.drop_first_element_list previnternalvar.internalbetween)
      in
      newbtwrest @ [(i, newinternalout)]
  ) in

  { previnternalvar with
    internalout = newinternalout;
    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_live_variables (cfg: RISCCfg.t) : DVCfg.t =
  DVCfg.from_cfg cfg
  |> DVCfg.fixed_point ~init:init ~update:update



module VariableMap = struct
  include Map.Make(Variable)

  let first_empty next start m l =
    let bindings =
      List.fold_left (
        fun acc x ->
          match find_opt x m with
          | None -> acc
          | Some x -> x :: acc) [] l |> List.sort Variable.compare in

    let rec aux x =
      if List.mem x bindings
      then aux (next x)
      else x
    in
    aux start

  let first_empty_Variable m l =
    let next = fun x -> x |> int_of_string |> (+) 1 |> string_of_int in
    let start = "1" in
    first_empty next start m l

  let get_or_set_mapping m l r =
    match find_opt r m with
    | None -> (
        let newr = first_empty_Variable m l in
        let newm = add r newr m in
        (newm, newr)
      )
    | Some r -> (m, r)
end


(* just rename the registers that dont share live status *)
let optimize_cfg (t: DVCfg.t) : DVCfg.t =
  let replace_code ((vin, vout): Variable.t list * Variable.t list)
      (a: Variable.t VariableMap.t)
      (code: DVCfg.elt)
    : (Variable.t VariableMap.t * DVCfg.elt) =
    match code with
    | Nop -> (
        (a, Nop)
      )
    | BRegOp (brop, r1, r2, r3) -> (
        let (newa, newr1) = VariableMap.get_or_set_mapping a vin r1.index in
        let (newa, newr2) = VariableMap.get_or_set_mapping newa vin r2.index in
        let (newa, newr3) = VariableMap.get_or_set_mapping newa vout r3.index in
        (newa, BRegOp (brop, {index = newr1}, {index = newr2}, {index = newr3}))
      )
    | BImmOp (biop, r1, i, r3) -> (
        let (newa, newr1) = VariableMap.get_or_set_mapping a vin r1.index in
        let (newa, newr3) = VariableMap.get_or_set_mapping newa vout r3.index in
        (newa, BImmOp (biop, {index = newr1}, i, {index = newr3}))
      )
    | URegOp (urop, r1, r3) -> (
        let (newa, newr1) = VariableMap.get_or_set_mapping a vin r1.index in
        let (newa, newr3) = VariableMap.get_or_set_mapping newa vout r3.index in
        (newa, URegOp (urop, {index = newr1}, {index = newr3}))
      )
    | Load (r1, r3) -> (
        let (newa, newr1) = VariableMap.get_or_set_mapping a vin r1.index in
        let (newa, newr3) = VariableMap.get_or_set_mapping newa vout r3.index in
        (newa, Load ({index = newr1}, {index = newr3}))
      )
    | LoadI (i, r3) -> (
        let (newa, newr3) = VariableMap.get_or_set_mapping a vout r3.index in
        (newa, LoadI (i, {index = newr3}))
      )
    | Store (r1, r3) -> (
        let (newa, newr1) = VariableMap.get_or_set_mapping a vin r1.index in
        let (newa, newr3) = VariableMap.get_or_set_mapping newa vout r3.index in
        (newa, Store ({index = newr1}, {index = newr3}))
      )
  in

  let aux (assignments: Variable.t VariableMap.t) (t: DVCfg.t) (node: Cfg.Node.t)
    : (Variable.t VariableMap.t * DVCfg.t) =
    let livevars = Cfg.NodeMap.find node t.internalvar in
    let code =
      match Cfg.NodeMap.find_opt node t.t.content with
      | None -> []
      | Some x -> x
    in
    let newcode, newassignments =
      (List.fold_left2
         (fun (acc, assignments) btw code ->
            let na, nc = replace_code btw assignments code in
            (acc @ [nc], na)
         )
         ([], assignments)
         livevars.internalbetween
         code)
    in

    let newcontent = Cfg.NodeMap.add
        node
        newcode
        t.t.content
    in

    let newt = { t with t = { t.t with content = newcontent } } in
    (newassignments, newt)
  in

  (* ------------------- *)

  (* at least the input variable should be in the mapping *)
  let assignments =
    match t.t.inputOutputVar with
      None -> VariableMap.empty
    | Some (i, _o) -> (
        VariableMap.get_or_set_mapping VariableMap.empty [] i |> fst
      )
  in

  let all_variables = List.fold_left
      (fun acc (_, code) ->
         Utility.unique_union acc (variables_all code))
      []
      (Cfg.NodeMap.to_list t.t.content)
  in

  let mapping =
    (* for each variable we get the union of all in and out that contains it
       then we find a register such that it's not in conflict *)
    List.fold_left (fun assignments v -> (
          (* union of all in and out such that v is in the set *)
          let union : 'a list =
            List.fold_left
              (fun (acc: 'a list) (node, (x: DVCfg.internalnode)) ->
                (* not interested in internalin or internalout since information
                   is mirrored into internalbetween *)
                 List.fold_left2
                   (fun acc (i, o) code ->
                      (* we also consider the out set if we "use" v as a guard *)
                      match List.mem v i,
                            List.mem v o,
                            List.mem v (variables_defined code) with
                      | false, false, false -> acc
                      | true, false, false -> Utility.unique_union i acc
                      | false, false, true
                      | false, true, _ -> Utility.unique_union o acc
                      | true, false, true
                      | true, true, _ -> Utility.unique_union
                                           (Utility.unique_union i o) acc
                  )
                  acc
                  x.internalbetween
                  (Cfg.NodeMap.find_opt node t.t.content |>
                   Option.value ~default:[])
              )
              []
              (Cfg.NodeMap.to_list t.internalvar)
          in
          let assignments, _ =
            VariableMap.get_or_set_mapping assignments union v in
          assignments
        )) assignments all_variables
  in

  let mapping, newt =
    Cfg.NodeSet.fold (* for each node we replace all the variables with the
                        optimized ones *)
      (fun node (assign, t) -> aux assign t node)
      t.t.nodes
      (mapping, t)
  in

  { newt with
    t = { newt.t with
          inputOutputVar =
            match newt.t.inputOutputVar with
              None -> None
            | Some (i, o) -> (
                match VariableMap.find_opt i mapping,
                      VariableMap.find_opt o mapping with
                | None, None -> Some (i, o)
                | Some i, None -> Some (i, o)
                | None, Some o -> Some (i, o)
                | Some i, Some o -> Some (i, o)
              )
        }}


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