diff --git a/amaranth/lib/data.py b/amaranth/lib/data.py index ec479b12f..d38f22a46 100644 --- a/amaranth/lib/data.py +++ b/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 diff --git a/docs/changes.rst b/docs/changes.rst index 41d2a7367..17564df70 100644 --- a/docs/changes.rst +++ b/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 ============= diff --git a/tests/test_lib_data.py b/tests/test_lib_data.py index a82fbf678..bf2583eeb 100644 --- a/tests/test_lib_data.py +++ b/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))