compiler: Assume array()s are always rectangular

artiq/compiler/transforms/artiq_ir_generator.py

@@ -1561,9 +1561,8 @@ class ARTIQIRGenerator(algorithm.Visitor):
             self.current_block = after_invoke
 
     def _get_array_offset(self, shape, indices):
-        last_stride = None
         result = indices[0]
-        for dim, index in zip(shape[:-1], indices[1:]):
+        for dim, index in zip(shape[1:], indices[1:]):
             result = self.append(ir.Arith(ast.Mult(loc=None), result, dim))
             result = self.append(ir.Arith(ast.Add(loc=None), result, index))
         return result
@@ -2090,9 +2089,8 @@ class ARTIQIRGenerator(algorithm.Visitor):
                 result_elt = result_type["elt"].find()
                 num_dims = result_type["num_dims"].value
 
-                # Derive shape from first element on each level (currently, type
+                # Derive shape from first element on each level (and fail later if the
-                # inference make sure arrays are always rectangular; in the future, we
+                # array is in fact jagged).
-                # might want to insert a runtime check here).
                 first_elt = None
                 lengths = []
                 for dim_idx in range(num_dims):
@@ -2110,32 +2108,37 @@ class ARTIQIRGenerator(algorithm.Visitor):
                 buffer = self.append(
                     ir.Alloc([num_total_elts], result_type.attributes["buffer"]))
 
-                def body_gen(index):
+                def assign_elems(outer_indices, indexed_arg):
-                    # TODO: This is hilariously inefficient; we really want to emit a
+                    if len(outer_indices) == num_dims:
-                    # nested loop for the source and keep one running index for the
+                        dest_idx = self._get_array_offset(lengths, outer_indices)
-                    # target buffer.
+                        self.append(ir.SetElem(buffer, dest_idx, indexed_arg))
-                    indices = []
+                    else:
-                    mod_idx = index
+                        this_level_len = self.iterable_len(indexed_arg)
-                    for dim_idx in reversed(range(1, num_dims)):
+                        dim_idx = len(outer_indices)
-                        dim_len = self.append(ir.GetAttr(shape, dim_idx))
+                        if dim_idx > 0:
-                        indices.append(
+                            # Check for rectangularity (outermost index is never jagged,
-                            self.append(ir.Arith(ast.Mod(loc=None), mod_idx, dim_len)))
+                            # by definition).
-                        mod_idx = self.append(
+                            result_len = self.append(ir.GetAttr(shape, dim_idx))
-                            ir.Arith(ast.FloorDiv(loc=None), mod_idx, dim_len))
+                            self._make_check(
-                    indices.append(mod_idx)
+                                self.append(ir.Compare(ast.Eq(loc=None), this_level_len, result_len)),
-
+                                lambda a, b: self.alloc_exn(
-                    elt = arg
+                                    builtins.TException("ValueError"),
-                    for idx in reversed(indices):
+                                    ir.Constant(
-                        elt = self.iterable_get(elt, idx)
+                                        "arrays must be rectangular (lengths were {0} vs. {1})",
-                    self.append(ir.SetElem(buffer, index, elt))
+                                        builtins.TStr()), a, b),
-                    return self.append(
+                                params=[this_level_len, result_len],
-                        ir.Arith(ast.Add(loc=None), index,
+                                loc=node.loc)
-                                 ir.Constant(1, self._size_type)))
-
-                self._make_loop(
-                    ir.Constant(0, self._size_type), lambda index: self.append(
-                        ir.Compare(ast.Lt(loc=None), index, num_total_elts)), body_gen)
 
+                        def body_gen(index):
+                            elem = self.iterable_get(indexed_arg, index)
+                            assign_elems(outer_indices + [index], elem)
+                            return self.append(
+                                ir.Arith(ast.Add(loc=None), index,
+                                        ir.Constant(1, self._size_type)))
+                        self._make_loop(
+                            ir.Constant(0, self._size_type), lambda index: self.append(
+                                ir.Compare(ast.Lt(loc=None), index, this_level_len)), body_gen)
+                assign_elems([], arg)
                 return self.append(ir.Alloc([buffer, shape], node.type))
             else:
                 assert False

artiq/compiler/transforms/inferencer.py

@@ -8,30 +8,6 @@ from .. import asttyped, types, builtins
 from .typedtree_printer import TypedtreePrinter
 
 
-def match_rectangular_list(elts):
-    num_elts = None
-    elt_type = None
-    all_child_elts = []
-
-    for e in elts:
-        if elt_type is None:
-            elt_type = e.type.find()
-        if not isinstance(e, asttyped.ListT):
-            return elt_type, 0
-        if num_elts is None:
-            num_elts = len(e.elts)
-        elif num_elts != len(e.elts):
-            return elt_type, 0
-        all_child_elts += e.elts
-
-    if not all_child_elts:
-        # This ultimately turned out to be a list (of list, of ...) of empty lists.
-        return elt_type["elt"], 1
-
-    elt, num_dims = match_rectangular_list(all_child_elts)
-    return elt, num_dims + 1
-
-
 class Inferencer(algorithm.Visitor):
     """
     :class:`Inferencer` infers types by recursively applying the unification
@@ -862,29 +838,42 @@ class Inferencer(algorithm.Visitor):
             if len(node.args) == 1 and len(node.keywords) == 0:
                 arg, = node.args
 
-                if builtins.is_iterable(arg.type):
+                # In the absence of any other information (there currently isn't a way
-                    # KLUDGE: Support multidimensional arary creation if lexically
+                # to specify any), assume that all iterables are expandable into a
-                    # specified as a rectangular array of lists.
+                # (runtime-checked) rectangular array of the innermost element type.
-                    elt, num_dims = match_rectangular_list([arg])
+                elt = arg.type
-                    if num_dims == 0:
+                num_dims = 0
-                        # Not given as a list, so just default to 1 dimension.
+                result_dims = (node.type.find()["num_dims"].value
-                        elt = builtins.get_iterable_elt(arg.type)
+                               if builtins.is_array(node.type) else -1)
-                        num_dims = 1
+                while True:
-                    self._unify(node.type,
+                    if num_dims == result_dims:
-                                builtins.TArray(elt, types.TValue(num_dims)),
+                        # If we already know the number of dimensions of the result,
-                                node.loc, arg.loc)
+                        # stop so we can disambiguate the (innermost) element type of
-                elif types.is_var(arg.type):
+                        # the argument if it is still unknown (e.g. empty array).
-                    pass  # undetermined yet
+                        break
-                else:
+                    if types.is_var(elt):
-                    note = diagnostic.Diagnostic("note",
+                        return  # undetermined yet
-                        "this expression has type {type}",
+                    if not builtins.is_iterable(elt):
-                        {"type": types.TypePrinter().name(arg.type)},
+                        break
-                        arg.loc)
+                    num_dims += 1
-                    diag = diagnostic.Diagnostic("error",
+                    elt = builtins.get_iterable_elt(elt)
+
+                if num_dims == 0:
+                    note = diagnostic.Diagnostic(
+                        "note", "this expression has type {type}",
+                        {"type": types.TypePrinter().name(arg.type)}, arg.loc)
+                    diag = diagnostic.Diagnostic(
+                        "error",
                         "the argument of {builtin}() must be of an iterable type",
                         {"builtin": typ.find().name},
-                        node.func.loc, notes=[note])
+                        node.func.loc,
+                        notes=[note])
                     self.engine.process(diag)
+                    return
+
+                self._unify(node.type,
+                            builtins.TArray(elt, types.TValue(num_dims)),
+                            node.loc, arg.loc)
             else:
                 diagnose(valid_forms())
         elif types.is_builtin(typ, "list"):

artiq/test/lit/embedding/arrays.py

@@ -12,6 +12,7 @@ float_mat = array([[1.0, 2.0], [3.0, 4.0]])
 
 @kernel
 def entrypoint():
+    # TODO: These need to be runtime tests!
     assert int_vec.shape == (3, )
     assert int_vec[0] == 1
     assert int_vec[1] == 2
@@ -22,14 +23,14 @@ def entrypoint():
     assert float_vec[1] == 2.0
     assert float_vec[2] == 3.0
 
-    # assert int_mat.shape == (2, 2)
+    assert int_mat.shape == (2, 2)
-    # assert int_mat[0][0] == 1
+    assert int_mat[0][0] == 1
-    # assert int_mat[0][1] == 2
+    assert int_mat[0][1] == 2
-    # assert int_mat[1][0] == 3
+    assert int_mat[1][0] == 3
-    # assert int_mat[1][1] == 4
+    assert int_mat[1][1] == 4
 
-    # assert float_mat.shape == (2, 2)
+    assert float_mat.shape == (2, 2)
-    # assert float_mat[0][0] == 1.0
+    assert float_mat[0][0] == 1.0
-    # assert float_mat[0][1] == 2.0
+    assert float_mat[0][1] == 2.0
-    # assert float_mat[1][0] == 3.0
+    assert float_mat[1][0] == 3.0
-    # assert float_mat[1][1] == 4.0
+    assert float_mat[1][1] == 4.0

artiq/test/lit/inferencer/array.py

@@ -1,7 +1,9 @@
 # RUN: %python -m artiq.compiler.testbench.inferencer %s >%t
 # RUN: OutputCheck %s --file-to-check=%t
 
-# CHECK-L: numpy.array(elt='a, num_dims=1)
+# Nothing known, as there could be several more dimensions
+# hidden from view by the array being empty.
+# CHECK-L: ([]:list(elt='a)):'b
 array([])
 
 # CHECK-L: numpy.array(elt=numpy.int?, num_dims=1)
@@ -9,5 +11,6 @@ array([1, 2, 3])
 # CHECK-L: numpy.array(elt=numpy.int?, num_dims=2)
 array([[1, 2, 3], [4, 5, 6]])
 
-# CHECK-L: numpy.array(elt=list(elt=numpy.int?), num_dims=1)
+# Jagged arrays produce runtime failure:
+# CHECK-L: numpy.array(elt=numpy.int?, num_dims=2)
 array([[1, 2, 3], [4, 5]])

artiq/test/lit/integration/array.py

@@ -13,13 +13,12 @@ assert len(empty_array) == 0
 assert empty_array.shape == (0,)
 assert [x * x for x in empty_array] == []
 
-# Creating a list from a generic iterable always generates an 1D array, as we can't
+# Creating arrays from generic iterables, rectangularity is assumed (and ensured
-# check for rectangularity at compile time. (This could be changed to *assume*
+# with runtime checks).
-# rectangularity and insert runtime checks instead.)
 list_of_lists = [[1, 2], [3, 4]]
 array_of_lists = array(list_of_lists)
-assert array_of_lists.shape == (2,)
+assert array_of_lists.shape == (2, 2)
-assert [x for x in array_of_lists] == list_of_lists
+assert [[y for y in x] for x in array_of_lists] == list_of_lists
 
 matrix = array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
 assert len(matrix) == 2