use and return module instead of m.d in pipeline introspection

diff --git a/src/add/pipeline.py b/src/add/pipeline.py
index 56e8e33cbfa4a0c9b0a3847fc740c8469b089608..c4afc9d862f7722bb3855da4b5d1e7585110fb0f 100644 (file)
--- 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 3629526538adcb8c443cbf2131377cf0f4014d03..e1806b302b5f81903c8db4de837cf3b6c4992637 100644 (file)
--- 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