open Ast
open Util

let rec expand_dims = function
    (* Flatten Block nodes returned by transformations below*)
    | FunDef (export, ret_type, name, params, body, loc) as node ->
        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)

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

    (* Add additional parameters for array dimensions *)
    | Param (ArrayDec (_, dims), name, _) as node ->
        let rec do_expand = function
            | [] -> [node]
            | 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 (_, ArrayDec (_, dims), _)) as node ->
        let rec do_expand = function
            | [] -> [node]
            | Dim (name, _) :: tail ->
                Var (name, noloc) :: (do_expand tail)
            | _ -> raise InvalidNode
        in
        Block (do_expand dims)

    | node -> transform_children expand_dims node

(*
let rec array_init = function
    (* transform scalar assignment into nested for loops *)
    | Assign (name, ArrayScalar (value)) ->
        let rec add_loop indices = function
            | [] ->
                Assign (Deref (name, indices), value)
            | dim :: rest ->
                let counter = fresh_var "counter" in
                let ind = (indices @ [Var counter]) in
                For (counter, IntConst 0, dim, IntConst 1, add_loop ind rest)
        in
        add_loop [] dims

    | Assign (name, ArrayConst (dims)) -> Block []

    | node -> transform array_init node
*)

let rec phase input =
    prerr_endline "- Array dimension reduction";
    match input with
    | Ast (node, args) ->
        Ast (expand_dims node, args)
    | _ -> raise (InvalidInput "desugar")