compiler: Implement multi-dimensional indexing of arrays

This generates rather more code than necessary, but has
the advantage of automatically handling incomplete
multi-dimensional subscripts which still leave arrays
behind.

artiq/compiler/transforms/artiq_ir_generator.py

@@ -1092,16 +1092,31 @@ class ARTIQIRGenerator(algorithm.Visitor):
             finally:
                 self.current_assign = old_assign
 
-            length = self.iterable_len(value, index.type)
-            mapped_index = self._map_index(length, index,
-                                           loc=node.begin_loc)
-            if self.current_assign is None:
-                result = self.iterable_get(value, mapped_index)
-                result.set_name("{}.at.{}".format(value.name, _readable_name(index)))
-                return result
+            # For multi-dimensional indexes, just apply them sequentially. This
+            # works, as they are only supported for types where we do not
+            # immediately need to distinguish between the Get and Set cases
+            # (i.e. arrays, which are reference types).
+            if types.is_tuple(index.type):
+                num_idxs = len(index.type.find().elts)
+                indices = [
+                    self.append(ir.GetAttr(index, i)) for i in range(num_idxs)
+                ]
             else:
-                self.append(ir.SetElem(value, mapped_index, self.current_assign,
-                                       name="{}.at.{}".format(value.name, _readable_name(index))))
+                indices = [index]
+            indexed = value
+            for i, idx in enumerate(indices):
+                length = self.iterable_len(indexed, idx.type)
+                mapped_index = self._map_index(length, idx, loc=node.begin_loc)
+                if self.current_assign is None or i < len(indices) - 1:
+                    indexed = self.iterable_get(indexed, mapped_index)
+                    indexed.set_name("{}.at.{}".format(indexed.name,
+                                                       _readable_name(idx)))
+                else:
+                    self.append(ir.SetElem(indexed, mapped_index, self.current_assign,
+                                           name="{}.at.{}".format(value.name,
+                                                                  _readable_name(index))))
+            if self.current_assign is None:
+                return indexed
         else: # Slice
             length = self.iterable_len(value, node.slice.type)
 

artiq/compiler/transforms/inferencer.py

@@ -208,10 +208,9 @@ class Inferencer(algorithm.Visitor):
         self.generic_visit(node)
         value = node.value
         if types.is_tuple(value.type):
-            diag = diagnostic.Diagnostic("error",
-                "multi-dimensional slices are not supported", {},
-                node.loc, [])
-            self.engine.process(diag)
+            for elt in value.type.find().elts:
+                self._unify(elt, builtins.TInt(),
+                            value.loc, None)
         else:
             self._unify(value.type, builtins.TInt(),
                         value.loc, None)
@@ -237,12 +236,39 @@ class Inferencer(algorithm.Visitor):
     def visit_SubscriptT(self, node):
         self.generic_visit(node)
         if isinstance(node.slice, ast.Index):
-            self._unify_iterable(element=node, collection=node.value)
+            if types.is_tuple(node.slice.value.type):
+                if not builtins.is_array(node.value.type):
+                    diag = diagnostic.Diagnostic(
+                        "error",
+                        "multi-dimensional slices only supported for arrays, not {type}",
+                        {"type": types.TypePrinter().name(node.value.type)},
+                        node.loc, [])
+                    self.engine.process(diag)
+                    return
+                num_idxs = len(node.slice.value.type.find().elts)
+                array_type = node.value.type.find()
+                num_dims = array_type["num_dims"].value
+                remaining_dims = num_dims - num_idxs
+                if remaining_dims < 0:
+                    diag = diagnostic.Diagnostic(
+                        "error",
+                        "too many indices for array of dimension {num_dims}",
+                        {"num_dims": num_dims}, node.slice.loc, [])
+                    self.engine.process(diag)
+                    return
+                if remaining_dims == 0:
+                    self._unify(node.type, array_type["elt"], node.loc,
+                                node.value.loc)
+                else:
+                    self._unify(
+                        node.type,
+                        builtins.TArray(array_type["elt"], remaining_dims))
+            else:
+                self._unify_iterable(element=node, collection=node.value)
         elif isinstance(node.slice, ast.Slice):
-            self._unify(node.type, node.value.type,
-                        node.loc, node.value.loc)
-        else: # ExtSlice
-            pass # error emitted above
+            self._unify(node.type, node.value.type, node.loc, node.value.loc)
+        else:  # ExtSlice
+            pass  # error emitted above
 
     def visit_IfExpT(self, node):
         self.generic_visit(node)

artiq/test/lit/inferencer/error_array.py

@@ -5,5 +5,9 @@
 a = array()
 
 b = array([1, 2, 3])
+
+# CHECK-L: ${LINE:+1}: error: too many indices for array of dimension 1
+b[1, 2]
+
 # CHECK-L: ${LINE:+1}: error: array attributes cannot be assigned to
 b.shape = (5, )

artiq/test/lit/integration/array.py

@@ -26,21 +26,21 @@ assert matrix.shape == (2, 3)
 # FIXME: Need to decide on a solution for array comparisons —
 # NumPy returns an array of bools!
 # assert [x for x in matrix] == [array([1.0, 2.0, 3.0]), array([4.0, 5.0, 6.0])]
-assert matrix[0][0] == 1.0
-assert matrix[0][1] == 2.0
-assert matrix[0][2] == 3.0
-assert matrix[1][0] == 4.0
-assert matrix[1][1] == 5.0
-assert matrix[1][2] == 6.0
+assert matrix[0, 0] == 1.0
+assert matrix[0, 1] == 2.0
+assert matrix[0, 2] == 3.0
+assert matrix[1, 0] == 4.0
+assert matrix[1, 1] == 5.0
+assert matrix[1, 2] == 6.0
 
-matrix[0][0] = 7.0
-matrix[1][1] = 8.0
-assert matrix[0][0] == 7.0
-assert matrix[0][1] == 2.0
-assert matrix[0][2] == 3.0
-assert matrix[1][0] == 4.0
-assert matrix[1][1] == 8.0
-assert matrix[1][2] == 6.0
+matrix[0, 0] = 7.0
+matrix[1, 1] = 8.0
+assert matrix[0, 0] == 7.0
+assert matrix[0, 1] == 2.0
+assert matrix[0, 2] == 3.0
+assert matrix[1, 0] == 4.0
+assert matrix[1, 1] == 8.0
+assert matrix[1, 2] == 6.0
 
 three_tensor = array([[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]])
 assert len(three_tensor) == 1