From 345a725267235fa4d1695718a2c65fafd5bb12ad Mon Sep 17 00:00:00 2001
From: Luke Kenneth Casson Leighton <lkcl@lkcl.net>
Date: Sat, 30 Mar 2019 23:11:28 +0000
Subject: [PATCH] use and return module instead of m.d in pipeline
 introspection

---
 src/add/pipeline.py         | 14 +++++++-------
 src/add/pipeline_example.py | 10 +++++++---
 2 files changed, 14 insertions(+), 10 deletions(-)

diff --git a/src/add/pipeline.py b/src/add/pipeline.py
index 56e8e33c..c4afc9d8 100644
--- a/src/add/pipeline.py
+++ b/src/add/pipeline.py
@@ -74,8 +74,8 @@ class PipelineStage:
     """ Pipeline builder stage with auto generation of pipeline registers.
     """
 
-    def __init__(self, pipe, prev=None):
-        self._pipe = pipe
+    def __init__(self, m, prev=None):
+        self._m = m
         self._preg_map = {}
         self._prev_stage = prev
         if prev:
@@ -117,18 +117,18 @@ class PipelineStage:
         if next_stage not in self._preg_map:
             self._preg_map[next_stage] = {}
         self._preg_map[next_stage][name] = new_pipereg
-        self._pipe.sync += eq(new_pipereg, value)
+        self._m.d.sync += eq(new_pipereg, value)
 
 
 class PipeManager:
-    def __init__(self, pipe):
-        self._pipe = pipe
+    def __init__(self, m):
+        self.m = m
 
     @contextmanager
     def Stage(self, prev=None):
-        stage = PipelineStage(self._pipe, prev)
+        stage = PipelineStage(self.m, prev)
         try:
-            yield stage
+            yield stage, stage._m
         finally:
             pass
 
diff --git a/src/add/pipeline_example.py b/src/add/pipeline_example.py
index 36295265..e1806b30 100644
--- a/src/add/pipeline_example.py
+++ b/src/add/pipeline_example.py
@@ -88,7 +88,7 @@ class PipelineStageExample(PipeManager):
     def __init__(self):
         self.m = Module()
         self._loopback = Signal(4)
-        PipeManager.__init__(self, self.m.d)
+        PipeManager.__init__(self, self.m)
 
     def stage0(self):
         self.n = ~self._loopback
@@ -109,10 +109,14 @@ class PipelineStageExample(PipeManager):
 
     def get_fragment(self, platform=None):
 
-        with self.Stage() as p:
+        with self.Stage() as (p, m):
             p.n = ~self._loopback
-        with self.Stage(p) as p:
+        with self.Stage(p) as (p, m):
             p.n = p.n + 2
+        with self.Stage(p) as (p, m):
+            localv = Signal(4)
+            m.d.comb += localv.eq(2)
+            p.n = p.n << localv
 
         return self.m
 
-- 
2.30.2