open Ast
open Util

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

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

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

    | VarUse (Type (Deref (name, values, loc), t), (Array (_, dims) as ctype), depth) ->
        let reduced = [expand (List.rev dims) 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) values] in
        VarLet (Assign (name, reduced, dim_reduce value, loc), ctype, depth)

    | node -> transform_children dim_reduce node

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