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 VariableMap = Map.Make(Variable)

let variables_frequency (instr : RISCCfg.elt) : (Variable.t * int) list =
  let add_one = (fun x -> match x with None -> Some 1 | Some x -> Some (x + 1))
  in

  let aux (acc: int VariableMap.t) (instr: RISCCfg.elt) : int VariableMap.t =
    match instr with
    | Nop ->
      acc
    | BRegOp (_, r1, r2, r3) ->
      VariableMap.update r1.index add_one acc |>
      VariableMap.update r2.index add_one |>
      VariableMap.update r3.index add_one
    | BImmOp (_, r1, _, r3)
    | URegOp (_, r1, r3)
    | Load (r1, r3)
    | Store (r1, r3) ->
      VariableMap.update r1.index add_one acc |>
      VariableMap.update r3.index add_one
    | LoadI (_, r3) ->
      VariableMap.update r3.index add_one acc
  in

  aux VariableMap.empty instr |> VariableMap.to_list

(* computes syntactic frequency of all variables in the code *)
let variables_all_frequency (instructions : RISCCfg.elt list)
  : (Variable.t * int) list =
  List.fold_left
    ( fun (acc: int VariableMap.t) (instr: RISCCfg.elt) ->
        VariableMap.union
          (fun _v x y -> Some (x + y))
          acc (variables_frequency instr |> VariableMap.of_list) )
    VariableMap.empty instructions |> VariableMap.to_list


let reduce_registers (n: int) (cfg: RISCCfg.t) : RISCCfg.t =
  (if n < 4 then (
      failwith "ReduceRegisters: number of registers too small"
    ) else ());

  (* we get all the variables with associated frequency (only syntactic use) *)
  let all_variables = List.fold_left
      (fun acc (_, code) ->
         Utility.unique_union_assoc
           (fun _n x y -> x + y) acc (variables_all_frequency code))
      []
      (Cfg.NodeMap.to_list cfg.content)
  in

  (* add one to in and out variables *)
  let all_variables =
    match cfg.input_output_var with
    | None -> all_variables
    | Some (i, _o) -> (
        match List.assoc_opt i all_variables with
        | None -> (i, 1) :: all_variables
        | Some f -> (i, f+1) :: (List.remove_assoc i all_variables)
      )
  in

  let all_variables =
    match cfg.input_output_var with
    | None -> all_variables
    | Some (_i, o) -> (
        match List.assoc_opt o all_variables with
        | None -> (o, 1) :: all_variables
        | Some f -> (o, f+1) :: (List.remove_assoc o all_variables)
      )
  in

  (* replace each operation with a list of operations that have the new
     registers or load from memory *)
  let replace_registers
      (remappedregisters: Variable.t VariableMap.t)
      (memorymap: int VariableMap.t)
      (temporaryregisters: Variable.t list)
      (code: RISCCfg.elt list)
    : RISCCfg.elt list =

    let tmpreg1: CfgRISC.RISCSimpleStatements.register =
      {index = List.nth temporaryregisters 0} in
    let tmpreg2: CfgRISC.RISCSimpleStatements.register =
      {index = List.nth temporaryregisters 1} in

    let aux (instruction: RISCCfg.elt) : RISCCfg.elt list =
      match instruction with
      | Nop -> [Nop]
      | BRegOp (brop, r1, r2, r3) -> (
          match ( VariableMap.find_opt r1.index remappedregisters,
                  VariableMap.find_opt r2.index remappedregisters,
                  VariableMap.find_opt r3.index remappedregisters,
                  VariableMap.find_opt r1.index memorymap,
                  VariableMap.find_opt r2.index memorymap,
                  VariableMap.find_opt r3.index memorymap )
          with
          | Some r1, Some r2, Some r3, _, _, _ ->
            [BRegOp (brop, {index = r1}, {index = r2}, {index = r3})]
          | Some r1, None, Some r3, _, Some m2, _ ->
            [LoadI (m2, tmpreg2);
             Load (tmpreg2, tmpreg2);
             BRegOp (brop, {index = r1}, tmpreg2, {index = r3})]
          | None, Some r2, Some r3, Some m1, _, _ ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             BRegOp (brop, tmpreg1, {index = r2}, {index = r3})]
          | None, None, Some r3, Some m1, Some m2, _ ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             LoadI (m2, tmpreg2);
             Load (tmpreg2, tmpreg2);
             BRegOp (brop, tmpreg1, tmpreg2, {index = r3})]
          | Some r1, Some r2, None, _, _, Some m3 ->
            [BRegOp (brop, {index = r1}, {index = r2}, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | Some r1, None, None, _, Some m2, Some m3 ->
            [LoadI (m2, tmpreg2);
             Load (tmpreg2, tmpreg2);
             BRegOp (brop, {index = r1}, tmpreg2, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | None, Some r2, None, Some m1, _, Some m3 ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             BRegOp (brop, tmpreg1, {index = r2}, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | None, None, None, Some m1, Some m2, Some m3 ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             LoadI (m2, tmpreg2);
             Load (tmpreg2, tmpreg2);
             BRegOp (brop, tmpreg1, tmpreg2, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | _ -> [BRegOp (brop,
                          {index = r1.index},
                          {index = r2.index},
                          {index = r3.index})]
        )
      | BImmOp (biop, r1, i, r3) -> (
          match ( VariableMap.find_opt r1.index remappedregisters,
                  VariableMap.find_opt r3.index remappedregisters,
                  VariableMap.find_opt r1.index memorymap,
                  VariableMap.find_opt r3.index memorymap )
          with
          | Some r1, Some r3, _, _ ->
            [BImmOp (biop, {index = r1}, i, {index = r3})]
          | Some r1, None, _, Some m3 ->
            [BImmOp (biop, {index = r1}, i, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | None, Some r3, Some m1, _ ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             BImmOp (biop, tmpreg1, i, {index = r3})]
          | None, None, Some m1, Some m3 ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             BImmOp (biop, tmpreg1, i, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | _ -> failwith ("ReduceRegisters: fail to partition a set, some" ^
                           " registers have no binding.")
        )
      | URegOp (urop, r1, r3) ->(
          match ( VariableMap.find_opt r1.index remappedregisters,
                  VariableMap.find_opt r3.index remappedregisters,
                  VariableMap.find_opt r1.index memorymap,
                  VariableMap.find_opt r3.index memorymap )
          with
          | Some r1, Some r3, _, _ ->
            [URegOp (urop, {index = r1}, {index = r3})]
          | Some r1, None, _, Some m3 ->
            [URegOp (urop, {index = r1}, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | None, Some r3, Some m1, _ ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             URegOp (urop, tmpreg1, {index = r3})]
          | None, None, Some m1, Some m3 ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             URegOp (urop, tmpreg1, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | _ -> failwith ("ReduceRegisters: fail to partition a set, some" ^
                           " registers have no binding.")
        )
      | Load (r1, r3) -> (
          match ( VariableMap.find_opt r1.index remappedregisters,
                  VariableMap.find_opt r3.index remappedregisters,
                  VariableMap.find_opt r1.index memorymap,
                  VariableMap.find_opt r3.index memorymap )
          with
          | Some r1, Some r3, _, _ ->
            [Load ({index = r1}, {index = r3})]
          | Some r1, None, _, Some m3 ->
            [Load ({index = r1}, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | None, Some r3, Some m1, _ ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             Load (tmpreg1, {index = r3})]
          | None, None, Some m1, Some m3 ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             Load (tmpreg1, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | _ -> failwith ("ReduceRegisters: fail to partition a set, some" ^
                           " registers have no binding.")
        )
      | LoadI (i, r3) -> (
          (* we want to store an integer in memory immediately (strange, but
             unless better heuristic to choose the variables to replace we are
             stuck) *)
          match ( VariableMap.find_opt r3.index remappedregisters,
                  VariableMap.find_opt r3.index memorymap )
          with
          | Some r3, _ ->
            [LoadI (i, {index = r3})]
          | None, Some m3 ->
            [LoadI (i, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | _ -> failwith ("ReduceRegisters: fail to partition a set, some" ^
                           " registers have no binding.")
        )
      | Store (r1, r3) -> (
          (* we want to maybe store an address in memory (very confusing,
             but maybe possible) *)
          match ( VariableMap.find_opt r1.index remappedregisters,
                  VariableMap.find_opt r3.index remappedregisters,
                  VariableMap.find_opt r1.index memorymap,
                  VariableMap.find_opt r3.index memorymap )
          with
          | Some r1, Some r3, _, _ ->
            [Store ({index = r1}, {index = r3})]
          | Some r1, None, _, Some m3 ->
            [Store ({index = r1}, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | None, Some r3, Some m1, _ ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             Store (tmpreg1, {index = r3})]
          | None, None, Some m1, Some m3 ->
            [LoadI (m1, tmpreg1);
             Load (tmpreg1, tmpreg1);
             Store (tmpreg1, tmpreg2);
             LoadI (m3, tmpreg1);
             Store (tmpreg2, tmpreg1)]
          | _ -> failwith ("ReduceRegisters: fail to partition a set, some" ^
                           " registers have no binding.")
        )
    in

    List.map aux code |> List.concat
  in


  (* partition the variables into two sets, most frequent and least frequent
     then apply the transformation *)
  let aux (cfg: RISCCfg.t) (all_variables: (string * int) list) =
    (* we keep the first two variables free for immediate use *)
    let most_frequent, least_frequent =
      List.sort (fun (_a, fa) (_b, fb) -> Int.compare fb fa) all_variables
      |> Utility.take (n-2)
    in
    let most_frequent, _frequencies = List.split most_frequent in
    let least_frequent, _frequencies = List.split least_frequent in

    (* we map the most frequent to new registers, so that the first two are
       always free *)
    let most_frequent_mapping = (* +3 because starts at 0, but we want to start
                                   at 1*)
      List.mapi
        (fun n v -> (v, (string_of_int (n+3): Variable.t)))
        most_frequent
      |> VariableMap.of_list
    in
    (* we map the least to memory *)
    let least_frequent_mapping =
      List.mapi (fun n v -> (v, (n: int))) least_frequent
      |> VariableMap.of_list
    in

    (* we need to replace both at the same time, because we might have mapped
       some registers to already used registers, so a double pass might not
       differentiate the two *)
    (* special care must be taken for the in and out registers *)
    let newcfg = {
      cfg with
      content = Cfg.NodeMap.map
          ( fun x ->
              replace_registers
                most_frequent_mapping
                least_frequent_mapping
                ["1"; "2"]
                x
          ) cfg.content}
    in

    if newcfg.input_output_var = None
    then newcfg (* if no input or output variables we ignore *)
    else
      let i, o = Option.get newcfg.input_output_var in
      match (VariableMap.find_opt i most_frequent_mapping,
             VariableMap.find_opt o most_frequent_mapping,
             VariableMap.find_opt i least_frequent_mapping,
             VariableMap.find_opt o least_frequent_mapping,
             newcfg.initial,
             newcfg.terminal )
      with (*we check if in and out are simply remapped or are put in memory*)
      | Some i, Some o, _, _, _, _ ->
        { newcfg with input_output_var = Some (i, o) }
      | Some i, None, _, Some _, _, None ->
        (* we check for the terminal node, if not present we are very confused
           and dont modify the out variable *)
        { newcfg with input_output_var = Some (i, o)}
      | Some i, None, _, Some mo, _, Some n ->
        (* since the output simbol is in memory we need to first retrive it
           and then put the result in a temporary register *)
        let terminal_content = (
          match Cfg.NodeMap.find_opt n newcfg.content with
          | None -> []
          | Some x -> x
        ) @ [LoadI (mo, {index = "2"});
             Load ({index = "2"}, {index = "2"})]
        in
        let content =
          Cfg.NodeMap.add n terminal_content newcfg.content
        in
        { newcfg with
          input_output_var = Some (i, "2");
          content = content
        }
      | None, Some o, Some _, _, _, None ->
        { newcfg with input_output_var = Some (i, o) }
      | None, Some o, Some mi, _, _, Some n -> (
          (* the input simbol should be stored in memory *)
          let initialcontent =
            [(LoadI (mi, {index = "2"}) : RISCCfg.elt);
             Store ({index = "1"}, {index = "2"})] @ (
              match Cfg.NodeMap.find_opt n newcfg.content with
              | None -> []
              | Some x -> x
            )
          in
          let content = Cfg.NodeMap.add n initialcontent newcfg.content in
          { newcfg with
            input_output_var = Some ("1", o);
            content = content
          }
        )
      | None, None, Some _, Some _, None, None ->
        { newcfg with input_output_var = Some (i, o) }
      | None, None, Some _, Some mo, None, Some n ->
        (* both simbols should be in memory *)
        let terminalcontent = (
          match Cfg.NodeMap.find_opt n newcfg.content with
          | None -> []
          | Some x -> x
        ) @ [LoadI (mo, {index = "2"});
             Load ({index = "2"}, {index = "2"})]
        in
        let content =
          Cfg.NodeMap.add n terminalcontent newcfg.content
        in
        { newcfg with
          input_output_var = Some (i, "2");
          content = content
        }
      | None, None, Some mi, Some _, Some n, None ->
        (* both simbols should be in memory *)
        let initialcontent =
          [(LoadI (mi, {index = "2"}) : RISCCfg.elt);
           Store ({index = "1"}, {index = "2"})] @ (
            match Cfg.NodeMap.find_opt n newcfg.content with
            | None -> []
            | Some x -> x
          )
        in
        let content = Cfg.NodeMap.add n initialcontent newcfg.content in
        { newcfg with
          input_output_var = Some ("1", o);
          content = content
        }
      | None, None, Some mi, Some mo, Some ni, Some no ->
        (* both simbols should be in memory *)
        let initialcontent =
          [(LoadI (mi, {index = "2"}) : RISCCfg.elt);
           Store ({index = "1"}, {index = "2"})] @ (
            match Cfg.NodeMap.find_opt ni newcfg.content with
            | None -> []
            | Some x -> x
          )
        in
        let terminalcontent = (
          match Cfg.NodeMap.find_opt no newcfg.content with
          | None -> []
          | Some x -> x
        ) @ [LoadI (mo, {index = "2"});
             Load ({index = "2"}, {index = "2"})]
        in
        let content =
          Cfg.NodeMap.add ni initialcontent newcfg.content
        in
        let content =
          Cfg.NodeMap.add no terminalcontent content
        in
        { newcfg with
          input_output_var = Some ("1", "2");
          content = content
        }
      | _ -> failwith ("ReduceRegisters: fail to partition a set, some" ^
                       " registers have no binding.")
  in

  ( if List.length all_variables <= n
    then cfg
    else aux cfg all_variables )