open Printf
open Types
open Util

(* Create new constant variables for all assigned array values so that they are
 * only evaluated once *)
let rec consts_to_vars node =
  let rec create_vars new_vars values = function
    | [] -> (new_vars, values)
    | hd :: tl ->
      let (new_vars, value) = match hd with
        | ArrayConst (values, ann) ->
          let (new_vars, values) = create_vars new_vars [] values in
          (new_vars, ArrayConst (values, ann))
        | value ->
          let index = fresh_const "const" in
          (new_vars @ [(index, value)], Var (index, None, annof value))
      in
      create_vars new_vars (values @ [value]) tl
  in
  match node with
  (* Add vardecs for values in arrayconst *)
  | VarDec (ArrayDims _ as ctype, name, Some (ArrayConst (values, vann)), ann) ->
    let (new_vars, values) = create_vars [] [] values in
    let value = ArrayConst (values, vann) in
    let create_vardec (name, value) =
      VarDec (basetypeof node, name, Some value, annof value)
    in
    let new_vardecs = List.map create_vardec new_vars in
    Block (new_vardecs @ [VarDec (ctype, name, Some value, ann)])

  (* Add vardec for scalar value *)
  | VarDec (ArrayDims _ as ctype, name, Some value, ann) as node ->
    let scalar_name = fresh_const "scalar" in
    Block [
      VarDec (basetypeof node, scalar_name, Some value, ann);
      VarDec (ctype, name, Some (Var (scalar_name, None, annof value)), ann);
    ]

  | node -> traverse_unit consts_to_vars node

let make_dims make_dimname values make_dec =
  let names = mapi make_dimname values in

  let decs = List.map2 make_dec values names in

  let make_dim value name = Dim (name, annof value) in
  let dims = List.map2 make_dim values names in

  (decs, dims)

(* Generate new variables for array dimensions, to avoid re-evalutation when
 * array dimensions are used (e.g., after array dimension reduction). *)
let rec array_dims node =
  match node with
  | VarDec (ArrayDims (ctype, values), name, init, ann) ->
    (* Names for VarDec dimensions must be unique to avoid weid errors when
     * during context analysis, when an array variable is redeclared within the
     * same scope *)
    let make_dimname i _ = fresh_const (name ^ "_" ^ string_of_int (i + 1)) in

    let make_dec value name = VarDec (Int, name, Some value, []) in
    let (decs, dims) = make_dims make_dimname values make_dec in
    Block (decs @ [VarDec (ArrayDims (ctype, dims), name, init, ann)])

  | GlobalDef (export, ArrayDims (ctype, values), name, init, ann) ->
    (* For global decs, the name must be derived from the array base name, but
     * not constant (no trailing __) since the variable must exist for exporting
     * (and not pruned during constant propagation) *)
    let make_dimname i _ = generate_id name (i + 1) in

    let make_dec value name = GlobalDef (export, Int, name, Some value, []) in
    let (decs, dims) = make_dims make_dimname values make_dec in
    Block (decs @ [GlobalDef (export, ArrayDims (ctype, dims), name, init, ann)])

  (* DISABLED, this is also done in extern.ml
  | GlobalDec (ArrayDims (ctype, values), name, ann) ->
    (*
    let rec make_decs = function
      | [] -> []
      | Dim (name, ann) :: tl -> GlobalDec (Int, name, ann) :: (make_decs tl)
      | _ -> raise InvalidNode
    in
    let decs = make_decs values in
    Block (decs @ [GlobalDec (ArrayDims (ctype, dims), name, ann)])
    *)

    let make_dec value name = GlobalDec (Int, name, []) in
    let (decs, dims) = make_dims name values make_dec in
    Block (decs @ [GlobalDec (ArrayDims (ctype, dims), name, ann)])
  *)

  | node -> traverse_unit array_dims node

(* Split variable initialisation into declaration and assignment *)
let rec split_inits = function
  (* Wrap array initialisation in ArrayInit to pass dimensions *)
  | VarDec (ArrayDims (_, dims) as ctype, name, Some value, ann) ->
    Block [
      VarDec (ctype, name, None, ann);
      Assign (name, None, ArrayInit (value, dims), ann);
    ]

  | GlobalDef (export, (ArrayDims (_, dims) as ctype), name, Some value, ann) ->
    Block [
      GlobalDef (export, ctype, name, None, ann);
      Assign (name, None, ArrayInit (value, dims), ann);
    ]

  | VarDec (ctype, name, Some init, ann) ->
    Block [
      VarDec (ctype, name, None, ann);
      Assign (name, None, init, ann);
    ]

  | GlobalDef (export, ctype, name, Some init, ann) ->
    Block [
      GlobalDef (export, ctype, name, None, ann);
      Assign (name, 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, _) -> Var (name, 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 extract_inits lst =
  let rec trav inits = function
    | [] ->
      (List.rev inits, [])
    | (Assign _ as hd) :: tl
    | (Allocate _ as hd) :: tl ->
      trav (hd :: inits) tl
    | hd :: tl ->
      let (inits, tl) = trav inits tl in
      (inits, (hd :: tl))
  in trav [] lst

let rec move_inits = function
  (* Move global initialisations to __init function *)
  | Program (decls, ann) ->
    let decls = List.map move_inits decls in
    begin match extract_inits decls with
    | ([], _) -> Program (decls, ann)
    | (inits, decls) ->
      let body = Block (VarDecs [] :: LocalFuns [] :: inits) in
      let init_func = FunDef (true, Void, "__init", [], body, []) in
      Program (init_func :: decls, ann)
    end

  (* Split local variable initialisations in declaration and assignment *)
  | FunDef (export, ret_type, name, params, Block body, ann) ->
    let rec place_inits inits = function
      | VarDecs lst :: tl ->
        let (inits, decs) = extract_inits lst in
        VarDecs decs :: (place_inits inits tl)
      | LocalFuns _ as hd :: tl ->
        hd :: inits @ tl
      | _ -> raise InvalidNode
    in
    let body = Block (place_inits [] body) in
    FunDef (export, ret_type, name, params, body, ann)

  | node -> traverse_unit move_inits node

let for_to_while node =
  let rec replace_var var replacement node =
    let trav = (replace_var var replacement) in
    match node with
    | Var (name, None, ann) when name = var ->
      Var (replacement, None, ann)
    | For (counter, start, stop, step, body, ann) when counter = var ->
      For (replacement, trav start, trav stop, trav step, trav body, ann)
    | node ->
      traverse_unit trav node
  in
  let rec traverse new_vars = function
    | FunDef (export, ret_type, name, params, body, ann) ->
      let rec place_decs decs = function
        | Block (VarDecs lst :: tl) -> Block (VarDecs (decs @ lst) :: tl)
        | _ -> raise InvalidNode
      in
      let new_vars = ref [] in
      let body = traverse new_vars body in
      let create_vardec name = VarDec (Int, name, None, []) in
      let new_vardecs = List.map create_vardec !new_vars in
      let body = place_decs new_vardecs body in
      FunDef (export, ret_type, name, params, body, ann)

    (* Transform for-loops to while-loops *)
    | For (counter, start, stop, step, body, ann) ->
      let _i = fresh_id counter in
      let _stop = fresh_const "stop" in
      let _step = fresh_const "step" in
      new_vars := !new_vars @ [_i; _stop; _step];

      let vi = Var (_i, None, []) in
      let vstop = Var (_stop, None, annof stop) in
      let vstep = Var (_step, None, annof step) in
      let cond = Cond (
        Binop (Gt, vstep, Const (IntVal 0, []), []),
        Binop (Lt, vi, vstop, []),
        Binop (Gt, vi, vstop, []),
        []
      ) in
      Block [
        Assign (_i, None, start, annof start);
        Assign (_stop, None, stop, annof stop);
        Assign (_step, None, step, annof step);
        traverse new_vars (While (cond, (Block (
          block_body (replace_var counter _i body) @
          [Assign (_i, None, Binop (Add, vi, vstep, []), [])]
        )), ann));
      ]

    (* Transform while-loops to do-while loops in if-statements *)
    (* DISABLED, while-loops are explicitly supported by the assembly phase
    | While (cond, body, ann) ->
      let cond = traverse new_vars cond in
      let body = traverse new_vars body in
      Block [If (cond, Block [DoWhile (cond, body, ann)], ann)]
    *)

    | node -> traverse_unit (traverse new_vars) node
  in
  traverse (ref []) node

let rec sublist n = function
  | [] when n > 0  -> raise (Invalid_argument "n")
  | []             -> []
  | lst when n = 0 -> lst
  | _ :: tl        -> sublist (n - 1) tl

let rec array_init = function
  (* Transform array constant initialisation into separate assign statements
   * for all entries in the constant array *)
  | Assign (name, None, ArrayInit (ArrayConst _ as value, dims), ann) ->
    let ndims = List.length dims in
    let rec make_assigns depth i indices = function
      | [] -> []
      | hd :: tl ->
        let assigns = traverse depth (i :: indices) hd in
        make_assigns depth (i + 1) indices tl @ assigns
    and traverse depth indices = function
      | ArrayConst (values, _) ->
        make_assigns (depth + 1) 0 indices values
      | value when depth = ndims ->
        let indices = List.map (fun i -> Const (IntVal i, [])) indices in
        [Assign (name, Some (List.rev indices), value, ann)]
      | value when depth < ndims ->
        (* Use the for-loops constructed for scalar assignment *)
        let value = ArrayInit (value, dims) in
        let indices = List.map (fun i -> Const (IntVal i, [])) indices in
        [array_init (Assign (name, Some (List.rev indices), value, ann))]
      | node ->
        let msg = sprintf
          "dimension mismatch: expected %d nesting levels, got %d"
          ndims depth
        in
        raise (NodeError (node, msg))
    in
    Block (List.rev (traverse 0 [] value))

  (* Replace no indices with empty indices to have a list below *)
  | Assign (name, None, (ArrayInit _ as value), ann) ->
    array_init (Assign (name, Some [], value, ann))

  | Assign (name, Some indices, ArrayInit (value, dims), ann) ->
    let rec add_loop indices = function
      | [] ->
        array_init (Assign (name, Some indices, value, ann))
      | dim :: rest ->
        let counter = fresh_id "i" in
        let start = Const (IntVal 0, []) in
        let step = Const (IntVal 1, []) in
        let body = Block [add_loop (indices @ [Var (counter, None, [])]) rest] in
        let stop = match dim with
        | Dim (name, ann) -> Var (name, None, ann)
        | _ -> dim
        in
        For (counter, start, stop, step, body, [])
    in
    let dims_left = sublist (List.length indices) dims in
    add_loop indices dims_left

  | node -> traverse_unit array_init node

let phase = function
  | Ast node ->
    (* Generate variable declarations for expressions that must be evaluated
     * once and used multiple times *)
    let node = consts_to_vars (array_dims node) in

    (* Split variable initialisations into declarations and assignments, and
     * move the assignments to the function body *)
    let node = move_inits (add_allocs (split_inits node)) in

    (* Transform ArrayConst assignment to assignments in for-loops, and
     * transform all for-loops to while-loops afterwards *)
    Ast (for_to_while (array_init (node)))

  | _ -> raise (InvalidInput "desugar")