civicaml/phases/dimreduce.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)) ->
        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 rec multiply_all = function
    | []       -> raise InvalidNode
    | [node]   -> node
    | hd :: tl -> Binop (Mul, hd, multiply_all tl, [])

let rec expand depth dims =
    let rec do_expand dims = function
        | []       -> raise InvalidNode
        | [node]   -> dim_reduce depth node
        | i :: j :: tl ->
            let parent_width = List.hd dims in
            let mul = Binop (Mul, dim_reduce depth i, parent_width, [Type Int]) in
            do_expand (List.tl dims) (Binop (Add, mul, j, [Type Int]) :: tl)
    in
    let use_dim = function
        | Dim _ as dim -> VarUse (dim, None, [Type Int; Depth depth])
        | node -> node
    in
    do_expand (List.map use_dim (List.tl dims))

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 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 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")