Update _NameSanitizer to try appending the current internal_index before generating an entirely new name. Fix a bug where a generated name could collide with an identifier name. Add tests.

Author: fdxmw

pyrtl/core.py

@@ -1193,13 +1193,39 @@ def next_index(self):
 
 
 class _NameSanitizer(_NameIndexer):
-    """Sanitizes the names so that names can be used in places that don't allow for
-    arbitrary names while not mangling valid names.
+    """Sanitizes names so they can be used in contexts that don't allow arbitrary names.
 
-    Put the values you want to validate into make_valid_string the first time you want
-    to sanitize a particular string (or before the first time), and retrieve from the
-    _NameSanitizer through indexing directly thereafter eg: sani["__&sfhs"] for
-    retrieval after the first time
+    For example, ``a.b`` is a valid ``WireVector`` name, but ``a.b`` is not a valid name
+    for a Python variable. If we want to generate Python code for ``a.b`` (like
+    ``FastSimulation``), the name must be sanitized.
+
+    Sanitization first attempts to replace non-word characters (anything that's not
+    alphanumeric or an underscore) with an underscore. If that didn't work, we try
+    appending a unique integer value. If that still doesn't work, we generate an
+    entirely new name consisting of ``internal_prefix`` followed by a unique integer
+    value.
+
+    ``make_valid_string`` must be called once to generate the sanitized version of a
+    name.
+
+    .. doctest only::
+
+        >>> import pyrtl
+
+    After ``make_valid_string`` has been called, the sanitized name can be retrieved
+    with ``__getitem__`` any number of times. For example::
+
+        >>> sanitizer = pyrtl.core._NameSanitizer(pyrtl.core._py_regex)
+
+        >>> sanitizer.make_valid_string("foo.bar")
+        'foo_bar'
+        >>> sanitizer["foo.bar"]
+        'foo_bar'
+        >>> sanitizer["foo.bar"]
+        'foo_bar'
+
+        >>> sanitizer.make_valid_string("foo_bar")
+        'foo_bar0'
     """
 
     def __init__(
@@ -1211,43 +1237,64 @@ def __init__(
         if identifier_regex_str[-1] != "$":
             identifier_regex_str += "$"
         self.identifier = re.compile(identifier_regex_str)
+        # Map from un-sanitized name to sanitized name.
         self.val_map = {}
+        # Set of all generated sanitized names.
+        self.sanitized_names = set()
         self.extra_checks = extra_checks
         super().__init__(internal_prefix)
 
-    def __getitem__(self, item):
-        """Get a value from the sanitizer"""
-        return self.val_map[item]
+    def __getitem__(self, name: str) -> str:
+        """Return the sanitized name for an un-sanitized name that was generated by
+        ``make_valid_string``.
+        """
+        return self.val_map[name]
 
-    def is_valid_str(self, string):
+    def is_valid_str(self, string: str) -> bool:
+        """Return ``True`` iff ``string`` matches ``identifier_regex_str`` and satisfies
+        ``extra_checks``.
+        """
         return self.identifier.match(string) and self.extra_checks(string)
 
-    def make_valid_string(self, string=""):
-        """Inputting a value for the first time."""
-        if self.is_valid_str(string):
-            self.val_map[string] = string
-            return string
-
+    def make_valid_string(self, string: str = "") -> str:
+        """Generate a sanitized name from an un-sanitized name."""
         if string in self.val_map:
             msg = f"Value {string} has already been given to the sanitizer"
             raise IndexError(msg)
-        # Try replacing non-word characters with ``_``.
-        internal_name = re.sub(r"\W", "_", string)
-
-        if not self.is_valid_str(internal_name) or internal_name in self.val_map:
-            # If that didn't work, try prepending ``_``.
-            internal_name = f"_{internal_name}"
-            while not self.is_valid_str(internal_name) or internal_name in self.val_map:
-                # If that didn't work, generate names starting with ``internal_prefix``
-                # until we find something acceptable.
-                internal_name = super().make_valid_string()
+
+        def is_usable(name: str) -> bool:
+            """Return ``True`` iff ``name`` can be used as a sanitized name.
+
+            A sanitized name is usable if it ``is_valid_str``, and isn't already in use.
+            """
+            return self.is_valid_str(name) and name not in self.sanitized_names
+
+        internal_name = string
+        if not is_usable(internal_name):
+            # Try replacing non-word characters with ``_``.
+            internal_name = re.sub(r"\W", "_", string)
+
+            if not is_usable(internal_name):
+                # If that didn't work, try appending the next ``internal_index``.
+                internal_name = f"{internal_name}{self.next_index()}"
+
+                if not is_usable(internal_name):
+                    # If that didn't work, generate an entirely new name starting with
+                    # ``internal_prefix``.
+                    internal_name = super().make_valid_string()
+
+                    if not is_usable(internal_name):
+                        msg = f"Could not generate a usable sanitized name for {string}"
+                        raise PyrtlError(msg)
+
         self.val_map[string] = internal_name
+        self.sanitized_names.add(internal_name)
         return internal_name
 
 
 class _PythonSanitizer(_NameSanitizer):
     """Name Sanitizer specifically built for Python identifers."""
 
-    def __init__(self, internal_prefix="_sani_temp", map_valid_vals=True):
-        super().__init__(_py_regex, internal_prefix, map_valid_vals)
+    def __init__(self, internal_prefix="_sani_temp"):
+        super().__init__(_py_regex, internal_prefix)
         self.extra_checks = lambda s: not keyword.iskeyword(s)

pyrtl/importexport.py

@@ -755,7 +755,7 @@ class _VerilogSanitizer(_NameSanitizer):
         unsigned use vectored wait wand weak0 weak1 while wire wor xnor xor
         """
 
-    def __init__(self, internal_prefix="_sani_temp"):
+    def __init__(self, internal_prefix):
         self._verilog_reserved_set = frozenset(self._verilog_reserved.split())
         super().__init__(self._ver_regex, internal_prefix, self._extra_checks)
 

tests/test_core.py

@@ -785,5 +785,37 @@ def test_wire_not_in_net_connections(self):
             dst_nets[w]
 
 
+class TestNameSanitizer(unittest.TestCase):
+    def test_name_collision(self):
+        # Test a sanitized name that collides with an unsanitized name.
+        sanitizer = pyrtl.core._NameSanitizer(
+            identifier_regex_str=pyrtl.core._py_regex,
+            internal_prefix="sanitized",
+        )
+        # This name should be sanitized by replacing the dot with an underscore.
+        sanitized_dot = sanitizer.make_valid_string("foo.bar")
+        self.assertEqual(sanitized_dot, "foo_bar")
+        self.assertEqual(sanitizer["foo.bar"], "foo_bar")
+
+        # This name collides with the sanitized name we just generated, so the first
+        # internal_index should be appended.
+        sanitized_underscore = sanitizer.make_valid_string("foo_bar")
+        self.assertEqual(sanitized_underscore, "foo_bar0")
+        self.assertEqual(sanitizer["foo_bar"], "foo_bar0")
+
+        # This name does not require sanitization, but it will collide with the next
+        # foo.bar variant.
+        sanitized_one = sanitizer.make_valid_string("foo_bar1")
+        self.assertEqual(sanitized_one, "foo_bar1")
+        self.assertEqual(sanitizer["foo_bar1"], "foo_bar1")
+
+        # Attempting to sanitize foo!bar by appending internal_index collides with the
+        # name we just registered. _NameSanitizer should give up and use
+        # internal_prefix.
+        sanitized_bang = sanitizer.make_valid_string("foo!bar")
+        self.assertEqual(sanitized_bang, "sanitized2")
+        self.assertEqual(sanitizer["foo!bar"], "sanitized2")
+
+
 if __name__ == "__main__":
     unittest.main()