open Ast
open Util

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

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

    | FunUse (funcall, fundef, depth) ->
        FunUse (expand_dims funcall, expand_dims fundef, depth)

    | FunCall (name, args, loc) ->
        FunCall (name, flatten_blocks (List.map expand_dims args), loc)

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

    (* Add additional function arguments for array dimensions *)
    | Arg (VarUse (var, (Array (_, dims) as ctype), depth)) ->
        let rec do_expand = function
            | []       -> [Arg (VarUse (var, ctype, depth))]
            | hd :: tl -> Arg (VarUse (hd, Int, depth)) :: (do_expand tl)
            | _        -> raise InvalidNode
        in
        Block (do_expand dims)

    | node -> transform_children expand_dims node

let rec phase input =
    prerr_endline "- Expand array dimensions";
    match input with
    | Ast node -> Ast (expand_dims node)
    | _ -> raise (InvalidInput "expand dimensions")