civicaml/phases/dim_reduce.ml

open Types
open Util

let rec expand_dims = function
    (* Flatten Block nodes returned by transformations below *)
    | FunDef (export, ret_type, name, params, body, ann) ->
        let params = flatten_blocks (List.map expand_dims params) in
        FunDef (export, ret_type, name, params, expand_dims body, ann)

    | FunDec (ret_type, name, params, ann) ->
        let params = flatten_blocks (List.map expand_dims params) in
        FunDec (ret_type, name, params, ann)

    | FunUse (dec, params, ann) ->
        FunUse (dec, flatten_blocks (List.map expand_dims params), ann)

    (* Add additional parameters for array dimensions *)
    | Param (ArrayDims (ctype, dims), name, ann) ->
        let rec do_expand = function
            | [] -> [Param (Array ctype, name, ann)]
            | Dim (name, ann) :: tail ->
                Param (Int, name, ann) :: (do_expand tail)
            | _ -> raise InvalidNode
        in
        Block (do_expand dims)

    (* Add additional function arguments for array dimensions *)
    | Arg (VarUse (VarDec (ArrayDims (ctype, dims), name, None, decann), None, ann)) as node ->
        let rec do_expand = function
            | [] ->
                (* Remove the (now obsolete dimensions fromt the type) *)
                let dec = VarDec (Array ctype, name, None, decann) in
                [VarUse (dec, None, ann)]
            | hd :: tl ->
                (* A VarDec node has been added for each dimension during
                 * desugaring, so we can safely reconstruct it here (we need no
                 * refrence because the type is immutable, yay!) *)
                let dimdec = VarDec (Int, nameof hd, None, annof hd) in
                Arg (VarUse (dimdec, None, [])) :: (do_expand tl)
        in
        Block (do_expand dims)

    (* Simplify array types in declarations *)
    | VarDec (ArrayDims (ctype, _), name, None, ann) ->
        VarDec (Array ctype, name, None, ann)

    | node -> transform_children expand_dims node

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 (dec, Some values, ann) as node ->
        (match typeof dec with
        | ArrayDims (_, dims) ->
            VarUse (dec, Some [expand depth (List.rev dims) values], ann)
        | _ -> node
        )

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

    | node -> transform_children (dim_reduce depth) node

let phase = function
    | Ast node -> Ast (dim_reduce 0 (expand_dims node))
    | _ -> raise (InvalidInput "dimension reduction")