lib.data: accept negative int and slice for ArrayLayout indexing.

As a side effect, makes iterating over `ArrayLayout` views work.

Fixes #1405.
Fixes #1418.

Author: wanda-phi

amaranth/lib/data.py

@@ -796,17 +796,36 @@ def __getitem__(self, key):
         :exc:`TypeError`
             If :meth:`.ShapeCastable.__call__` does not return a value or a value-castable object.
         """
-        if isinstance(key, slice):
-            raise TypeError(
-                "View cannot be indexed with a slice; did you mean to call `.as_value()` first?")
         if isinstance(self.__layout, ArrayLayout):
-            if not isinstance(key, (int, Value, ValueCastable)):
+            elem_width = Shape.cast(self.__layout.elem_shape).width
+            if isinstance(key, slice):
+                start, stop, stride = key.indices(self.__layout.length)
+                shape = ArrayLayout(self.__layout.elem_shape, len(range(start, stop, stride)))
+                if stride == 1:
+                    value = self.__target[start * elem_width:stop * elem_width]
+                else:
+                    value = Cat(self.__target[index * elem_width:(index + 1) * elem_width]
+                                for index in range(start, stop, stride))
+            elif isinstance(key, int):
+                if key not in range(-self.__layout.length, self.__layout.length):
+                    raise IndexError(f"Index {key} is out of range for array layout of length "
+                                     f"{self.__layout.length}")
+                if key < 0:
+                    key += self.__layout.length
+                shape = self.__layout.elem_shape
+                value = self.__target[key * elem_width:(key + 1) * elem_width]
+            elif isinstance(key, (int, Value, ValueCastable)):
+                shape = self.__layout.elem_shape
+                value = self.__target.word_select(key, elem_width)
+            else:
                 raise TypeError(
                     f"View with array layout may only be indexed with an integer or a value, "
                     f"not {key!r}")
-            shape = self.__layout.elem_shape
-            value = self.__target.word_select(key, Shape.cast(self.__layout.elem_shape).width)
         else:
+            if isinstance(key, slice):
+                raise TypeError(
+                    "Non-array view cannot be indexed with a slice; did you mean to call "
+                    "`.as_value()` first?")
             if isinstance(key, (Value, ValueCastable)):
                 raise TypeError(
                     f"Only views with array layout, not {self.__layout!r}, may be indexed with "

tests/test_lib_data.py

@@ -614,10 +614,52 @@ def test_getitem(self):
         self.assertEqual(v["q"].shape(), signed(1))
         self.assertRepr(v["r"][0], "(slice (slice (sig v) 0:4) 0:2)")
         self.assertRepr(v["r"][1], "(slice (slice (sig v) 0:4) 2:4)")
+        self.assertRepr(v["r"][-2], "(slice (slice (sig v) 0:4) 0:2)")
+        self.assertRepr(v["r"][-1], "(slice (slice (sig v) 0:4) 2:4)")
         self.assertRepr(v["r"][i], "(part (slice (sig v) 0:4) (sig i) 2 2)")
         self.assertRepr(v["t"][0]["u"], "(slice (slice (slice (sig v) 0:4) 0:2) 0:1)")
         self.assertRepr(v["t"][1]["v"], "(slice (slice (slice (sig v) 0:4) 2:4) 1:2)")
 
+    def test_getitem_slice(self):
+        a = Signal(data.ArrayLayout(unsigned(2), 5))
+        self.assertEqual(a[1:3].shape(), data.ArrayLayout(unsigned(2), 2))
+        self.assertRepr(a[1:3].as_value(), "(slice (sig a) 2:6)")
+        self.assertRepr(a[2:].as_value(), "(slice (sig a) 4:10)")
+        self.assertRepr(a[:-2].as_value(), "(slice (sig a) 0:6)")
+        self.assertRepr(a[-1:].as_value(), "(slice (sig a) 8:10)")
+        self.assertRepr(a[::-1].as_value(), """
+            (cat
+                (slice (sig a) 8:10)
+                (slice (sig a) 6:8)
+                (slice (sig a) 4:6)
+                (slice (sig a) 2:4)
+                (slice (sig a) 0:2)
+            )
+        """)
+        self.assertRepr(a[::2].as_value(), """
+            (cat
+                (slice (sig a) 0:2)
+                (slice (sig a) 4:6)
+                (slice (sig a) 8:10)
+            )
+        """)
+        self.assertRepr(a[1::2].as_value(), """
+            (cat
+                (slice (sig a) 2:4)
+                (slice (sig a) 6:8)
+            )
+        """)
+
+    def test_array_iter(self):
+        a = Signal(data.ArrayLayout(unsigned(2), 5))
+        l = list(a)
+        self.assertEqual(len(l), 5)
+        self.assertRepr(l[0], "(slice (sig a) 0:2)")
+        self.assertRepr(l[1], "(slice (sig a) 2:4)")
+        self.assertRepr(l[2], "(slice (sig a) 4:6)")
+        self.assertRepr(l[3], "(slice (sig a) 6:8)")
+        self.assertRepr(l[4], "(slice (sig a) 8:10)")
+
     def test_getitem_custom_call(self):
         class Reverser(ShapeCastable):
             def as_shape(self):
@@ -674,9 +716,17 @@ def test_index_wrong_struct_dynamic(self):
                 r"with a value$"):
             Signal(data.StructLayout({}))[Signal(1)]
 
+    def test_index_wrong_oob(self):
+        with self.assertRaisesRegex(IndexError,
+                r"^Index 2 is out of range for array layout of length 2$"):
+            Signal(data.ArrayLayout(unsigned(2), 2))[2]
+        with self.assertRaisesRegex(IndexError,
+                r"^Index -3 is out of range for array layout of length 2$"):
+            Signal(data.ArrayLayout(unsigned(2), 2))[-3]
+
     def test_index_wrong_slice(self):
         with self.assertRaisesRegex(TypeError,
-                r"^View cannot be indexed with a slice; did you mean to call `.as_value\(\)` "
+                r"^Non-array view cannot be indexed with a slice; did you mean to call `.as_value\(\)` "
                 r"first\?$"):
             Signal(data.StructLayout({}))[0:1]