From 4a903a86e3379b93c113daf72efdff7cc0becff1 Mon Sep 17 00:00:00 2001
From: whitequark <cz@m-labs.hk>
Date: Thu, 27 Dec 2018 16:02:31 +0000
Subject: [PATCH] hdl.dsl: add support for fsm.ongoing().

---
 examples/fsm.py             |  8 ++++++--
 nmigen/hdl/dsl.py           | 31 ++++++++++++++++++++++---------
 nmigen/test/test_hdl_dsl.py | 31 +++++++++++++++++++++++++++++--
 3 files changed, 57 insertions(+), 13 deletions(-)

diff --git a/examples/fsm.py b/examples/fsm.py
index 2f7ed50..3659fca 100644
--- a/examples/fsm.py
+++ b/examples/fsm.py
@@ -24,7 +24,7 @@ class UARTReceiver:
             m.d.sync += ctr.eq(ctr - 1)
 
         bit = Signal(3)
-        with m.FSM():
+        with m.FSM() as fsm:
             with m.State("START"):
                 with m.If(~self.i):
                     m.next = "DATA"
@@ -46,12 +46,16 @@ class UARTReceiver:
                         m.next = "DONE"
                     with m.Else():
                         m.next = "ERROR"
+
             with m.State("DONE"):
                 m.d.comb += self.rdy.eq(1)
                 with m.If(self.ack):
                     m.next = "START"
+
+            m.d.comb += self.err.eq(fsm.ongoing("ERROR"))
             with m.State("ERROR"):
-                m.d.comb += self.err.eq(1)
+                pass
+
         return m.lower(platform)
 
 
diff --git a/nmigen/hdl/dsl.py b/nmigen/hdl/dsl.py
index 9144ad1..944a3ec 100644
--- a/nmigen/hdl/dsl.py
+++ b/nmigen/hdl/dsl.py
@@ -92,7 +92,16 @@ class _ModuleBuilderDomainSet:
         self._builder._add_domain(domain)
 
 
-_FSM = namedtuple("_FSM", ("state", "encoding", "decoding"))
+class FSM:
+    def __init__(self, state, encoding, decoding):
+        self.state    = state
+        self.encoding = encoding
+        self.decoding = decoding
+
+    def ongoing(self, name):
+        if name not in self.encoding:
+            self.encoding[name] = len(self.encoding)
+        return self.state == self.encoding[name]
 
 
 class Module(_ModuleBuilderRoot):
@@ -221,16 +230,18 @@ class Module(_ModuleBuilderRoot):
         fsm_data = self._set_ctrl("FSM", {
             "name":     name,
             "signal":   Signal(name="{}_state".format(name)),
+            "reset":    reset,
             "domain":   domain,
             "encoding": OrderedDict(),
+            "decoding": OrderedDict(),
             "states":   OrderedDict(),
         })
-        if reset is not None:
-            fsm_data["encoding"][reset] = 0
+        self._generated[name] = fsm = \
+            FSM(fsm_data["signal"], fsm_data["encoding"], fsm_data["decoding"])
         try:
             self._ctrl_context = "FSM"
             self.domain._depth += 1
-            yield
+            yield fsm
         finally:
             self.domain._depth -= 1
             self._ctrl_context = None
@@ -301,17 +312,19 @@ class Module(_ModuleBuilderRoot):
             self._statements.append(Switch(switch_test, switch_cases))
 
         if name == "FSM":
-            fsm_signal, fsm_encoding, fsm_states = data["signal"], data["encoding"], data["states"]
+            fsm_signal, fsm_reset, fsm_encoding, fsm_decoding, fsm_states = \
+                data["signal"], data["reset"], data["encoding"], data["decoding"], data["states"]
             fsm_signal.nbits = bits_for(len(fsm_encoding) - 1)
+            if fsm_reset is None:
+                fsm_signal.reset = fsm_encoding[next(iter(fsm_states))]
+            else:
+                fsm_signal.reset = fsm_encoding[fsm_reset]
             # The FSM is encoded such that the state with encoding 0 is always the reset state.
-            fsm_decoding = OrderedDict({n: s for s, n in fsm_encoding.items()})
+            fsm_decoding.update((n, s) for s, n in fsm_encoding.items())
             fsm_signal.decoder = lambda n: "{}/{}".format(fsm_decoding[n], n)
             self._statements.append(Switch(fsm_signal,
                 OrderedDict((fsm_encoding[name], stmts) for name, stmts in fsm_states.items())))
 
-            fsm_name = data["name"]
-            self._generated[fsm_name] = _FSM(fsm_signal, fsm_encoding, fsm_decoding)
-
     def _add_statement(self, assigns, domain, depth, compat_mode=False):
         def domain_name(domain):
             if domain is None:
diff --git a/nmigen/test/test_hdl_dsl.py b/nmigen/test/test_hdl_dsl.py
index 8b72410..43c9672 100644
--- a/nmigen/test/test_hdl_dsl.py
+++ b/nmigen/test/test_hdl_dsl.py
@@ -365,12 +365,39 @@ class DSLTestCase(FHDLTestCase):
         self.assertRepr(m._statements, """
         (
             (switch (sig fsm_state)
-                (case 1
+                (case 0
                     (eq (sig a) (const 1'd0))
+                    (eq (sig fsm_state) (const 1'd1))
+                )
+                (case 1
                     (eq (sig fsm_state) (const 1'd0))
                 )
+            )
+        )
+        """)
+
+    def test_FSM_ongoing(self):
+        a = Signal()
+        b = Signal()
+        m = Module()
+        with m.FSM() as fsm:
+            m.d.comb += b.eq(fsm.ongoing("SECOND"))
+            with m.State("FIRST"):
+                pass
+            m.d.comb += a.eq(fsm.ongoing("FIRST"))
+            with m.State("SECOND"):
+                pass
+        m._flush()
+        self.assertEqual(m._generated["fsm"].state.reset, 1)
+        self.maxDiff = 10000
+        self.assertRepr(m._statements, """
+        (
+            (eq (sig b) (== (sig fsm_state) (const 1'd0)))
+            (eq (sig a) (== (sig fsm_state) (const 1'd1)))
+            (switch (sig fsm_state)
+                (case 1
+                )
                 (case 0
-                    (eq (sig fsm_state) (const 1'd1))
                 )
             )
         )
-- 
2.30.2