open Types
open Util

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

let use_dim depth = function
    | Dim _ as dim -> VarUse (dim, None, [Type Int; Depth depth])
    (*| VarUse (dim, None, ann) -> VarUse ()*)
    | node -> node

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

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

    (* Increase nesting depth when goiing into function *)
    | FunDef (export, ret_type, name, params, body, ann) ->
        let trav = dim_reduce (depth + 1) in
        FunDef (export, ret_type, name, List.map trav params, trav body, ann)

    (* Expand indices when dereferencing *)
    | VarUse (VarDec (Array (_, dims), _, _, _) as dec, Some values, ann) ->
        VarUse (dec, Some [expand depth (List.rev dims) values], ann)

    (* Expand indices when assigning to array index *)
    | VarLet (VarDec (Array (_, dims), _, _, _) as dec, Some values, value, ann) ->
        VarLet (dec, Some [expand depth (List.rev dims) values], value, ann)

    | node -> transform_children (dim_reduce depth) node

let rec simplify_decs = function
    | VarDec (Array (ctype, dims), name, init, ann) ->
        VarDec (FlatArray ctype, name, init, ann)

    | Param (Array (ctype, dims), name, ann) ->
        Param (FlatArray ctype, name, ann)

    | node -> transform_children simplify_decs node

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