lib.data: bring Const's ArrayLayout capabilities up to par with views.

Fixes #1486.

Author: wanda-phi

amaranth/lib/data.py

@@ -814,7 +814,7 @@ def __getitem__(self, key):
                     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)):
+            elif isinstance(key, (Value, ValueCastable)):
                 shape = self.__layout.elem_shape
                 value = self.__target.word_select(key, elem_width)
             else:
@@ -1044,16 +1044,38 @@ def __getitem__(self, key):
             :meth:`.ShapeCastable.from_bits`. Usually this will be a :exc:`ValueError`.
         """
         if isinstance(self.__layout, ArrayLayout):
-            if isinstance(key, (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) & ((1 << elem_width * (stop - start)) - 1)
+                else:
+                    value = 0
+                    pos = 0
+                    for index in range(start, stop, stride):
+                        elem_value = (self.__target >> index * elem_width) & ((1 << elem_width) - 1)
+                        value |= elem_value << pos
+                        pos += elem_width
+            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) & ((1 << elem_width) - 1)
+            elif isinstance(key, (Value, ValueCastable)):
                 return View(self.__layout, self.as_value())[key]
-            if not isinstance(key, int):
+            else:
                 raise TypeError(
                     f"Constant with array layout may only be indexed with an integer or a value, "
                     f"not {key!r}")
-            shape = self.__layout.elem_shape
-            elem_width = Shape.cast(self.__layout.elem_shape).width
-            value = (self.__target >> key * elem_width) & ((1 << elem_width) - 1)
         else:
+            if isinstance(key, slice):
+                raise TypeError(
+                    "Non-array constant cannot be indexed with a slice; did you mean to call "
+                    "`.as_value()` first?")
             if isinstance(key, (Value, ValueCastable)):
                 raise TypeError(
                     f"Only constants with array layout, not {self.__layout!r}, may be indexed with "
@@ -1096,6 +1118,12 @@ def __getattr__(self, name):
                 f"may only be accessed by indexing")
         return item
 
+    def __len__(self):
+        if not isinstance(self.__layout, ArrayLayout):
+            raise TypeError(
+                f"`len()` can only be used on constants of array layout, not {self.__layout!r}")
+        return self.__layout.length
+
     def __eq__(self, other):
         if isinstance(other, View) and self.__layout == other._View__layout:
             return self.as_value() == other._View__target

docs/changes.rst

@@ -66,6 +66,20 @@ Platform integration changes
 * Changed: the Gowin platform now uses ``nextpnr-himbaechel`` rather than ``nextpnr-gowin``.
 
 
+Version 0.5.2 (unreleased)
+==========================
+
+
+Standard library changes
+------------------------
+
+.. currentmodule:: amaranth.lib
+
+* Added: constants of :class:`amaranth.lib.data.ArrayLayout` can be indexed with negative integers or slices.
+* Added: :py:`len()` works on constants of :class:`amaranth.lib.data.ArrayLayout`.
+* Added: constants of :class:`amaranth.lib.data.ArrayLayout` are iterable.
+
+
 Version 0.5.1
 =============
 

tests/test_lib_data.py

@@ -939,12 +939,28 @@ def test_getitem(self):
         self.assertEqual(v["q"], -1)
         self.assertEqual(v["r"][0], 3)
         self.assertEqual(v["r"][1], 2)
+        self.assertEqual(v["r"][-2], 3)
+        self.assertEqual(v["r"][-1], 2)
         self.assertRepr(v["r"][i], "(part (const 4'd11) (sig i) 2 2)")
         self.assertEqual(v["t"][0], data.Const(l, 2))
         self.assertEqual(v["t"][1], data.Const(l, 2))
         self.assertEqual(v["t"][0]["u"], 0)
         self.assertEqual(v["t"][1]["v"], 1)
 
+    def test_getitem_slice(self):
+        def A(n):
+            return data.ArrayLayout(unsigned(4), n)
+        v = data.Const(data.ArrayLayout(unsigned(4), 5), 0xabcde)
+        self.assertEqual(v[1:3], data.Const(A(2), 0xcd))
+        self.assertEqual(v[2:], data.Const(A(3), 0xabc))
+        self.assertEqual(v[:-2], data.Const(A(3), 0xcde))
+        self.assertEqual(v[-1:], data.Const(A(1), 0xa))
+        self.assertEqual(v[::-1], data.Const(A(5), 0xedcba))
+
+    def test_array_iter(self):
+        v = data.Const(data.ArrayLayout(unsigned(4), 5), 0xabcde)
+        self.assertEqual(list(v), [0xe, 0xd, 0xc, 0xb, 0xa])
+
     def test_getitem_custom_call(self):
         class Reverser(ShapeCastable):
             def as_shape(self):
@@ -1105,6 +1121,14 @@ def test_compare(self):
                 r"a constant with the same layout, not .*$"):
             c5 != [0,1,2,3,4]
 
+    def test_len(self):
+        c1 = data.Const(data.StructLayout({"a": unsigned(2)}), 2)
+        with self.assertRaisesRegex(TypeError,
+                r"^`len\(\)` can only be used on constants of array layout, not StructLayout.*$"):
+            len(c1)
+        c2 = data.Const(data.ArrayLayout(2, 3), 0x12)
+        self.assertEqual(len(c2), 3)
+
     def test_operator(self):
         s1 = data.Const(data.StructLayout({"a": unsigned(2)}), 2)
         s2 = Signal(unsigned(2))