open Types
open Util
open Globals

let rec expand_dims = function
  (* Flatten Block nodes returned by transformations below *)
  | FunDef (export, ret_type, name, params, body, ann) ->
    let params = flatten_blocks (List.map expand_dims params) in
    FunDef (export, ret_type, name, params, expand_dims body, ann)

  | FunDec (ret_type, name, params, ann) ->
    let params = flatten_blocks (List.map expand_dims params) in
    FunDec (ret_type, name, params, ann)

  | FunUse (dec, params, ann) ->
    FunUse (dec, flatten_blocks (List.map expand_dims params), ann)

  (* Add additional parameters for array dimensions *)
  | Param (ArrayDims (ctype, dims), name, ann) ->
    let rec do_expand = function
      | [] -> [Param (Array ctype, name, ann)]
      | Dim (name, ann) :: tail ->
        Param (Int, name, ann) :: (do_expand tail)
      | _ -> raise InvalidNode
    in
    Block (do_expand dims)

  (* Add additional function arguments for array dimensions *)
  | Arg (VarUse (VarDec (ArrayDims (ctype, dims), name, None, decann), None, ann)) ->
    let rec do_expand = function
      | [] ->
        (* Remove the (now obsolete dimensions fromt the type) *)
        let dec = VarDec (Array ctype, name, None, decann) in
        [VarUse (dec, None, ann)]
      | hd :: tl ->
        (* A VarDec node has been added for each dimension during
         * desugaring, so we can safely reconstruct it here (we need no
         * refrence because the type is immutable, yay!) *)
        let dimdec = VarDec (Int, nameof hd, None, annof hd) in
        Arg (VarUse (dimdec, None, [])) :: (do_expand tl)
    in
    Block (do_expand dims)

  (* Simplify array types in declarations *)
  | VarDec (ArrayDims (ctype, _), name, None, ann) ->
    VarDec (Array ctype, name, None, ann)

  | node -> transform_children expand_dims node

let rec multiply = function
  | []       -> raise InvalidNode
  | [node]   -> node
  | hd :: tl -> Binop (Mul, hd, multiply tl, [Type Int])

let rec multiply_all = function
  | []       -> raise InvalidNode
  | [node]   -> node
  | hd :: tl -> Binop (Mul, hd, multiply_all tl, [])

let rec expand depth dims =
  let rec do_expand dims = function
    | []       -> raise InvalidNode
    | [node]   -> dim_reduce depth node
    | i :: j :: tl ->
      let parent_width = List.hd dims in
      let mul = Binop (Mul, dim_reduce depth i, parent_width, [Type Int]) in
      do_expand (List.tl dims) (Binop (Add, mul, j, [Type Int]) :: tl)
  in
  let use_dim = function
    | Dim _ as dim -> VarUse (dim, None, [Type Int; Depth depth])
    | node -> node
  in
  do_expand (List.map use_dim (List.tl dims))

and dim_reduce depth = function
  | Allocate (dec, dims, ann) ->
    Allocate (dec, [multiply dims], ann)

  (* Increase nesting depth when goiing into function *)
  | FunDef (export, ret_type, name, params, body, ann) ->
    let trav = dim_reduce (depth + 1) in
    FunDef (export, ret_type, name, List.map trav params, trav body, ann)

  (* Expand indices when dereferencing *)
  | VarUse (dec, Some values, ann) as node ->
    begin match typeof dec with
    | ArrayDims (_, dims) ->
      VarUse (dec, Some [expand depth dims values], ann)
    | _ -> node
    end

  (* Expand indices when assigning to array index *)
  | VarLet (dec, Some values, value, ann) as node ->
    begin match typeof dec with
    | ArrayDims (_, dims) ->
      VarLet (dec, Some [expand depth dims values], value, ann)
    | _ -> node
    end

  | node -> transform_children (dim_reduce depth) node

let phase = function
  | Ast node -> Ast (dim_reduce 0 (expand_dims node))
  | _ -> raise (InvalidInput "dimension reduction")