open Ast
open Util

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

let rec expand dims depth = function
    | []       -> raise InvalidNode
    | [node]   -> dim_reduce depth node
    | hd :: tl -> let trav = dim_reduce depth in
                  let mul = Binop (Mul, trav hd, trav (List.hd dims), noloc) in
                  Binop (Mul, mul, expand (List.tl dims) depth tl, noloc)

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

    | FunDef (export, ret_type, name, params, body, loc) ->
        let trav = dim_reduce (depth + 1) in
        FunDef (export, ret_type, name, List.map trav params, trav body, loc)

    | Dim (name, loc) ->
        VarUse (Var (name, loc), Int, depth)

    | VarUse (Type (Deref (name, values, loc), t), (Array (_, dims) as ctype), depth) ->
        let reduced = [expand (List.rev dims) depth values] in
        VarUse (Type (Deref (name, reduced, loc), t), ctype, depth)

    | VarLet (Assign (name, Some values, value, loc), (Array (_, dims) as ctype), depth) ->
        let reduced = Some [expand (List.rev dims) depth values] in
        VarLet (Assign (name, reduced, dim_reduce depth value, loc), ctype, depth)

    | node -> transform_children (dim_reduce depth) node

let rec phase input =
    prerr_endline "- Array dimension reduction";
    match input with
    | Ast node -> Ast (dim_reduce 0 node)
    | _ -> raise (InvalidInput "dimension reduction")