Repository: amaranth-lang/amaranth
Pull request #1316 (2024-04-11 06:57:59)

hdl._ast: add `Format.Enum`, `Format.Struct`, `Format.Array`.


Patch:
diff --git a/amaranth/hdl/_ast.py b/amaranth/hdl/_ast.py
index 77b3a05..4b56d07 100644
--- a/amaranth/hdl/_ast.py
+++ b/amaranth/hdl/_ast.py
@@ -2559,8 +2559,27 @@ class Initial(Value):
         return "(initial)"
 
 
+class _FormatLike:
+    def _as_format(self) -> "Format":
+        raise NotImplementedError # :nocov:
+
+    def __add__(self, other):
+        if not isinstance(other, _FormatLike):
+            return NotImplemented
+        return Format._from_chunks(self._as_format()._chunks + other._as_format()._chunks)
+
+    def __format__(self, format_desc):
+        """Forbidden formatting.
+
+        ``Format`` objects cannot be directly formatted for the same reason as the ``Value``s
+        they contain.
+        """
+        raise TypeError(f"Format object {self!r} cannot be converted to string. Use `repr` "
+                        f"to print the AST, or pass it to the `Print` statement.")
+
+
 @final
-class Format:
+class Format(_FormatLike):
     def __init__(self, format, *args, **kwargs):
         fmt = string.Formatter()
         chunks = []
@@ -2615,17 +2634,17 @@ class Format:
                     shape = obj.shape()
                     if isinstance(shape, ShapeCastable):
                         fmt = shape.format(obj, format_spec)
-                        if not isinstance(fmt, Format):
+                        if not isinstance(fmt, _FormatLike):
                             raise TypeError(f"`ShapeCastable.format` must return a 'Format' instance, not {fmt!r}")
-                        chunks += fmt._chunks
+                        chunks += fmt._as_format()._chunks
                     else:
                         obj = Value.cast(obj)
                         self._parse_format_spec(format_spec, obj.shape())
                         chunks.append((obj, format_spec))
-                elif isinstance(obj, Format):
+                elif isinstance(obj, _FormatLike):
                     if format_spec != "":
                         raise ValueError(f"Format specifiers ({format_spec!r}) cannot be used for 'Format' objects")
-                    chunks += obj._chunks
+                    chunks += obj._as_format()._chunks
                 else:
                     chunks.append(fmt.format_field(obj, format_spec))
 
@@ -2638,6 +2657,9 @@ class Format:
 
         self._chunks = self._clean_chunks(chunks)
 
+    def _as_format(self):
+        return self
+
     @classmethod
     def _from_chunks(cls, chunks):
         res = object.__new__(cls)
@@ -2671,25 +2693,11 @@ class Format:
                     format_string.append("{}")
         return ("".join(format_string), tuple(args))
 
-    def __add__(self, other):
-        if not isinstance(other, Format):
-            return NotImplemented
-        return Format._from_chunks(self._chunks + other._chunks)
-
     def __repr__(self):
         format_string, args = self._to_format_string()
         args = "".join(f" {arg!r}" for arg in args)
         return f"(format {format_string!r}{args})"
 
-    def __format__(self, format_desc):
-        """Forbidden formatting.
-
-        ``Format`` objects cannot be directly formatted for the same reason as the ``Value``s
-        they contain.
-        """
-        raise TypeError(f"Format object {self!r} cannot be converted to string. Use `repr` "
-                        f"to print the AST, or pass it to the `Print` statement.")
-
     _FORMAT_SPEC_PATTERN = re.compile(r"""
         (?:
             (?P<fill>.)?
@@ -2760,6 +2768,90 @@ class Format:
         return res
 
 
+    class Enum(_FormatLike):
+        def __init__(self, value, /, variants):
+            self._value = Value.cast(value)
+            if isinstance(variants, EnumMeta):
+                self._variants = {Const.cast(member.value).value: member.name for member in variants}
+            else:
+                self._variants = dict(variants)
+            for val, name in self._variants.items():
+                if not isinstance(val, int):
+                    raise TypeError(f"Variant values must be integers, not {val!r}")
+                if not isinstance(name, str):
+                    raise TypeError(f"Variant names must be strings, not {name!r}")
+
+        def _as_format(self):
+            def str_val(name):
+                name = name.encode()
+                return Const(int.from_bytes(name, "little"), len(name) * 8)
+            value = SwitchValue(self._value, [
+                (val, str_val(name))
+                for val, name in self._variants.items()
+            ] + [(None, str_val("[unknown]"))])
+            return Format("{:s}", value)
+
+        def __repr__(self):
+            variants = "".join(
+                f" ({val!r} {name!r})"
+                for val, name in self._variants.items()
+            )
+            return f"(format-enum {self._value!r}{variants})"
+
+
+    class Struct(_FormatLike):
+        def __init__(self, value, /, fields):
+            self._value = Value.cast(value)
+            self._fields: dict[str, _FormatLike] = dict(fields)
+            for name, format in self._fields.items():
+                if not isinstance(name, str):
+                    raise TypeError(f"Field names must be strings, not {name!r}")
+                if not isinstance(format, _FormatLike):
+                    raise TypeError(f"Field format must be a 'Format', not {format!r}")
+
+        def _as_format(self):
+            chunks = ["{"]
+            for idx, (name, field) in enumerate(self._fields.items()):
+                if idx != 0:
+                    chunks.append(", ")
+                chunks.append(f"{name}=")
+                chunks += field._as_format()._chunks
+            chunks.append("}")
+            return Format._from_chunks(chunks)
+
+        def __repr__(self):
+            fields = "".join(
+                f" ({name!r} {field!r})"
+                for name, field in self._fields.items()
+            )
+            return f"(format-struct {self._value!r}{fields})"
+
+
+    class Array(_FormatLike):
+        def __init__(self, value, /, fields):
+            self._value = Value.cast(value)
+            self._fields = list(fields)
+            for format in self._fields:
+                if not isinstance(format, (Format, Format.Enum, Format.Struct, Format.Array)):
+                    raise TypeError(f"Field format must be a 'Format', not {format!r}")
+
+        def _as_format(self):
+            chunks = ["["]
+            for idx, field in enumerate(self._fields):
+                if idx != 0:
+                    chunks.append(", ")
+                chunks += field._as_format()._chunks
+            chunks.append("]")
+            return Format._from_chunks(chunks)
+
+        def __repr__(self):
+            fields = "".join(
+                f" {field!r}"
+                for field in self._fields
+            )
+            return f"(format-array {self._value!r}{fields})"
+
+
 class _StatementList(list):
     def __repr__(self):
         return "({})".format(" ".join(map(repr, self)))
@@ -2872,8 +2964,10 @@ class Property(Statement, MustUse):
         self._test   = Value.cast(test)
         if isinstance(message, str):
             message = Format._from_chunks([message])
-        if message is not None and not isinstance(message, Format):
-            raise TypeError(f"Property message must be None, str, or Format, not {message!r}")
+        if message is not None:
+            if not isinstance(message, _FormatLike):
+                raise TypeError(f"Property message must be None, str, or Format, not {message!r}")
+            message = message._as_format()
         self._message = message
         del self._MustUse__silence
 
diff --git a/amaranth/lib/enum.py b/amaranth/lib/enum.py
index 6a2fce7..f957222 100644
--- a/amaranth/lib/enum.py
+++ b/amaranth/lib/enum.py
@@ -14,14 +14,14 @@ for _member in py_enum.__all__:
 del _member
 
 
-class EnumMeta(ShapeCastable, py_enum.EnumMeta):
-    """Subclass of the standard :class:`enum.EnumMeta` that implements the :class:`ShapeCastable`
+class EnumType(ShapeCastable, py_enum.EnumMeta):
+    """Subclass of the standard :class:`enum.EnumType` that implements the :class:`ShapeCastable`
     protocol.
 
     This metaclass provides the :meth:`as_shape` method, making its instances
     :ref:`shape-like <lang-shapelike>`, and accepts a ``shape=`` keyword argument
     to specify a shape explicitly. Other than this, it acts the same as the standard
-    :class:`enum.EnumMeta` class; if the ``shape=`` argument is not specified and
+    :class:`enum.EnumType` class; if the ``shape=`` argument is not specified and
     :meth:`as_shape` is never called, it places no restrictions on the enumeration class
     or the values of its members.
 
@@ -180,32 +180,37 @@ class EnumMeta(ShapeCastable, py_enum.EnumMeta):
         yield Repr(FormatEnum(cls), value)
 
 
+# In 3.11, Python renamed EnumMeta to EnumType. Like Python itself, we support both for
+# compatibility.
+EnumMeta = EnumType
+
+
 class Enum(py_enum.Enum):
-    """Subclass of the standard :class:`enum.Enum` that has :class:`EnumMeta` as
+    """Subclass of the standard :class:`enum.Enum` that has :class:`EnumType` as
     its metaclass and :class:`EnumView` as its view class."""
 
 
 class IntEnum(py_enum.IntEnum):
-    """Subclass of the standard :class:`enum.IntEnum` that has :class:`EnumMeta` as
+    """Subclass of the standard :class:`enum.IntEnum` that has :class:`EnumType` as
     its metaclass."""
 
 
 class Flag(py_enum.Flag):
-    """Subclass of the standard :class:`enum.Flag` that has :class:`EnumMeta` as
+    """Subclass of the standard :class:`enum.Flag` that has :class:`EnumType` as
     its metaclass and :class:`FlagView` as its view class."""
 
 
 class IntFlag(py_enum.IntFlag):
-    """Subclass of the standard :class:`enum.IntFlag` that has :class:`EnumMeta` as
+    """Subclass of the standard :class:`enum.IntFlag` that has :class:`EnumType` as
     its metaclass."""
 
 
 # Fix up the metaclass after the fact: the metaclass __new__ requires these classes
 # to already be present, and also would not install itself on them due to lack of shape.
-Enum.__class__ = EnumMeta
-IntEnum.__class__ = EnumMeta
-Flag.__class__ = EnumMeta
-IntFlag.__class__ = EnumMeta
+Enum.__class__ = EnumType
+IntEnum.__class__ = EnumType
+Flag.__class__ = EnumType
+IntFlag.__class__ = EnumType
 
 
 class EnumView(ValueCastable):
@@ -219,7 +224,7 @@ class EnumView(ValueCastable):
         """Constructs a view with the given enum type and target
         (a :ref:`value-like <lang-valuelike>`).
         """
-        if not isinstance(enum, EnumMeta) or not hasattr(enum, "_amaranth_shape_"):
+        if not isinstance(enum, EnumType) or not hasattr(enum, "_amaranth_shape_"):
             raise TypeError(f"EnumView type must be an enum with shape, not {enum!r}")
         try:
             cast_target = Value.cast(target)
diff --git a/amaranth/sim/_pyrtl.py b/amaranth/sim/_pyrtl.py
index f6c4943..5505cfb 100644
--- a/amaranth/sim/_pyrtl.py
+++ b/amaranth/sim/_pyrtl.py
@@ -153,12 +153,16 @@ class _ValueCompiler(ValueVisitor, _Compiler):
 
 
 class _RHSValueCompiler(_ValueCompiler):
-    def __init__(self, state, emitter, *, mode, inputs=None):
+    def __init__(self, state, emitter, *, mode, inputs=None, rrhs=None):
         super().__init__(state, emitter)
         assert mode in ("curr", "next")
         self.mode = mode
         # If not None, `inputs` gets populated with RHS signals.
         self.inputs = inputs
+        # When this compiler is used to grab the "next" value from within _LHSValueCompiler,
+        # we still need to use "curr" mode for reading part offsets etc. Allow setting a separate
+        # _RhsValueCompiler for these contexts.
+        self.rrhs = rrhs or self
 
     def sign(self, value):
         value_mask = (1 << len(value)) - 1
@@ -251,7 +255,7 @@ class _RHSValueCompiler(_ValueCompiler):
 
     def on_Part(self, value):
         offset_mask = (1 << len(value.offset)) - 1
-        offset = f"({value.stride} * ({offset_mask:#x} & {self(value.offset)}))"
+        offset = f"({value.stride} * ({offset_mask:#x} & {self.rrhs(value.offset)}))"
         return f"({(1 << value.width) - 1} & " \
                f"{self(value.value)} >> {offset})"
 
@@ -267,7 +271,7 @@ class _RHSValueCompiler(_ValueCompiler):
         return f"0"
 
     def on_SwitchValue(self, value):
-        gen_test = self.emitter.def_var("test", f"{(1 << len(value.test)) - 1:#x} & {self(value.test)}")
+        gen_test = self.emitter.def_var("test", f"{(1 << len(value.test)) - 1:#x} & {self.rrhs(value.test)}")
         gen_value = self.emitter.def_var("rhs_switch", "0")
         def case_handler(patterns, elem):
             self.emitter.append(f"{gen_value} = {self.sign(elem)}")
@@ -290,7 +294,7 @@ class _LHSValueCompiler(_ValueCompiler):
         self.rrhs = rhs
         # `lrhs` is used to translate the read part of a read-modify-write cycle during partial
         # update of an lvalue.
-        self.lrhs = _RHSValueCompiler(state, emitter, mode="next", inputs=None)
+        self.lrhs = _RHSValueCompiler(state, emitter, mode="next", inputs=None, rrhs=rhs)
         # If not None, `outputs` gets populated with signals on LHS.
         self.outputs = outputs
 
diff --git a/docs/stdlib/enum.rst b/docs/stdlib/enum.rst
index 9887158..52d507e 100644
--- a/docs/stdlib/enum.rst
+++ b/docs/stdlib/enum.rst
@@ -114,7 +114,7 @@ It is also possible to define a custom view class for a given enum:
 Metaclass
 =========
 
-.. autoclass:: EnumMeta()
+.. autoclass:: EnumType()
 
 
 Base classes


Test patch:
diff --git a/tests/test_hdl_ast.py b/tests/test_hdl_ast.py
index 2b6c3cb..d39da7f 100644
--- a/tests/test_hdl_ast.py
+++ b/tests/test_hdl_ast.py
@@ -1736,6 +1736,97 @@ class FormatTestCase(FHDLTestCase):
             f"{fmt}"
 
 
+class FormatEnumTestCase(FHDLTestCase):
+    def test_construct(self):
+        a = Signal(3)
+        fmt = Format.Enum(a, {1: "A", 2: "B", 3: "C"})
+        self.assertRepr(fmt, "(format-enum (sig a) (1 'A') (2 'B') (3 'C'))")
+        self.assertRepr(Format("{}", fmt), """
+            (format '{:s}' (switch-value (sig a)
+                (case 001 (const 8'd65))
+                (case 010 (const 8'd66))
+                (case 011 (const 8'd67))
+                (default (const 72'd1723507152241428428123))
+            ))
+        """)
+
+        class MyEnum(Enum):
+            A = 0
+            B = 3
+            C = 4
+
+        fmt = Format.Enum(a, MyEnum)
+        self.assertRepr(fmt, "(format-enum (sig a) (0 'A') (3 'B') (4 'C'))")
+        self.assertRepr(Format("{}", fmt), """
+            (format '{:s}' (switch-value (sig a)
+                (case 000 (const 8'd65))
+                (case 011 (const 8'd66))
+                (case 100 (const 8'd67))
+                (default (const 72'd1723507152241428428123))
+            ))
+        """)
+
+    def test_construct_wrong(self):
+        a = Signal(3)
+        with self.assertRaisesRegex(TypeError,
+                r"^Variant values must be integers, not 'a'$"):
+            Format.Enum(a, {"a": "B"})
+        with self.assertRaisesRegex(TypeError,
+                r"^Variant names must be strings, not 123$"):
+            Format.Enum(a, {1: 123})
+
+
+class FormatStructTestCase(FHDLTestCase):
+    def test_construct(self):
+        sig = Signal(3)
+        fmt = Format.Struct(sig, {"a": Format("{}", sig[0]), "b": Format("{}", sig[1:3])})
+        self.assertRepr(fmt, """
+        (format-struct (sig sig)
+            ('a' (format '{}' (slice (sig sig) 0:1)))
+            ('b' (format '{}' (slice (sig sig) 1:3)))
+        )
+        """)
+        self.assertRepr(Format("{}", fmt), """
+            (format '{{a={}, b={}}}'
+                (slice (sig sig) 0:1)
+                (slice (sig sig) 1:3)
+            )
+        """)
+
+    def test_construct_wrong(self):
+        sig = Signal(3)
+        with self.assertRaisesRegex(TypeError,
+                r"^Field names must be strings, not 1$"):
+            Format.Struct(sig, {1: Format("{}", sig[1:3])})
+        with self.assertRaisesRegex(TypeError,
+                r"^Field format must be a 'Format', not \(slice \(sig sig\) 1:3\)$"):
+            Format.Struct(sig, {"a": sig[1:3]})
+
+
+class FormatArrayTestCase(FHDLTestCase):
+    def test_construct(self):
+        sig = Signal(4)
+        fmt = Format.Array(sig, [Format("{}", sig[0:2]), Format("{}", sig[2:4])])
+        self.assertRepr(fmt, """
+        (format-array (sig sig)
+            (format '{}' (slice (sig sig) 0:2))
+            (format '{}' (slice (sig sig) 2:4))
+        )
+        """)
+        self.assertRepr(Format("{}", fmt), """
+            (format '[{}, {}]'
+                (slice (sig sig) 0:2)
+                (slice (sig sig) 2:4)
+            )
+        """)
+
+    def test_construct_wrong(self):
+        sig = Signal(3)
+        with self.assertRaisesRegex(TypeError,
+                r"^Field format must be a 'Format', not \(slice \(sig sig\) 1:3\)$"):
+            Format.Array(sig, [sig[1:3]])
+
+
 class PrintTestCase(FHDLTestCase):
     def test_construct(self):
         a = Signal()
diff --git a/tests/test_lib_enum.py b/tests/test_lib_enum.py
index af4b006..806c31e 100644
--- a/tests/test_lib_enum.py
+++ b/tests/test_lib_enum.py
@@ -5,7 +5,7 @@ import unittest
 
 from amaranth import *
 from amaranth.hdl import *
-from amaranth.lib.enum import Enum, EnumMeta, Flag, IntEnum, EnumView, FlagView
+from amaranth.lib.enum import Enum, EnumType, Flag, IntEnum, EnumView, FlagView
 
 from .utils import *
 
@@ -100,7 +100,7 @@ class EnumTestCase(FHDLTestCase):
             Z = 0
             A = 10
             B = 20
-        self.assertNotIsInstance(EnumA, EnumMeta)
+        self.assertNotIsInstance(EnumA, EnumType)
         self.assertIsInstance(EnumA, py_enum.EnumMeta)
 
     def test_const_shape(self):
diff --git a/tests/test_sim.py b/tests/test_sim.py
index fa4f881..6b9affd 100644
--- a/tests/test_sim.py
+++ b/tests/test_sim.py
@@ -366,6 +366,16 @@ class SimulatorUnitTestCase(FHDLTestCase):
         self.assertStatement(stmt, [C(1), C(0b010)], C(0b001110101, 9))
         self.assertStatement(stmt, [C(2), C(0b100)], C(0b001001001, 9))
 
+    def test_array_lhs_heterogenous_slice(self):
+        l = Signal(1, init=1)
+        m = Signal(3, init=4)
+        n = Signal(5, init=7)
+        array = Array([l, m, n])
+        stmt = lambda y, a, b: [array[a].as_value()[2:].eq(b), y.eq(Cat(*array))]
+        self.assertStatement(stmt, [C(0), C(0b000)], C(0b001111001, 9))
+        self.assertStatement(stmt, [C(1), C(0b010)], C(0b001110001, 9))
+        self.assertStatement(stmt, [C(2), C(0b100)], C(0b100111001, 9))
+
     def test_array_lhs_oob(self):
         l = Signal(3)
         m = Signal(3)