From 65864f0fa481b2716239068dc01163857b9a1e67 Mon Sep 17 00:00:00 2001
From: whitequark <cz@m-labs.hk>
Date: Fri, 11 Oct 2019 11:16:00 +0000
Subject: [PATCH] hdl.ast: simplify enum handling.

---
 nmigen/hdl/ast.py           | 23 +++++++++--------------
 nmigen/test/test_hdl_ast.py |  2 +-
 2 files changed, 10 insertions(+), 15 deletions(-)

diff --git a/nmigen/hdl/ast.py b/nmigen/hdl/ast.py
index 64e7082..db1a70b 100644
--- a/nmigen/hdl/ast.py
+++ b/nmigen/hdl/ast.py
@@ -40,11 +40,6 @@ def _enum_shape(enum_type):
     return (width, signed)
 
 
-def _enum_to_bits(enum_value):
-    width, signed = _enum_shape(type(enum_value))
-    return format(enum_value.value & ((1 << width) - 1), "b").rjust(width, "0")
-
-
 class Value(metaclass=ABCMeta):
     @staticmethod
     def cast(obj):
@@ -300,14 +295,14 @@ class Value(metaclass=ABCMeta):
                               .format(pattern, len(self)),
                               SyntaxWarning, stacklevel=3)
                 continue
-            if isinstance(pattern, int):
-                matches.append(self == pattern)
-            elif isinstance(pattern, (str, Enum)):
-                if isinstance(pattern, Enum):
-                    pattern = _enum_to_bits(pattern)
+            if isinstance(pattern, str):
                 mask    = int(pattern.replace("0", "1").replace("-", "0"), 2)
                 pattern = int(pattern.replace("-", "0"), 2)
                 matches.append((self & mask) == pattern)
+            elif isinstance(pattern, int):
+                matches.append(self == pattern)
+            elif isinstance(pattern, Enum):
+                matches.append(self == pattern.value)
             else:
                 assert False
         if not matches:
@@ -1304,12 +1299,12 @@ class Switch(Statement):
             # Map: 2 -> "0010"; "0010" -> "0010"
             new_keys = ()
             for key in keys:
-                if isinstance(key, (bool, int)):
-                    key = "{:0{}b}".format(key, len(self.test))
-                elif isinstance(key, str):
+                if isinstance(key, str):
                     pass
+                elif isinstance(key, int):
+                    key = format(key, "b").rjust(len(self.test), "0")
                 elif isinstance(key, Enum):
-                    key = _enum_to_bits(key)
+                    key = format(key.value, "b").rjust(len(self.test), "0")
                 else:
                     raise TypeError("Object '{!r}' cannot be used as a switch key"
                                     .format(key))
diff --git a/nmigen/test/test_hdl_ast.py b/nmigen/test/test_hdl_ast.py
index 362e2a8..1a2bb17 100644
--- a/nmigen/test/test_hdl_ast.py
+++ b/nmigen/test/test_hdl_ast.py
@@ -346,7 +346,7 @@ class OperatorTestCase(FHDLTestCase):
     def test_matches_enum(self):
         s = Signal.enum(SignedEnum)
         self.assertRepr(s.matches(SignedEnum.FOO), """
-        (== (& (sig s) (const 2'd3)) (const 2'd3))
+        (== (sig s) (const 1'sd-1))
         """)
 
     def test_matches_width_wrong(self):
-- 
2.30.2