open Analysis

let globalCounterLabel = ref 0

let nextLabel () : string =
  globalCounterLabel := !globalCounterLabel + 1;
  "l" ^ (string_of_int !globalCounterLabel)

module RISCAssembly = struct
  type register = {
    index : string
  }

  type label = string

  type risci =
    | Nop
    | BRegOp of brop * register * register * register
    | BImmOp of biop * register * int * register
    | URegOp of urop * register * register
    | Load of register * register
    | LoadI of register * int
    | Store of register * register
    | Jump of label
    | CJump of register * label * label
    | Label of label
  and brop =
    | Add
    | Sub
    | Mult
    | Div
    | Mod
    | Pow
    | And
    | Or
    | Eq
    | Less
    | LessEq
    | More
    | MoreEq
  and biop =
    | AddI
    | SubI
    | MultI
    | DivI
    | ModI
    | PowI
    | AndI
    | OrI
    | EqI
    | LessI
    | LessEqI
    | MoreI
    | MoreEqI
  and urop =
    | Not
    | Copy
    | Rand

  type t = {
    code : risci list;
    inputval: int option
  }

  let pp_risci (ppf: out_channel) (v: risci) : unit =
    let rec pp_risci (ppf: out_channel) (v: risci) : unit =
      match v with
        Nop -> Printf.fprintf ppf "\tNop\n"
      | BRegOp (b, r1, r2, r3) -> Printf.fprintf ppf "\t%a r%s r%s => r%s\n" pp_brop b r1.index r2.index r3.index
      | BImmOp (b, r1, i, r3)  -> Printf.fprintf ppf "\t%a r%s %d => r%s\n" pp_biop b r1.index i r3.index
      | URegOp (u, r1, r2) -> Printf.fprintf ppf "\t%a r%s => r%s\n" pp_urop u r1.index r2.index
      | Load (r1, r2) -> Printf.fprintf ppf "\tLoad r%s => r%s\n" r1.index r2.index
      | LoadI (r2, i) -> Printf.fprintf ppf "\tLoadI %d => r%s\n" i r2.index
      | Store (r1, r2) -> Printf.fprintf ppf "\tStore r%s => r%s\n" r1.index r2.index
      | Jump (label) -> Printf.fprintf ppf "\tJump %s\n" label
      | CJump (r, l1, l2) -> Printf.fprintf ppf "\tCJump r%s => %s, %s\n" r.index l1 l2
      | Label (label) -> Printf.fprintf ppf "%s:" label
    and pp_brop (ppf: out_channel) (v: brop) : unit =
      match v with
        Add    -> Printf.fprintf ppf "Add"
      | Sub    -> Printf.fprintf ppf "Sub"
      | Mult   -> Printf.fprintf ppf "Mult"
      | Div    -> Printf.fprintf ppf "Div"
      | Mod    -> Printf.fprintf ppf "Mod"
      | Pow    -> Printf.fprintf ppf "Pow"
      | And    -> Printf.fprintf ppf "And"
      | Or     -> Printf.fprintf ppf "Or"
      | Eq     -> Printf.fprintf ppf "Eq"
      | Less   -> Printf.fprintf ppf "Less"
      | LessEq -> Printf.fprintf ppf "LessEq"
      | More   -> Printf.fprintf ppf "More"
      | MoreEq -> Printf.fprintf ppf "MoreEq"
    and pp_biop (ppf: out_channel) (v: biop) : unit =
      match v with
        AddI    -> Printf.fprintf ppf "AddI"
      | SubI    -> Printf.fprintf ppf "SubI"
      | MultI   -> Printf.fprintf ppf "MultI"
      | DivI    -> Printf.fprintf ppf "DivI"
      | ModI    -> Printf.fprintf ppf "ModI"
      | PowI    -> Printf.fprintf ppf "PowI"
      | AndI    -> Printf.fprintf ppf "AndI"
      | OrI     -> Printf.fprintf ppf "OrI"
      | EqI     -> Printf.fprintf ppf "EqI"
      | LessI   -> Printf.fprintf ppf "LessI"
      | LessEqI -> Printf.fprintf ppf "LessEqI"
      | MoreI   -> Printf.fprintf ppf "MoreI"
      | MoreEqI -> Printf.fprintf ppf "MoreEqI"
    and pp_urop (ppf: out_channel) (v: urop) : unit =
      match v with
        Not     -> Printf.fprintf ppf "Nop"
      | Copy    -> Printf.fprintf ppf "Copy"
      | Rand    -> Printf.fprintf ppf "Rand"
    in
    pp_risci ppf v

  let pp (ppf: out_channel) (t: t) : unit =
    Printf.fprintf ppf "Input Val: ";
    match t.inputval with
      None -> Printf.fprintf ppf "None\n"
    | Some i -> Printf.fprintf ppf "Some %d\n" i;
    Printf.fprintf ppf "Code:\n";
    List.iter (pp_risci ppf) t.code
end

let convert_cfgrisc_risci (i: CfgRISC.RISCSimpleStatements.t list) : (RISCAssembly.risci list) =
  let rec helper (i: CfgRISC.RISCSimpleStatements.t) : RISCAssembly.risci =
    match i with
    | Nop                        -> Nop
    | BRegOp (brop, r1, r2, r3)  -> BRegOp (helper_brop brop,
                                            {index = r1.index},
                                            {index = r2.index},
                                            {index = r3.index})
    | BImmOp (biop, r1, imm, r3) -> BImmOp (helper_biop biop,
                                            {index = r1.index},
                                            imm,
                                            {index = r3.index})
    | URegOp (urop, r1, r3)      -> URegOp (helper_urop urop,
                                            {index = r1.index},
                                            {index = r3.index})
    | Load   (r1, r3)            -> Load   ({index = r1.index},
                                            {index = r3.index})
    | LoadI  (r3, imm)           -> LoadI  ({index = r3.index},
                                            imm)
    | Store  (r1, r3)            -> Store  ({index = r1.index},
                                            {index = r3.index})
  and helper_brop (brop: CfgRISC.RISCSimpleStatements.brop) : RISCAssembly.brop =
    match brop with
    | Add     -> Add
    | Sub     -> Sub
    | Mult    -> Mult
    | Div     -> Div
    | Mod     -> Mod
    | Pow     -> Pow
    | And     -> And
    | Or      -> Or
    | Eq      -> Eq
    | Less    -> Less
    | LessEq  -> LessEq
    | More    -> More
    | MoreEq  -> MoreEq
  and helper_biop (biop: CfgRISC.RISCSimpleStatements.biop) : RISCAssembly.biop =
    match biop with
    | AddI    -> AddI
    | SubI    -> SubI
    | MultI   -> MultI
    | DivI    -> DivI
    | ModI    -> ModI
    | PowI    -> PowI
    | AndI    -> AndI
    | OrI     -> OrI
    | EqI     -> EqI
    | LessI   -> LessI
    | LessEqI -> LessEqI
    | MoreI   -> MoreI
    | MoreEqI -> MoreEqI
  and helper_urop (urop: CfgRISC.RISCSimpleStatements.urop) : RISCAssembly.urop =
    match urop with
    | Not     -> Not
    | Copy    -> Copy
    | Rand    -> Rand
  in
  List.map helper i

let nextCommonSuccessor (prg: CfgRISC.RISCCfg.t) (node1: Cfg.Node.t) (node2: Cfg.Node.t) : Cfg.Node.t option =
  (* Assume the two input nodes are the two branches of an if then else
     statement, then create the two lists that represent the runs until the
     terminal node by choosing always the false statement in guard statements
     (if the guard is for a while statement it gets ignored, if it is for an
     if then else it chooses one of the branches) then find the first common
     node in the lists
  *)
  let rec walk (node: Cfg.Node.t) : Cfg.Node.t list =
    node :: match Cfg.NodeMap.find_opt node prg.edges with
    | None -> []
    | Some (edge, None) -> (walk edge)
    | Some (_, Some edge) -> (walk edge)
  in

  let list1 = walk node1 in
  let list2 = walk node2 in
  let common = List.filter (fun x -> List.mem x list2) list1 in
  match common with
    [] -> None
  | a::_ -> Some a


let rec helper
    (prg: CfgRISC.RISCCfg.t)
    (currentnode: Cfg.Node.t)
    (alreadyVisited: Cfg.Node.t list)
  : (RISCAssembly.risci list) * (Cfg.Node.t list)  =
  (* takes the program, the current node and a list of already visited nodes to
     compute the linearized three address instructions and the list of
     previoulsy visited nodes plus the newly visited nodes. Stops as soon if
     node has already been visited or if no nodes are next *)
  if List.mem currentnode alreadyVisited then
    ([], alreadyVisited)
  else (
    let nextnodes = (Cfg.NodeMap.find_opt currentnode prg.edges) in
    let currentcode =
      (match (Cfg.NodeMap.find_opt currentnode prg.content) with
      | None -> []
      | Some x -> convert_cfgrisc_risci x)
    in
    match nextnodes with
    | Some (nextnode1, None) ->
      let res, vis = (helper prg nextnode1) (currentnode :: alreadyVisited) in
      (currentcode @ res, vis)
    | Some (nextnode1, Some nextnode2) -> (
      let ncs = nextCommonSuccessor prg nextnode1 nextnode2 in
      match ncs with
      | None -> (* should never happen since the terminal node should always be
                   rechable *)
        failwith "Topology got a little mixed up"
      | Some ncs -> (
          if (ncs.id = nextnode2.id)
          then (* while or for loop, three labels are necessary *)
            let label1 = nextLabel () in
            let label2 = nextLabel () in
            let label3 = nextLabel () in

            let res1, _ = (helper prg nextnode1 (currentnode :: nextnode2 :: alreadyVisited)) in
            let res2, vis2 = (helper prg nextnode2 (currentnode :: nextnode1 :: alreadyVisited)) in

            match List.nth currentcode ((List.length currentcode) - 1) with
            | BRegOp (_, _, _, r)
            | BImmOp (_, _, _, r)
            | URegOp (_, _, r)
            | Load (_, r)
            | LoadI (r, _) -> (([Label label1] : RISCAssembly.risci list) @
                                currentcode @
                               ([CJump (r, label2, label3); Label label2] : RISCAssembly.risci list) @
                               res1 @
                               ([Jump label1; Label label3] : RISCAssembly.risci list) @
                               res2
                              , vis2)
            | _ -> failwith "Missing instruction at branch"
          else (* if branches, three labels are necessary *)
            let label1 = nextLabel () in
            let label2 = nextLabel () in
            let label3 = nextLabel () in

            let res1, vis1 = (helper prg nextnode1 (currentnode :: ncs :: alreadyVisited)) in
            let res2, _ = (helper prg nextnode2 vis1) in
            let res3, vis3 = (helper prg ncs (currentnode :: alreadyVisited)) in
            match List.nth currentcode ((List.length currentcode) - 1) with
            | BRegOp (_, _, _, r)
            | BImmOp (_, _, _, r)
            | URegOp (_, _, r)
            | Load (_, r)
            | LoadI (r, _) -> (currentcode @
                               ([CJump (r, label1, label2); Label label1] : RISCAssembly.risci list) @
                               res1 @
                               ([Jump label3; Label label2] : RISCAssembly.risci list) @
                               res2 @
                               ([Label label3] : RISCAssembly.risci list) @
                               res3
                               , vis3)
            | _ -> failwith "Missing instruction at branch"
        )
    )
    | None -> (currentcode, currentnode :: alreadyVisited)
  )

let convert (prg: CfgRISC.RISCCfg.t) : RISCAssembly.t =
  {code = (helper prg (Option.get prg.initial) [] |> fst |>
           List.append ([Label "main"] : RISCAssembly.risci list));
   inputval = prg.inputVal}