Add src parameter to WireVector, Output, and Register

pyrtl/wire.py

@@ -207,7 +207,9 @@ class WireVector:
         existing wire. If you were to do ``a = b`` it would lose the old value of ``a``
         and simply overwrite it with a new value, in this case with a reference to
         :class:`WireVector` ``b``. In contrast ``a <<= b`` does not overwrite ``a``, but
-        simply wires the two together.
+        simply wires the two together. Alternatively, ``src`` can be passed at
+        construction time: ``a = WireVector(src=b)`` is equivalent to
+        ``a = WireVector()`` followed by ``a <<= b``.
 
     .. _wirevector_coercion:
 
@@ -391,7 +393,11 @@ class WireVector:
     _code = "W"
 
     def __init__(
-        self, bitwidth: int | None = None, name: str = "", block: Block = None
+        self,
+        bitwidth: int | None = None,
+        name: str = "",
+        src: WireVectorLike | None = None,
+        block: Block = None,
     ):
         """Construct a generic :class:`WireVector`.
 
@@ -422,14 +428,23 @@ def __init__(
             >>> temp.bitwidth
             8
 
+            >>> # `result` is connected to `data` at construction time using `src`.
+            >>> result = pyrtl.WireVector(name="result", src=data)
+            >>> result.bitwidth
+            8
+
         :param bitwidth: If no ``bitwidth`` is provided, it will be set to the minimum
             number of bits needed to represent this wire.
-        :param block: The :class:`Block` under which the wire should be placed.
-            Defaults to the :ref:`working_block`.
         :param name: The name of the wire. Must be unique. If empty, a name will be
             autogenerated. If non-empty, the wire's value can be inspected with
             :meth:`Simulation.inspect`, and this wire will appear in traces generated
             by :meth:`SimulationTrace.render_trace`.
+        :param src: An optional source to connect to this wire at construction
+            time. ``WireVector(src=x)`` is equivalent to creating the wire and
+            then doing ``wire <<= x``. Accepts any
+            :ref:`WireVectorLike<wirevector_coercion>` value.
+        :param block: The :class:`Block` under which the wire should be placed.
+            Defaults to the :ref:`working_block`.
         """
         self._name = None
 
@@ -441,6 +456,9 @@ def __init__(
         if core._setting_keep_wirevector_call_stack:
             self.init_call_stack = traceback.format_stack()
 
+        if src is not None:
+            self <<= src
+
     @property
     def name(self) -> str:
         """A property holding the :class:`WireVector`'s unique name.
@@ -1665,14 +1683,35 @@ class Output(WireVector):
         ...
         pyrtl.pyrtlexceptions.PyrtlInternalError: error, Outputs cannot be arguments for
         a net
+
+    An :class:`Output` can be connected at construction time using ``src``::
+
+        >>> inp = pyrtl.Input(name="inp", bitwidth=8)
+        >>> out = pyrtl.Output(name="out", src=inp)
+        >>> out.bitwidth
+        8
     """
 
     _code = "O"
 
     def __init__(
-        self, bitwidth: int | None = None, name: str = "", block: Block = None
+        self,
+        bitwidth: int | None = None,
+        name: str = "",
+        src: WireVectorLike | None = None,
+        block: Block = None,
     ):
-        super().__init__(bitwidth, name, block)
+        """Construct an :class:`Output` wire.
+
+        :param bitwidth: Number of bits for this output.
+        :param name: The name of the output. Must be unique.
+        :param src: An optional source to connect to this output at construction
+            time. ``Output(name="x", src=w)`` is equivalent to creating the output
+            and then doing ``x <<= w``.
+        :param block: The :class:`Block` under which the wire should be placed.
+            Defaults to the :ref:`working_block`.
+        """
+        super().__init__(bitwidth=bitwidth, name=name, src=src, block=block)
 
 
 class Const(WireVector):
@@ -1810,6 +1849,20 @@ 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.
+
+    .. doctest only::
+
+        >>> pyrtl.reset_working_block()
+
+    A ``Register`` can also have its next value connected at construction time using
+    ``src``, which is equivalent to setting ``reg.next <<= src``::
+
+        >>> inp = pyrtl.Input(name="inp", bitwidth=3)
+        >>> reg = pyrtl.Register(name="reg", bitwidth=3, src=inp)
+
+    Note that self-referencing patterns like ``reg = Register(src=reg + 1)`` are not
+    possible because ``reg`` does not exist yet when ``src`` is evaluated. Use the
+    traditional ``reg.next <<= reg + 1`` pattern for those cases.
     """
 
     _code = "R"
@@ -1930,8 +1983,9 @@ def __init__(
         bitwidth: int | None = None,
         name: str = "",
         reset_value: int | None = None,
-        block: Block | None = None,
+        src: WireVectorLike | None = None,
         State: type[enum.IntEnum] | None = None,
+        block: Block | None = None,
     ):
         """Construct a ``Register``.
 
@@ -1983,15 +2037,18 @@ def __init__(
         :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 :class:`Block` under which the wire should be placed. Defaults
-            to the :ref:`working_block`.
+        :param src: An optional source to connect as the register's next value at
+            construction time. ``Register(src=x)`` is equivalent to creating the
+            register and then doing ``reg.next <<= x``.
         :param State: An :class:`~enum.IntEnum` defining all possible states for the
             ``Register``. This should be an :class:`~enum.IntEnum` class, like
             ``MyState`` in the example above. If ``bitwidth`` is ``None``, the largest
             value in the :class:`~enum.IntEnum` determines the ``Register``'s
             ``bitwidth``. When ``State`` is not ``None``,
             :meth:`~.SimulationTrace.render_trace` defaults to displaying enumeration
             names rather than hex values.
+        :param block: The :class:`Block` under which the wire should be placed. Defaults
+            to the :ref:`working_block`.
 
         :raises PyrtlError: If the ``reset_value`` or ``State`` cannot fit into the
             specified ``bitwidth`` for this register.
@@ -2027,6 +2084,9 @@ def __init__(
                 raise PyrtlError(msg)
         self.reset_value = reset_value
 
+        if src is not None:
+            self.next <<= src
+
     @property
     def next(self):
         """Sets the Register's value for the next cycle (it is before the D-Latch)."""

tests/test_wire.py

@@ -53,6 +53,35 @@ def test_rename(self):
         self.assertIn("testJohn", block.wirevector_by_name)
         self.assertIn(w, block.wirevector_set)
 
+    def test_src_connects_wire(self):
+        inp = pyrtl.Input(bitwidth=8, name="inp")
+        w = pyrtl.WireVector(bitwidth=8, name="w", src=inp)
+        out = pyrtl.Output(name="out", bitwidth=8)
+        out <<= w
+        sim = pyrtl.Simulation()
+        sim.step(provided_inputs={"inp": 42})
+        self.assertEqual(sim.inspect("out"), 42)
+
+    def test_src_infers_bitwidth(self):
+        inp = pyrtl.Input(bitwidth=8, name="inp")
+        w = pyrtl.WireVector(name="w", src=inp)
+        self.assertEqual(w.bitwidth, 8)
+
+    def test_src_with_int(self):
+        w = pyrtl.WireVector(bitwidth=8, name="w", src=42)
+        out = pyrtl.Output(name="out", bitwidth=8)
+        out <<= w
+        sim = pyrtl.Simulation()
+        sim.step(provided_inputs={})
+        self.assertEqual(sim.inspect("out"), 42)
+
+    def test_src_none_does_nothing(self):
+        w = pyrtl.WireVector(bitwidth=8, name="w", src=None)
+        # No net should be created for this wire
+        block = pyrtl.working_block()
+        src_dict, _ = block.net_connections(include_virtual_nodes=False)
+        self.assertNotIn(w, src_dict)
+
 
 class TestWireVectorNames(unittest.TestCase):
     def setUp(self):
@@ -341,6 +370,42 @@ def test_invalid_reset_value_not_an_integer(self):
             pyrtl.Register(4, reset_value="hello")
 
 
+class TestRegisterSrc(unittest.TestCase):
+    def setUp(self):
+        pyrtl.reset_working_block()
+
+    def test_register_src_sets_next(self):
+        inp = pyrtl.Input(bitwidth=3, name="inp")
+        r = pyrtl.Register(bitwidth=3, src=inp)
+        self.assertIsNotNone(r.reg_in)
+
+    def test_register_src_simulates(self):
+        inp = pyrtl.Input(bitwidth=3, name="inp")
+        r = pyrtl.Register(bitwidth=3, name="r", src=inp)
+        pyrtl.Output(name="out", bitwidth=3, src=r)
+        sim = pyrtl.Simulation()
+        sim.step(provided_inputs={"inp": 5})
+        self.assertEqual(sim.inspect("out"), 0)  # reset value on first step
+        sim.step(provided_inputs={"inp": 7})
+        self.assertEqual(sim.inspect("out"), 5)  # gets previous input
+
+    def test_register_src_with_reset_value(self):
+        inp = pyrtl.Input(bitwidth=3, name="inp")
+        r = pyrtl.Register(bitwidth=3, name="r", reset_value=3, src=inp)
+        pyrtl.Output(name="out", bitwidth=3, src=r)
+        sim = pyrtl.Simulation()
+        sim.step(provided_inputs={"inp": 5})
+        self.assertEqual(sim.inspect("out"), 3)  # reset value on first step
+        sim.step(provided_inputs={"inp": 7})
+        self.assertEqual(sim.inspect("out"), 5)
+
+    def test_register_src_double_assign_error(self):
+        inp = pyrtl.Input(bitwidth=3, name="inp")
+        r = pyrtl.Register(bitwidth=3, src=inp)
+        with self.assertRaises(pyrtl.PyrtlError):
+            r.next <<= inp
+
+
 class TestStateRegister(unittest.TestCase):
     def setUp(self):
         pyrtl.reset_working_block()
@@ -501,6 +566,36 @@ def test_slice_output(self):
         with self.assertRaises(pyrtl.PyrtlInternalError):
             _ = o[0]
 
+    def test_output_src_connects(self):
+        inp = pyrtl.Input(bitwidth=8, name="inp")
+        pyrtl.Output(name="out", src=inp)
+        sim = pyrtl.Simulation()
+        sim.step(provided_inputs={"inp": 99})
+        self.assertEqual(sim.inspect("out"), 99)
+
+    def test_output_src_infers_bitwidth(self):
+        inp = pyrtl.Input(bitwidth=8, name="inp")
+        out = pyrtl.Output(name="out", src=inp)
+        self.assertEqual(out.bitwidth, 8)
+
+    def test_output_src_with_expression(self):
+        a = pyrtl.Input(bitwidth=8, name="a")
+        b = pyrtl.Input(bitwidth=8, name="b")
+        pyrtl.Output(name="out", src=a + b)
+        sim = pyrtl.Simulation()
+        sim.step(provided_inputs={"a": 10, "b": 20})
+        self.assertEqual(sim.inspect("out"), 30)
+
+    def test_input_rejects_src(self):
+        w = pyrtl.WireVector(bitwidth=8)
+        with self.assertRaises(TypeError):
+            pyrtl.Input(bitwidth=8, name="a", src=w)
+
+    def test_const_rejects_src(self):
+        w = pyrtl.WireVector(bitwidth=8)
+        with self.assertRaises(TypeError):
+            pyrtl.Const(val=5, src=w)
+
 
 class TestKeepingCallStack(unittest.TestCase):
     def setUp(self):