Allow indexing tuples in kernel code

This only allows for indexing with a constant value (e.g. x[0]).

While slices would be possible to implement, it's not clear how to
preserve type inference here. The current typing rule is:

  Γ ⊢ x : τ  Γ ⊢ a : Int  Γ ⊢ b : Int
  ------------------------------------
             Γ ⊢ x[a:b] : τ

However, tuples would require a different typing rule, and so we'd need
to defer type inference if τ is a tyvar. I'm not confident that this
won't change behaviour, so we leave as-is for now.

Signed-off-by: Jonathan Coates <jonathan.coates@oxionics.com>

artiq/compiler/transforms/artiq_ir_generator.py

@@ -1198,7 +1198,27 @@ class ARTIQIRGenerator(algorithm.Visitor):
         finally:
             self.current_assign = old_assign
 
-        if isinstance(node.slice, ast.Index):
+        if types.is_tuple(node.value.type):
+            assert isinstance(node.slice, ast.Index), \
+                "Internal compiler error: tuple index should be an Index"
+            assert isinstance(node.slice.value, ast.Num), \
+                "Internal compiler error: tuple index should be a constant"
+
+            if self.current_assign is not None:
+                diag = diagnostic.Diagnostic("error",
+                    "cannot assign to a tuple element",
+                    {}, node.loc)
+                self.engine.process(diag)
+
+            index = node.slice.value.n
+            indexed = self.append(
+                ir.GetAttr(value, index, name="{}.e{}".format(value.name, index)),
+                loc=node.loc
+            )
+
+            return indexed
+
+        elif isinstance(node.slice, ast.Index):
             try:
                 old_assign, self.current_assign = self.current_assign, None
                 index = self.visit(node.slice.value)

artiq/compiler/transforms/inferencer.py

@@ -259,7 +259,31 @@ class Inferencer(algorithm.Visitor):
 
     def visit_SubscriptT(self, node):
         self.generic_visit(node)
-        if isinstance(node.slice, ast.Index):
+
+        if types.is_tuple(node.value.type):
+            if (not isinstance(node.slice, ast.Index) or
+                not isinstance(node.slice.value, ast.Num)):
+                diag = diagnostic.Diagnostic(
+                    "error", "tuples can only be indexed by a constant", {},
+                    node.slice.loc, []
+                )
+                self.engine.process(diag)
+                return
+
+            tuple_type = node.value.type.find()
+            index = node.slice.value.n
+            if index < 0 or index >= len(tuple_type.elts):
+                diag = diagnostic.Diagnostic(
+                    "error",
+                    "index {index} is out of range for tuple of size {size}",
+                    {"index": index, "size": len(tuple_type.elts)},
+                    node.slice.loc, []
+                )
+                self.engine.process(diag)
+                return
+
+            self._unify(node.type, tuple_type.elts[index], node.loc, node.value.loc)
+        elif isinstance(node.slice, ast.Index):
             if types.is_tuple(node.slice.value.type):
                 if types.is_var(node.value.type):
                     return

artiq/test/lit/embedding/error_tuple_index_assign.py

@@ -0,0 +1,15 @@
+# RUN: %python -m artiq.compiler.testbench.embedding +diag %s 2>%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+from artiq.language.core import *
+from artiq.language.types import *
+
+@kernel
+def modify(x):
+    # CHECK-L: ${LINE:+1}: error: cannot assign to a tuple element
+    x[0] = 2
+
+@kernel
+def entrypoint():
+    modify((1, "foo", True))
+    modify((2, "bar", False))

artiq/test/lit/embedding/index_tuple.py

@@ -0,0 +1,16 @@
+# RUN: env ARTIQ_DUMP_LLVM=%t %python -m artiq.compiler.testbench.embedding +compile %s
+# RUN: OutputCheck %s --file-to-check=%t.ll
+
+from artiq.language.core import *
+from artiq.language.types import *
+
+# CHECK-L: void @_Z16testbench.unpackzz({ i32, { i8*, i32 }, i1 } %ARG.x)
+
+@kernel
+def unpack(x):
+    print(x[0])
+
+@kernel
+def entrypoint():
+    unpack((1, "foo", True))
+    unpack((2, "bar", False))

artiq/test/lit/inferencer/error_tuple_index.py

@@ -0,0 +1,14 @@
+# RUN: %python -m artiq.compiler.testbench.inferencer +diag %s >%t
+# RUN: OutputCheck %s --file-to-check=%t
+
+i = 0
+x = (1, "foo", True)
+
+# CHECK-L: ${LINE:+1}: error: tuples can only be indexed by a constant
+x[i]
+
+# CHECK-L: ${LINE:+1}: error: tuples can only be indexed by a constant
+x[0:2]
+
+# CHECK-L: ${LINE:+1}: error: index 3 is out of range for tuple of size 3
+x[3]

artiq/test/lit/integration/tuple_index.py

@@ -0,0 +1,22 @@
+# RUN: %python -m artiq.compiler.testbench.jit %s
+# RUN: %python %s
+
+# Basic indexing
+a = (1, "xyz", True)
+
+assert a[0] == 1
+assert a[1] == "xyz"
+assert a[2] == True
+
+# Nested indexing
+b = (a, 2, (3, "abc", a))
+
+assert b[0][0] == 1
+assert b[1] == 2
+assert b[2][2][0] == 1
+
+# Usage on the LHS of an assignment
+c = (1, 2, [1, 2, 3])
+
+c[2][0] = 456
+assert c[2][0] == 456