From 934546e63360de082f3300665f9bf3f834efa4d5 Mon Sep 17 00:00:00 2001
From: whitequark <whitequark@whitequark.org>
Date: Wed, 26 Dec 2018 12:39:05 +0000
Subject: [PATCH] hdl.dsl: provide generated values for FSMs.

---
 nmigen/hdl/dsl.py           | 13 +++++++++++--
 nmigen/test/test_hdl_dsl.py | 14 ++++++++++++++
 2 files changed, 25 insertions(+), 2 deletions(-)

diff --git a/nmigen/hdl/dsl.py b/nmigen/hdl/dsl.py
index eec3f6e..6ed5c02 100644
--- a/nmigen/hdl/dsl.py
+++ b/nmigen/hdl/dsl.py
@@ -1,4 +1,4 @@
-from collections import OrderedDict
+from collections import OrderedDict, namedtuple
 from collections.abc import Iterable
 from contextlib import contextmanager
 
@@ -92,6 +92,9 @@ class _ModuleBuilderDomainSet:
         self._builder._add_domain(domain)
 
 
+_FSM = namedtuple("_FSM", ("state", "encoding", "decoding"))
+
+
 class Module(_ModuleBuilderRoot):
     def __init__(self):
         _ModuleBuilderRoot.__init__(self, self, depth=0)
@@ -105,6 +108,7 @@ class Module(_ModuleBuilderRoot):
         self._driving      = SignalDict()
         self._submodules   = []
         self._domains      = []
+        self._generated    = {}
 
     def _check_context(self, construct, context):
         if self._ctrl_context != context:
@@ -215,6 +219,7 @@ class Module(_ModuleBuilderRoot):
     def FSM(self, reset=None, domain="sync", name="fsm"):
         self._check_context("FSM", context=None)
         fsm_data = self._set_ctrl("FSM", {
+            "name":     name,
             "signal":   Signal(name="{}_state".format(name)),
             "domain":   domain,
             "encoding": OrderedDict(),
@@ -298,11 +303,14 @@ class Module(_ModuleBuilderRoot):
             fsm_signal, fsm_encoding, fsm_states = data["signal"], data["encoding"], data["states"]
             fsm_signal.nbits = bits_for(len(fsm_encoding) - 1)
             # The FSM is encoded such that the state with encoding 0 is always the reset state.
-            fsm_decoding = {n: s for s, n in fsm_encoding.items()}
+            fsm_decoding = OrderedDict({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:
@@ -354,6 +362,7 @@ class Module(_ModuleBuilderRoot):
         for signal, domain in self._driving.items():
             fragment.add_driver(signal, domain)
         fragment.add_domains(self._domains)
+        fragment.generated.update(self._generated)
         return fragment
 
     get_fragment = lower
diff --git a/nmigen/test/test_hdl_dsl.py b/nmigen/test/test_hdl_dsl.py
index 0a84b71..7087bc5 100644
--- a/nmigen/test/test_hdl_dsl.py
+++ b/nmigen/test/test_hdl_dsl.py
@@ -1,3 +1,5 @@
+from collections import OrderedDict
+
 from ..hdl.ast import *
 from ..hdl.cd import *
 from ..hdl.dsl import *
@@ -338,6 +340,18 @@ class DSLTestCase(FHDLTestCase):
             "(sig b)": "sync",
         })
 
+        frag = m.lower(platform=None)
+        fsm  = frag.find_generated("fsm")
+        self.assertIsInstance(fsm.state, Signal)
+        self.assertEqual(fsm.encoding, OrderedDict({
+            "FIRST": 0,
+            "SECOND": 1,
+        }))
+        self.assertEqual(fsm.decoding, OrderedDict({
+            0: "FIRST",
+            1: "SECOND"
+        }))
+
     def test_FSM_reset(self):
         a = Signal()
         m = Module()
-- 
2.30.2