civicaml/phases/dimreduce.ml

open Types
open Util

let flatten_type = function
  | GlobalDef (export, ArrayDims (ctype, _), name, None, ann) ->
    GlobalDef (export, Array ctype, name, None, ann)

  | GlobalDec (ArrayDims (ctype, _), name, ann) ->
    GlobalDec (Array ctype, name, ann)

  | VarDec (ArrayDims (ctype, _), name, None, ann) ->
    VarDec (Array ctype, name, None, ann)

  | Param (ArrayDims (ctype, _), name, ann) ->
    Param (Array ctype, name, ann)

  | _ -> raise InvalidNode

(* Pass array dimensions explicitly to functions *)
let rec expand_dims = function
  (* For extern arrays, also add a new variables for each dimension with a
   * consistent naming scheme so that they can be exported with the same name by
   * another module *)
  | GlobalDec (ArrayDims (_, dims) as ctype, name, ann) ->
    let rec gendims decs = function
      | [] -> decs
      | Dim (name, ann) as dim :: tl ->
        gendims (GlobalDec (Int, name, ann) :: decs) tl
      | _ -> raise InvalidNode
    in
    Block (List.rev (GlobalDec (ctype, name, ann) :: gendims [] dims))

  (* 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 (_, dims) as ctype, name, ann) ->
    let rec do_expand = function
      | [] -> [Param (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 (dec, None, ann)) when is_array dec ->
    let make_dimdec = function
      | Dim (name, ann) -> Param (Int, name, ann)
      | _ -> raise InvalidNode
    in
    let rec do_expand = function
      | [] ->
        (* Remove the (now obsolete dimensions from the type) *)
        [VarUse (flatten_type dec, None, ann)]

      | hd :: tl ->
        (* A declaration has been added for each dimension during earlier
         * phases, so we can safely reconstruct it here *)
        Arg (VarUse (make_dimdec hd, None, [])) :: do_expand tl
    in
    let dims =
      match typeof dec with
      | ArrayDims (_, dims) -> dims
      | _ -> raise InvalidNode
    in
    Block (do_expand dims)

  | node -> traverse_unit expand_dims node

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

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

(* Transform multi-dimensional arrays into one-dimensional arrays in row-major
 * order *)
and dim_reduce depth = function
  (* Simplify array types in declarations *)
  | GlobalDef (export, ArrayDims (ctype, _), name, None, ann) ->
    GlobalDef (export, Array ctype, name, None, ann)

  | GlobalDec (ArrayDims (ctype, _), name, ann) ->
    GlobalDec (Array ctype, name, ann)

  | Param (ArrayDims (ctype, _), name, ann) ->
    Param (Array ctype, name, ann)

  | VarDec (ArrayDims (ctype, _), name, None, ann) ->
    VarDec (Array ctype, name, None, ann)

  (* Allocate in rw-major order with the array size being the cartegian product
   * of all dimensions *)
  | Allocate (dec, dims, ann) -> Allocate (dec, [multiply dims], ann)

  (* Increase nesting depth when going 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) ->
    begin
      match typeof dec with
      | ArrayDims (_, dims) ->
        VarUse (dec, Some [expand depth dims values], ann)
      | _ -> raise InvalidNode
    end

  (* Expand indices when assigning to array index *)
  | VarLet (dec, Some values, value, ann) ->
    begin match typeof dec with
    | ArrayDims (_, dims) ->
      let value = dim_reduce depth value in
      VarLet (dec, Some [expand depth dims values], value, ann)
    | _ -> raise InvalidNode
    end

  | node -> traverse_unit (dim_reduce depth) node

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