Updating clone_wire to be safe when used on same block.

Author: mdko

pyrtl/transform.py

@@ -21,7 +21,7 @@
 import functools
 
 from .core import set_working_block, LogicNet, working_block
-from .wire import Const, Input, Output, WireVector, Register
+from .wire import Const, Input, Output, WireVector, Register, next_tempvar_name
 
 
 def net_transform(transform_func, block=None, **kwargs):
@@ -242,17 +242,18 @@ def clone_wire(old_wire, name=None):
     :param old_wire: The wire to clone
     :param name: A name for the new wire
 
-    Note that this function is mainly intended to be used when the
-    two wires are from different blocks. Making two wires with the
-    same name in the same block is not allowed.
+    Naming the newly cloned wire the same as another existing wire in the same
+    block will cause the other wire to be given a new internally created name. This
+    function is mainly intended to be used when the two wires are from different blocks.
     """
+    name = old_wire.name if name is None else name
+    if name in working_block().wirevector_by_name and (working_block() is old_wire._block):
+        w = working_block().wirevector_by_name[name]
+        w.name = next_tempvar_name()
+
     if isinstance(old_wire, Const):
-        if name is None:
-            return Const(old_wire.val, old_wire.bitwidth, name=old_wire.name)
         return Const(old_wire.val, old_wire.bitwidth, name=name)
     else:
-        if name is None:
-            return old_wire.__class__(old_wire.bitwidth, name=old_wire.name)
         return old_wire.__class__(old_wire.bitwidth, name=name)
 
 

tests/test_transform.py

@@ -1,3 +1,4 @@
+from pyrtl.core import set_working_block
 import unittest
 import pyrtl
 from pyrtl import transform
@@ -257,6 +258,70 @@ def test_replace_only_dst_wire(self):
         self.assertEqual(w4_dst_net.args, (w4, w2))
         self.assertEqual(w4_dst_net.dests, w1_dst_net.dests)
 
+class TestCloning(unittest.TestCase):
+    def setUp(self):
+        pyrtl.reset_working_block()
+
+    def test_clone_wire_no_or_same_name_same_block(self):
+        for clone_name in (None, 'a'):
+            a = pyrtl.WireVector(1, 'a')
+            self.assertEqual(a.name, 'a')
+            self.assertEqual(pyrtl.working_block().wirevector_set, {a})
+            self.assertIs(pyrtl.working_block().wirevector_by_name['a'], a)
+
+            w = pyrtl.clone_wire(a, name=clone_name)
+            self.assertTrue(a.name.startswith("tmp"))
+            self.assertEqual(w.name, 'a')
+            self.assertIs(pyrtl.working_block().wirevector_by_name['a'], w)
+            self.assertEqual(pyrtl.working_block().wirevector_set, {a, w})
+
+            w.name = 'w'
+            self.assertEqual(w.name, 'w')
+            self.assertIs(pyrtl.working_block().wirevector_by_name['w'], w)
+
+            pyrtl.working_block().remove_wirevector(a)
+            self.assertEqual(pyrtl.working_block().wirevector_set, {w})
+            pyrtl.reset_working_block()
+
+    def test_clone_wire_no_or_same_name_different_block(self):
+        for clone_name in (None, 'a'):
+            a = pyrtl.WireVector(1, 'a')
+            b = pyrtl.Block()
+            with pyrtl.set_working_block(b):
+                w = pyrtl.clone_wire(a, name=clone_name)
+
+            self.assertEqual(a.name, 'a')
+            self.assertIs(pyrtl.working_block().wirevector_by_name['a'], a)
+            self.assertEqual(pyrtl.working_block().wirevector_set, {a})
+
+            self.assertEqual(w.name, 'a')
+            self.assertIs(b.wirevector_by_name['a'], w)
+            self.assertEqual(b.wirevector_set, {w})
+            pyrtl.reset_working_block()
+
+    def test_clone_wire_with_different_name_same_block(self):
+        a = pyrtl.WireVector(1, 'a')
+        w = pyrtl.clone_wire(a, 'w')
+        self.assertEqual(a.name, 'a')
+        self.assertEqual(w.name, 'w')
+        self.assertIs(pyrtl.working_block().wirevector_by_name['a'], a)
+        self.assertEqual(pyrtl.working_block().wirevector_set, {a, w})
+
+        pyrtl.working_block().remove_wirevector(a)
+        self.assertEqual(pyrtl.working_block().wirevector_set, {w})
+
+    def test_clone_wire_with_different_name_different_block(self):
+        a = pyrtl.WireVector(1, 'a')
+        b = pyrtl.Block()
+        with set_working_block(b):
+            w = pyrtl.clone_wire(a, 'w')
+        self.assertEqual(a.name, 'a')
+        self.assertEqual(w.name, 'w')
+        self.assertIs(pyrtl.working_block().wirevector_by_name['a'], a)
+        self.assertEqual(pyrtl.working_block().wirevector_set, {a})
+        self.assertIs(b.wirevector_by_name['w'], w)
+        self.assertEqual(b.wirevector_set, {w})
+
 
 # this code needs mocking from python 3's unittests to work
 """