From c04d21aa31bef2971e50ad36fa8c626bbff34f99 Mon Sep 17 00:00:00 2001
From: Donald Sebastian Leung <dl@m-labs.hk>
Date: Tue, 27 Oct 2020 13:33:41 +0800
Subject: [PATCH] Restructure rtio.sed.output_driver to follow nMigen
 convention

---
 README.md                 |  4 ++--
 rtio/sed/output_driver.py | 36 +++++++++++++++++++++++++-----------
 2 files changed, 27 insertions(+), 13 deletions(-)

diff --git a/README.md b/README.md
index a4092cb..1399e1f 100644
--- a/README.md
+++ b/README.md
@@ -30,10 +30,10 @@ $ python -m rtio.test.sed.output_network
 - - [ ] `rtio.cri` (`Interface` and `CRIDecoder` only)
 - - [x] `rtio.sed.output_network` - Sorting network (high priority)
 - - [ ] `rtio.sed.output_driver`
-- [ ] Restructure to code to follow nMigen convention and re-validate existing assertions (if any)
+- [x] Restructure to code to follow nMigen convention and re-validate existing assertions (if any)
 - - [x] `rtio.cri` (`Interface` and `CRIDecoder` only)
 - - [x] `rtio.sed.output_network`
-- - [ ] `rtio.sed.output_driver`
+- - [x] `rtio.sed.output_driver`
 
 ## License
 
diff --git a/rtio/sed/output_driver.py b/rtio/sed/output_driver.py
index 935edd3..5d5a687 100644
--- a/rtio/sed/output_driver.py
+++ b/rtio/sed/output_driver.py
@@ -10,8 +10,11 @@ __all__ = ["OutputDriver"]
 
 class OutputDriver(Elaboratable):
     def __init__(self, channels, glbl_fine_ts_width, lane_count, seqn_width):
-        m = Module()
-        self.m = m
+        self.channels = channels
+        self.glbl_fine_ts_width = glbl_fine_ts_width
+        self.lane_count = lane_count
+        self.seqn_width = seqn_width
+
         self.collision = Signal()
         self.collision_channel = Signal(range(len(channels)), reset_less=True)
         self.busy = Signal()
@@ -19,9 +22,8 @@ class OutputDriver(Elaboratable):
 
         # output network
         layout_on_payload = layouts.output_network_payload(channels, glbl_fine_ts_width)
-        output_network = OutputNetwork(lane_count, seqn_width, layout_on_payload)
-        m.submodules += output_network
-        self.input = output_network.input
+        self.output_network = OutputNetwork(lane_count, seqn_width, layout_on_payload)
+        self.input = self.output_network.input
 
         # detect collisions (adds one pipeline stage)
         layout_lane_data = [
@@ -29,8 +31,22 @@ class OutputDriver(Elaboratable):
             ("collision", 1),
             ("payload", layout_on_payload)
         ]
-        lane_datas = [Record(layout_lane_data, reset_less=True) for _ in range(lane_count)]
-        en_replaces = [channel.interface.o.enable_replace for channel in channels]
+        self.lane_datas = [Record(layout_lane_data) for _ in range(lane_count)]
+        for lane_data in self.lane_datas:
+            lane_data.valid.reset_less = True
+            lane_data.collision.reset_less = True
+            lane_data.payload.reset_less = True
+        self.en_replaces = [channel.interface.o.enable_replace for channel in channels]
+
+    def elaborate(self, platform):
+        m = Module()
+
+        channels = self.channels
+        output_network = self.output_network
+        lane_datas = self.lane_datas
+        en_replaces = self.en_replaces
+
+        m.submodules += self.output_network
         for lane_data, on_output in zip(lane_datas, output_network.output):
             lane_data.valid.reset_less = False
             lane_data.collision.reset_less = False
@@ -67,7 +83,7 @@ class OutputDriver(Elaboratable):
                 m.d.comb += selected.eq(lane_data.valid & ~lane_data.collision & (lane_data.payload.channel == n))
                 onehot_stb.append(selected)
                 if hasattr(lane_data.payload, "fine_ts") and hasattr(oif, "fine_ts"):
-                    ts_shift = len(lane_data.ayload.fine_ts) - len(oif.fine_ts)
+                    ts_shift = len(lane_data.payload.fine_ts) - len(oif.fine_ts)
                     onehot_fine_ts.append(Mux(selected, lane_data.payload.fine_ts[ts_shift:], 0))
                 if hasattr(lane_data.payload, "address"):
                     onehot_address.append(Mux(selected, lane_datapayload.address, 0))
@@ -95,6 +111,4 @@ class OutputDriver(Elaboratable):
                 m.d.sync += self.busy.eq(1)
                 m.d.sync += self.busy_channel.eq(channel_r)
 
-
-    def elaborate(self, platform):
-        return self.m
+        return m