open Types
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 (dec, params, loc) ->
        FunUse (dec, flatten_blocks (List.map expand_dims params), 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 (VarDec (Array (_, dims), _, _, _), None, _)) as node ->
        let rec do_expand = function
            | []       -> [node]
            | hd :: tl -> Arg (VarUse (hd, None, [])) :: (do_expand tl)
        in
        Block (do_expand dims)

    | node -> transform_children expand_dims node

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