open Printf
open Types
open Util

(* Generate new variables for array dimensions, to avoid re-evalutation when
 * array dimensions are used (e.g., after array dimension reduction). *)
let move_array_dims node =
  let patch_dims basename values make_decs make_dimname =
    let names = mapi (fun i _ -> make_dimname basename i) values in
    let decs = List.concat (List.map2 make_decs values names) in
    let make_dim value name = Dim (name, annof value) in
    let dims = List.map2 make_dim values names in
    (decs, dims)
  in
  let fresh_dim name _ = fresh_const name in
  let rec trav = function
    | VarDec (ArrayDims (ctype, values), name, init, ann) ->
      let make_decs value name = [VarDec (Int, name, Some value, [])] in
      let decs, dims = patch_dims name values make_decs fresh_dim in
      Block (decs @ [VarDec (ArrayDims (ctype, dims), name, init, ann)])

    (* Omit the trailing "_" for exported variables since they should not be
     * pruned by optimisations *)
    | GlobalDef (export, ArrayDims (ctype, values), name, init, ann) ->
      let make_dimname = if export then generate_array_dim else fresh_dim in
      let make_decs value name = [GlobalDef (export, Int, name, Some value, [])] in
      let decs, dims = patch_dims name values make_decs make_dimname in
      Block (decs @ [GlobalDef (export, ArrayDims (ctype, dims), name, init, ann)])

    | node -> traverse_unit trav node
  in
  trav node

(* Create new constant variables for scalar initialisations on arrays so that
 * they are only evaluated once *)
let rec move_scalar_inits = function
  (* Prevent next match for ArrayConst initialisations *)
  | VarDec (ArrayDims _, _, Some (ArrayConst _), _) as node ->
    node

  (* Add vardec for scalar value *)
  | VarDec (ArrayDims (ctype, dims) as atype, name, Some value, ann) ->
    let scalar_dec = VarDec (ctype, fresh_const "scalar", Some value, ann) in
    let scalar_use = VarUse (scalar_dec, None, annof value) in
    Block [scalar_dec; VarDec (atype, name, Some scalar_use, ann)]

  | node -> traverse_unit move_scalar_inits node

(* Split variable initialisation into declaration and assignment *)
let rec split_inits = function
  | VarDec (ctype, name, Some init, ann) ->
    let dec = VarDec (ctype, name, None, ann) in
    Block [dec; VarLet (dec, None, init, ann)]

  | GlobalDef (export, ctype, name, Some init, ann) ->
    let dec = GlobalDef (export, ctype, name, None, ann) in
    Block [dec; VarLet (dec, None, init, ann)]

  | node -> traverse_unit split_inits node

(* Add __allocate statements after array declarations *)
let rec add_allocs node =
  let create_dimvar = function
    | Dim (name, _) as dim -> VarUse (dim, None, [])
    | _ -> raise InvalidNode
  in
  match node with
  | VarDec (ArrayDims (_, dims), _, _, ann) ->
    Block [node; Allocate (node, List.map create_dimvar dims, ann)]

  | GlobalDef (_, ArrayDims (_, dims), _, _, ann) ->
    Block [node; Allocate (node, List.map create_dimvar dims, ann)]

  | node -> traverse_unit add_allocs node

let dimsof = function
  | GlobalDef (_, ArrayDims (_, dims), _, _, _)
  | VarDec (ArrayDims (_, dims), _, _, _) -> dims
  | _ -> raise InvalidNode

let rec array_init = function
  (* Transform array constant initialisation into separate assign statements
   * for all entries in the array literal *)
  | VarLet (dec, None, (ArrayConst _ as value), ann) ->
    let intconst i = Const (IntVal (Int32.of_int i), []) in
    let ndims = List.length (dimsof dec) in
    let rec make_assigns depth i indices = function
      | [] -> []
      | hd :: tl ->
        let assigns = trav depth (i :: indices) hd in
        make_assigns depth (i + 1) indices tl @ assigns
    and trav depth indices = function
      | ArrayConst (values, _) ->
        make_assigns (depth + 1) 0 indices values
      | value when depth = ndims ->
        let indices = List.map intconst indices in
        [VarLet (dec, Some (List.rev indices), value, ann)]
      | node ->
        raise (FatalError (NodeMsg (node, sprintf
          "dimension mismatch: expected %d nesting levels, got %d"
          ndims depth)))
    in
    Block (List.rev (trav 0 [] value))

  (* Scalar initialisation *)
  | VarLet (dec, None, scalar, ann) when is_array dec ->
    let rec nest_loops indices = function
      | [] -> Block [VarLet (dec, Some (List.rev indices), scalar, [])]
      | dim :: tl ->
        let counter = fresh_id "i" in
        let start = Const (IntVal 0l, []) in
        let stop = VarUse (dim, None, ann) in
        let step = Const (IntVal 1l, []) in
        let body = nest_loops (Var (counter, None, []) :: indices) tl in
        For (counter, start, stop, step, body, [])
    in
    nest_loops [] (dimsof dec)

  | node -> traverse_unit array_init node

let rec for_to_while = function
  (* Transform for-loops to while-loops *)
  | For (counter, start, stop, step, body, ann) ->
    let dec name init = VarDec (Int, name, Some init, []) in
    let _i = dec counter start in
    let _stop = dec (fresh_const "stop") stop in
    let _step = dec (fresh_const "step") step in

    let vi = VarUse (_i, None, []) in
    let vstop = VarUse (_stop, None, annof stop) in
    let vstep = VarUse (_step, None, annof step) in
    let cond =
      Cond (
        Binop (Gt, vstep, Const (IntVal 0l, []), []),
        Binop (Lt, vi, vstop, []),
        Binop (Gt, vi, vstop, []),
      [])
    in
    Block [
      _i; _stop; _step;
      While (cond, (Block (
        [body; VarLet (_i, None, Binop (Add, vi, vstep, []), [])]
      )), ann) |> for_to_while;
    ]

  (* Transform while-loops to do-while loops in if-statements *)
  (* DISABLED, while-loops are explicitly supported by the assembly phase
  | While (cond, body, ann) ->
    If (cond, DoWhile (cond, for_to_while body, ann), ann)
  *)

  | node -> traverse_unit for_to_while node

let rec move_vardecs = function
  | FunDef (export, ret_type, name, params, body, ann) ->
    let rec trav = function
      | FunDef _ as node -> (move_vardecs node, [])
      | VarDec _ as node -> (DummyNode, [node])
      | node -> traverse_list trav node
    in
    let body, decs = traverse_list trav body in
    let body = Block (decs @ (block_body body)) in
    FunDef (export, ret_type, name, params, body, ann)

  | node -> traverse_unit move_vardecs node

let rec move_inits = function
  | Block (VarDecs decs :: (LocalFuns _ as funs) :: body)->
    let rec extract_stats decs stats = function
      | VarDec _ as dec :: tl -> extract_stats (dec :: decs) stats tl
      | stat :: tl -> extract_stats decs (stat :: stats) tl
      | [] -> List.rev decs, List.rev stats
    in
    let decs, stats = extract_stats [] [] decs in
    Block (VarDecs decs :: funs :: stats @ body)
  | node -> traverse_unit move_inits node

let rec move_global_inits = function
  (* Move global initialisations to __init function *)
  | Program (decls, ann) ->
    let decls = List.map move_global_inits decls in
    let rec extract_inits inits = function
      | [] ->
        (List.rev inits, [])
      | (VarLet _ as hd) :: tl
      | (Allocate _ as hd) :: tl ->
        extract_inits (hd :: inits) tl
      | hd :: tl ->
        let inits, tl = extract_inits inits tl in
        (inits, (hd :: tl))
    in
    begin match extract_inits [] decls with
    | ([], _) -> Program (decls, ann)
    | (inits, decls) ->
      let init_func = FunDef (true, Void, "__init", [], Block inits, []) in
      Program (init_func :: decls, ann)
    end

  | node -> traverse_unit move_global_inits node

let rec group_vardecs = function
  | FunDef (export, ret_type, name, params, Block body, ann) ->
    let rec create = function
      | (VarDec _ as hd) :: tl -> VarDecs [hd] :: create tl
      | tl -> tl
    in
    let rec merge = function
      | VarDecs [a] :: VarDecs b :: tl -> merge (VarDecs (a :: b) :: tl)
      | VarDecs a :: VarDecs b :: tl -> merge (VarDecs (a @ b) :: tl)
      | tl -> tl
    in
    let body = Block (create body |> merge) |> group_vardecs in
    FunDef (export, ret_type, name, params, body, ann)

  | node -> traverse_unit group_vardecs node

let phase = function
  | Ast node ->
    Ast begin
      (* Move array dimensions and scalar initialisations into new variables as
       * initialisations, so that they are evaluated exactly once, and so that
       * dimension names are consistent with the array name *)
      move_array_dims node |> move_scalar_inits |>

      (* Split variable initialisations into declarations and assignments *)
      split_inits |> add_allocs |>

      (* Transform ArrayConst assignment to assignments into for-loops *)
      array_init |>

      (* Transform for-loops to while-loops *)
      for_to_while |> split_inits |>

      (* Move initialization assignments to the function body *)
      move_inits |>

      (* Move variable declarations to the beginning of the function *)
      move_vardecs |>

      (* Move global initialisation assignments to __init *)
      move_global_inits |>

      (* Create and merge VarDecs nodes at the start of each function *)
      group_vardecs |>

      (* Propagate new declaration properties to uses (since we have no
       * pointers) *)
      Context.analyse false
    end

  | _ -> raise InvalidInput