Fixed bug where only VarDec parameters would have their dimensions expanded in function calls

phases/dimreduce.ml

@@ -1,6 +1,18 @@
 open Types
 open Util
 
+let flatten_type = function
+  | GlobalDef (export, ArrayDims (ctype, _), name, None, ann) ->
+    GlobalDef (export, Array ctype, name, None, 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
   (* Flatten Block nodes returned by transformations below *)
@@ -12,7 +24,7 @@ let rec expand_dims = function
     let params = flatten_blocks (List.map expand_dims params) in
     FunDec (ret_type, name, params, ann)
 
-  | FunUse (dec, params, ann) ->
+  | FunUse (dec, params, ann) as node ->
     FunUse (dec, flatten_blocks (List.map expand_dims params), ann)
 
   (* Add additional parameters for array dimensions *)
@@ -26,18 +38,24 @@ let rec expand_dims = function
     Block (do_expand dims)
 
   (* Add additional function arguments for array dimensions *)
-  | Arg (VarUse (VarDec (ArrayDims (ctype, dims), name, None, decann), None, ann)) ->
+  | 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) *)
-        let dec = VarDec (Array ctype, name, None, decann) in
-        [VarUse (dec, None, ann)]
+        [VarUse (flatten_type 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)
+        (* 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)
 
@@ -90,7 +108,7 @@ and dim_reduce depth = function
     FunDef (export, ret_type, name, List.map trav params, trav body, ann)
 
   (* Expand indices when dereferencing *)
-  | VarUse (dec, Some values, ann) as node ->
+  | VarUse (dec, Some values, ann) ->
     begin match typeof dec with
     | ArrayDims (_, dims) ->
       VarUse (dec, Some [expand depth dims values], ann)
@@ -98,7 +116,7 @@ and dim_reduce depth = function
     end
 
   (* Expand indices when assigning to array index *)
-  | VarLet (dec, Some values, value, ann) as node ->
+  | VarLet (dec, Some values, value, ann) ->
     begin match typeof dec with
     | ArrayDims (_, dims) ->
       let value = dim_reduce depth value in