Initial StateRegister implementation.

docs/basic.rst

@@ -23,6 +23,7 @@ hard-wired values and :class:`.Register` is how sequential elements are created
                          pyrtl.Output
                          pyrtl.Const
                          pyrtl.Register
+                         pyrtl.StateRegister
     :parts: 1
 
 WireVector

docs/index.rst

@@ -92,6 +92,7 @@ this is all a user needs to create a functional hardware design.
                          pyrtl.Output
                          pyrtl.Const
                          pyrtl.Register
+                         pyrtl.StateRegister
     :parts: 1
 
 After specifying a hardware design, there are then options to simulate your

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=uv run python3
+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.
@@ -1956,7 +1956,7 @@ def print_perf_counters(self, *trace_names: str, file=sys.stdout):
 
 
 def enum_name(EnumClass: type) -> Callable[[int], str]:
-    """Returns a function that returns the name of an :class:`enum.IntEnum` value.
+    """Returns a function that returns the name of an :class:`~enum.IntEnum` value.
 
     .. doctest only::
 
@@ -1965,32 +1965,44 @@ def enum_name(EnumClass: type) -> Callable[[int], str]:
         >>> pyrtl.reset_working_block()
 
     Use ``enum_name`` as a ``repr_func`` or ``repr_per_name`` for
-    :meth:`SimulationTrace.render_trace` to display :class:`enum.IntEnum` names in
+    :meth:`~SimulationTrace.render_trace` to display :class:`~enum.IntEnum` names in
     traces, instead of their numeric value. Example::
 
-        >>> class State(enum.IntEnum):
+        >>> class Option(enum.IntEnum):
         ...     FOO = 0
         ...     BAR = 1
-        >>> state = pyrtl.Input(name="state", bitwidth=1)
+        >>> pyrtl.enum_name(Option)(1)
+        'BAR'
+
+        >>> option = pyrtl.Input(name="option", bitwidth=1)
 
         >>> sim = pyrtl.Simulation()
-        >>> sim.step_multiple({"state": [State.FOO, State.BAR]})
-        >>> sim.tracer.render_trace(repr_per_name={"state": pyrtl.enum_name(State)})
+        >>> sim.step_multiple({"option": [Option.FOO, Option.BAR]})
+        >>> sim.tracer.render_trace(repr_per_name={"option": pyrtl.enum_name(Option)})
 
     Which prints::
 
-             │0  │1
+              │0  │1
+
+        option FOO│BAR
+
+    .. note::
 
-        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``.
+    :param EnumClass: ``enum`` to convert. This is the enum class, like ``Option``, not
+        an enum value, like ``Option.FOO`` or ``1``.
 
-    :return: A function that accepts an enum value, like ``State.FOO`` or ``1``, and
-             returns the value's name as a string, like ``"FOO"``.
+    :return: A function that accepts an enum value, like ``Option.FOO`` or ``1``, and
+             returns the value's name as a string, like ``"FOO"``. Unknown values will
+             be converted to string with :class:`hex`.
     """
 
     def value_to_name(value: int) -> str:
-        return EnumClass(value).name
+        try:
+            return EnumClass(value).name
+        except ValueError:
+            return hex(value)
 
     return value_to_name