From 1c7c75a254c35da5058eec913e55214931062dbc Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Mon, 24 Dec 2018 02:02:59 +0000
Subject: [PATCH] hdl.xfrm: implement SwitchCleaner, for pruning empty
 switches.

---
 nmigen/back/rtlil.py         |  3 ++-
 nmigen/hdl/xfrm.py           | 18 ++++++++++++++++--
 nmigen/test/test_hdl_xfrm.py | 27 +++++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 3 deletions(-)

diff --git a/nmigen/back/rtlil.py b/nmigen/back/rtlil.py
index 91bd1c1..aca7ff2 100644
--- a/nmigen/back/rtlil.py
+++ b/nmigen/back/rtlil.py
@@ -620,6 +620,7 @@ def convert_fragment(builder, fragment, name, top):
         rhs_compiler   = _RHSValueCompiler(compiler_state)
         lhs_compiler   = _LHSValueCompiler(compiler_state)
         stmt_compiler  = _StatementCompiler(compiler_state, rhs_compiler, lhs_compiler)
+        switch_cleaner = xfrm.SwitchCleaner()
 
         verilog_trigger = None
         verilog_trigger_sync_emitted = False
@@ -720,7 +721,7 @@ def convert_fragment(builder, fragment, name, top):
                     stmt_compiler._group = group_signals
                     stmt_compiler._case = case
                     stmt_compiler._has_rhs = False
-                    stmt_compiler(fragment.statements)
+                    stmt_compiler(switch_cleaner(fragment.statements))
 
                     # Verilog `always @*` blocks will not run if `*` does not match anythng, i.e.
                     # if the implicit sensitivity list is empty. We check this while translating,
diff --git a/nmigen/hdl/xfrm.py b/nmigen/hdl/xfrm.py
index 06e5c28..2b219b8 100644
--- a/nmigen/hdl/xfrm.py
+++ b/nmigen/hdl/xfrm.py
@@ -13,7 +13,7 @@ __all__ = ["ValueVisitor", "ValueTransformer",
            "StatementVisitor", "StatementTransformer",
            "FragmentTransformer",
            "DomainRenamer", "DomainLowerer",
-           "LHSGroupAnalyzer",
+           "SwitchCleaner", "LHSGroupAnalyzer",
            "ResetInserter", "CEInserter"]
 
 
@@ -165,7 +165,7 @@ class StatementTransformer(StatementVisitor):
         return Assign(self.on_value(stmt.lhs), self.on_value(stmt.rhs))
 
     def on_Switch(self, stmt):
-        cases = OrderedDict((k, self.on_statement(v)) for k, v in stmt.cases.items())
+        cases = OrderedDict((k, self.on_statement(s)) for k, s in stmt.cases.items())
         return Switch(self.on_value(stmt.test), cases)
 
     def on_statements(self, stmts):
@@ -280,6 +280,20 @@ class DomainLowerer(FragmentTransformer, ValueTransformer, StatementTransformer)
         return cd.rst
 
 
+class SwitchCleaner(StatementVisitor):
+    def on_Assign(self, stmt):
+        return stmt
+
+    def on_Switch(self, stmt):
+        cases = OrderedDict((k, self.on_statement(s)) for k, s in stmt.cases.items())
+        if any(len(s) for s in stmt.cases.values()):
+            return Switch(stmt.test, cases)
+
+    def on_statements(self, stmts):
+        stmts = flatten(self.on_statement(stmt) for stmt in stmts)
+        return _StatementList(stmt for stmt in stmts if stmt is not None)
+
+
 class LHSGroupAnalyzer(StatementVisitor):
     def __init__(self):
         self.signals = SignalDict()
diff --git a/nmigen/test/test_hdl_xfrm.py b/nmigen/test/test_hdl_xfrm.py
index d89db07..1c6d6c0 100644
--- a/nmigen/test/test_hdl_xfrm.py
+++ b/nmigen/test/test_hdl_xfrm.py
@@ -158,6 +158,33 @@ class DomainLowererTestCase(FHDLTestCase):
             DomainLowerer({"sync": sync})(f)
 
 
+class SwitchCleanerTestCase(FHDLTestCase):
+    def test_clean(self):
+        a = Signal()
+        b = Signal()
+        c = Signal()
+        stmts = [
+            Switch(a, {
+                1: a.eq(0),
+                0: [
+                    b.eq(1),
+                    Switch(b, {1: []})
+                ]
+            })
+        ]
+
+        self.assertRepr(SwitchCleaner()(stmts), """
+        (
+            (switch (sig a)
+                (case 1
+                    (eq (sig a) (const 1'd0)))
+                (case 0
+                    (eq (sig b) (const 1'd1)))
+            )
+        )
+        """)
+
+
 class LHSGroupAnalyzerTestCase(FHDLTestCase):
     def test_no_group_unrelated(self):
         a = Signal()
-- 
2.30.2