Initial StateRegister implementation

Author: fdxmw

docs/regmem.rst

@@ -9,6 +9,11 @@ Registers
     :show-inheritance:
     :special-members: __init__
 
+.. autoclass:: pyrtl.StateRegister
+    :members:
+    :show-inheritance:
+    :special-members: __init__
+
 Memories
 --------
 

examples/Makefile

@@ -1,4 +1,4 @@
-PYTHON=python
+PYTHON=uv run
 PY_FILES=$(wildcard *.py)
 IPYNB_FILES=$(addprefix ../ipynb-examples/, $(PY_FILES:.py=.ipynb))
 

examples/example3-statemachine.py

@@ -15,10 +15,8 @@
 dispense = pyrtl.Output(1, "dispense")
 refund = pyrtl.Output(1, "refund")
 
-state = pyrtl.Register(3, "state")
 
-
-# First new step, let's enumerate a set of constants to serve as our states
+# First new step, let's enumerate a set of constants for all possible states.
 class State(enum.IntEnum):
     WAIT = 0  # Waiting for first token.
     TOK1 = 1  # Received first token, waiting for second token.
@@ -28,6 +26,11 @@ class State(enum.IntEnum):
     RFND = 5  # Issue refund.
 
 
+# Define a `StateRegister`, which is just like a `Register`, except that it calculates
+# the `Register`'s bitwidth from the largest possible `State`. `StateRegister`s also
+# display state names in traces by default.
+state = pyrtl.StateRegister(State, "state")
+
 # Now we could build a state machine using just the `Registers` and logic discussed in
 # prior examples, but doing operations **conditionally** on some input is a pretty
 # fundamental operation in hardware design. PyRTL provides `conditional_assignment` to
@@ -114,11 +117,9 @@ class State(enum.IntEnum):
 sim.step_multiple(sim_inputs)
 
 # Also, to make our input/output easy to reason about let's specify an order to the
-# traces with `trace_list`. We also use `enum_name` to display the state names (`WAIT`,
-# `TOK1`, ...) rather than their numbers (0, 1, ...).
+# traces with `trace_list`.
 sim.tracer.render_trace(
-    trace_list=["token_in", "req_refund", "state", "dispense", "refund"],
-    repr_per_name={"state": pyrtl.enum_name(State)},
+    trace_list=["token_in", "req_refund", "state", "dispense", "refund"]
 )
 
 # Finally, suppose you want to simulate your design and verify its output matches your

ipynb-examples/example3-statemachine.ipynb

@@ -48,16 +48,14 @@
     "req_refund = pyrtl.Input(1, \"req_refund\")\n",
     "\n",
     "dispense = pyrtl.Output(1, \"dispense\")\n",
-    "refund = pyrtl.Output(1, \"refund\")\n",
-    "\n",
-    "state = pyrtl.Register(3, \"state\")\n"
+    "refund = pyrtl.Output(1, \"refund\")\n"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    " First new step, let's enumerate a set of constants to serve as our states\n"
+    " First new step, let's enumerate a set of constants for all possible states.\n"
    ]
   },
   {
@@ -77,6 +75,26 @@
     "    RFND = 5  # Issue refund.\n"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    " Define a `StateRegister`, which is just like a [Register](https://pyrtl.readthedocs.io/en/latest/basic.html#pyrtl.Register), except that it calculates\n",
+    " the [Register](https://pyrtl.readthedocs.io/en/latest/basic.html#pyrtl.Register)'s bitwidth from the largest possible `State`. `StateRegister`s also\n",
+    " display state names in traces by default.\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "collapsed": true
+   },
+   "outputs": [],
+   "source": [
+    "state = pyrtl.StateRegister(State, \"state\")\n"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -216,8 +234,7 @@
    "metadata": {},
    "source": [
     " Also, to make our input/output easy to reason about let's specify an order to the\n",
-    " traces with `trace_list`. We also use `enum_name` to display the state names (`WAIT`,\n",
-    " `TOK1`, ...) rather than their numbers (0, 1, ...).\n"
+    " traces with `trace_list`.\n"
    ]
   },
   {
@@ -229,8 +246,7 @@
    "outputs": [],
    "source": [
     "sim.tracer.render_trace(\n",
-    "    trace_list=[\"token_in\", \"req_refund\", \"state\", \"dispense\", \"refund\"],\n",
-    "    repr_per_name={\"state\": pyrtl.enum_name(State)},\n",
+    "    trace_list=[\"token_in\", \"req_refund\", \"state\", \"dispense\", \"refund\"]\n",
     ")\n"
    ]
   },

pyrtl/__init__.py

@@ -17,7 +17,7 @@
 )
 
 # convenience classes for building hardware
-from .wire import WireVector, Input, Output, Const, Register
+from .wire import WireVector, Input, Output, Const, Register, StateRegister
 
 from .gate_graph import GateGraph, Gate
 
@@ -161,6 +161,7 @@
     "Output",
     "Const",
     "Register",
+    "StateRegister",
     # gate_graph
     "GateGraph",
     "Gate",

pyrtl/simulation.py

@@ -20,7 +20,7 @@
 from pyrtl.importexport import _VerilogSanitizer
 from pyrtl.memory import MemBlock, RomBlock
 from pyrtl.pyrtlexceptions import PyrtlError, PyrtlInternalError
-from pyrtl.wire import Const, Input, Output, Register, WireVector
+from pyrtl.wire import Const, Input, Output, Register, StateRegister, WireVector
 
 # ----------------------------------------------------------------
 #    __                         ___    __
@@ -1106,6 +1106,8 @@ def invoke_f(f, value):
 
         if f is not None:
             return invoke_f(f, value)
+        if isinstance(wire, StateRegister):
+            return invoke_f(enum_name(wire.States), value)
         return invoke_f(repr_func, value)
 
     def render_val(
@@ -1141,20 +1143,18 @@ def render_val(
             _prev_line* fields in RendererConstants.
         :param is_last: If True, current_val is in the last cycle.
         """
-        if len(w) > 1 or w.name in repr_per_name:
+        if len(w) > 1 or w.name in repr_per_name or isinstance(w, StateRegister):
             # Render values in boxes for multi-bit wires ("bus"), or single-bit wires
             # with a specific representation.
             #
             # We display multi-wire zero values as a centered horizontal line when a
             # specific `repr_per_name` is not requested for this trace, and a standard
             # numeric format is requested.
-            flat_zero = w.name not in repr_per_name and (
-                repr_func is hex
-                or repr_func is oct
-                or repr_func is int
-                or repr_func is str
-                or repr_func is bin
-                or repr_func is val_to_signed_integer
+            numeric_formats = [hex, oct, int, str, bin, val_to_signed_integer]
+            flat_zero = (
+                w.name not in repr_per_name
+                and not isinstance(w, StateRegister)
+                and repr_func in numeric_formats
             )
             if prev_line:
                 # Bus wires are currently never rendered across multiple lines.
@@ -1980,6 +1980,9 @@ def enum_name(EnumClass: type) -> Callable[[int], str]:
 
         state FOO│BAR
 
+    When using ``enum_name`` with a :class:`.Register`, consider using
+    :class:`.StateRegister` instead.
+
     :param EnumClass: ``enum`` to convert. This is the enum class, like ``State``, not
         an enum value, like ``State.FOO`` or ``1``.
 

pyrtl/wire.py

@@ -16,6 +16,7 @@
 
 from __future__ import annotations
 
+import enum
 import numbers
 import re
 import traceback
@@ -1661,15 +1662,15 @@ def __ior__(self, _):
 
 
 class Register(WireVector):
-    """A WireVector with an embedded register state element.
+    """A :class:`WireVector` with an embedded register state element.
 
-    Registers only update their outputs on the rising edges of an implicit clock signal.
-    The "value" in the current cycle can be accessed by referencing the Register itself.
-    To set the value for the next cycle (after the next rising clock edge), set the
-    :attr:`Register.next` property with the ``<<=`` (:meth:`~WireVector.__ilshift__`)
-    operator.
+    ``Registers`` only update their outputs on the rising edges of an implicit clock
+    signal. The "value" in the current cycle can be accessed by referencing the
+    ``Register`` itself. To set the value for the next cycle (after the next rising
+    clock edge), set the :attr:`Register.next` property with the ``<<=``
+    (:meth:`~WireVector.__ilshift__`) operator.
 
-    Registers reset to zero by default, and reside in the same clock domain.
+    ``Registers`` reset to zero by default, and reside in the same clock domain.
 
     .. doctest only::
 
@@ -1689,6 +1690,10 @@ class Register(WireVector):
     This builds a zero-initialized 2-bit counter. The second line sets the counter's
     value in the next cycle (``counter.next``) to the counter's value in the current
     cycle (``counter``), plus one.
+
+    .. note::
+
+        Consider using :class:`StateRegister` for state machine ``Registers``.
     """
 
     _code = "R"
@@ -1811,20 +1816,20 @@ def __init__(
         reset_value: int | None = None,
         block: Block = None,
     ):
-        """Construct a register.
+        """Construct a ``Register``.
 
         It is an error if the ``reset_value`` cannot fit into the specified ``bitwidth``
         for this register.
 
-        :param bitwidth: Number of bits to represent this register.
-        :param name: The name of the register's current value (``reg``, not
+        :param bitwidth: Number of bits to represent this ``Register``.
+        :param name: The name of the ``Register``'s current value (``reg``, not
             ``reg.next``). Must be unique. If none is provided, one will be
             autogenerated.
-        :param reset_value: Value to initialize this register to during simulation and
-            in any code (e.g. Verilog) that is exported. Defaults to 0. Can be
+        :param reset_value: Value to initialize this ``Register`` to during simulation
+            and in any code (e.g. Verilog) that is exported. Defaults to 0. Can be
             overridden at simulation time.
-        :param block: The block under which the wire should be placed. Defaults to the
-            :ref:`working_block`.
+        :param block: The :class:`Block` under which the wire should be placed. Defaults
+            to the :ref:`working_block`.
         """
         from pyrtl.helperfuncs import infer_val_and_bitwidth
 
@@ -1871,6 +1876,74 @@ def _build(self, next):
         working_block().add_net(net)
 
 
+class StateRegister(Register):
+    """A :class:`Register` containing an :class:`~enum.IntEnum` state.
+
+    ``StateRegister`` functions identically to :class:`Register`, except that the
+    :class:`Register`'s bitwidth is calculated from an :class:`~enum.IntEnum`'s largest
+    value, and :meth:`~.SimulationTrace.render_trace` displays state names by
+    default.
+
+    .. doctest only::
+
+        >>> import pyrtl
+        >>> pyrtl.reset_working_block()
+
+    Example::
+
+        >>> import enum
+        >>> class MyStates(enum.IntEnum):
+        ...     ZERO = 0
+        ...     ONE = 1
+        ...     TWO = 2
+        ...     THREE = 3
+
+        >>> state = pyrtl.StateRegister(name="state", States=MyStates)
+        >>> state.bitwidth
+        2
+
+        >>> state.next <<= state + 1
+
+        >>> sim = pyrtl.Simulation()
+        >>> sim.step_multiple(nsteps=4)
+        >>> sim.tracer.render_trace()
+
+    Which prints::
+
+             │0    │1    │2    │3
+        state ZERO │ONE  │TWO  │THREE
+    """
+
+    def __init__(
+        self,
+        States: enum.IntEnum,
+        name: str = "",
+        reset_value: int | None = None,
+        block: Block = None,
+    ):
+        """Constructs a :class:`Register` containing an :class:`~enum.IntEnum` value.
+
+        :param States: An :class:`~enum.IntEnum` containing all possible states for the
+            ``StateRegister``. The largest value in the :class:`~enum.IntEnum`
+            determines the :class:`Register`'s :attr:`~WireVector.bitwidth`.
+        :param name: The name of the ``StateRegister``'s current value (``state``, not
+            ``state.next``). Must be unique. If none is provided, one will be
+            autogenerated.
+        :param reset_value: Value to initialize this ``StateRegister`` to during
+            simulation and in any code (e.g. Verilog) that is exported. Defaults to 0.
+            Can be overridden at simulation time.
+        :param block: The :class:`Block` under which the wire should be placed. Defaults
+            to the :ref:`working_block`.
+        """
+        from pyrtl.helperfuncs import infer_val_and_bitwidth
+
+        self.States = States
+        bitwidth = infer_val_and_bitwidth(max(States)).bitwidth
+        super().__init__(
+            bitwidth=bitwidth, name=name, reset_value=reset_value, block=block
+        )
+
+
 class WrappedWireVector:
     """Wraps a WireVector. Forwards all method calls and attribute accesses.
 

tests/test_simulation.py

@@ -273,7 +273,11 @@ class State(enum.IntEnum):
             renderer=self.renderer,
             repr_per_name={state.name: pyrtl.enum_name(State)},
         )
-        expected = "     |0  |1  \n      \nstate FOO|BAR\n"
+        expected = (
+            "     |0  |1  \n"
+            "      \n"
+            "state FOO|BAR\n"
+        )  # fmt: skip
         self.assertEqual(buff.getvalue(), expected)
 
     def test_val_to_signed_integer(self):
@@ -286,7 +290,11 @@ def test_val_to_signed_integer(self):
         sim.tracer.render_trace(
             file=buff, renderer=self.renderer, repr_func=pyrtl.val_to_signed_integer
         )
-        expected = "       |0 |1 |2 |3 \n        \ncounter --|1 |-2|-1\n"
+        expected = (
+            "       |0 |1 |2 |3 \n"
+            "        \n"
+            "counter --|1 |-2|-1\n"
+        )  # fmt: skip
         self.assertEqual(buff.getvalue(), expected)
 
     def test_custom_repr_per_wire(self):