civicaml/phases/expand_dims.ml

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 (ArrayDec (ctype, dims), name, loc) ->
        let rec do_expand = function
            | [] ->
                [Param (ArraySpec (ctype, list_size dims), 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, ArrayDec (ctype, dims), depth)) ->
        let rec do_expand = function
            | [] ->
                let spec = ArraySpec (ctype, list_size dims) in
                [Arg (VarUse (var, spec, depth))]
            | Dim (name, _) :: tl ->
                Arg (VarUse (Var (name, noloc), 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, args) -> Ast (expand_dims node, args)
    | _ -> raise (InvalidInput "expand dimensions")