From 33a446b402d730e8dab452761b6857f9e7ad85fe Mon Sep 17 00:00:00 2001
From: Luke Kenneth Casson Leighton <lkcl@lkcl.net>
Date: Mon, 29 Apr 2019 22:54:11 +0100
Subject: [PATCH] store ospecfn and ispecfn in stagehelper

---
 src/add/singlepipe.py | 10 +++++-----
 src/add/stageapi.py   | 21 +++++++++++++++++----
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/src/add/singlepipe.py b/src/add/singlepipe.py
index 3675ea90..c9faa510 100644
--- a/src/add/singlepipe.py
+++ b/src/add/singlepipe.py
@@ -198,13 +198,12 @@ class ControlBase(StageHelper, Elaboratable):
 
         # set up the input and output data
         if stage is not None:
-            self._new_data(self, self, "data")
+            self._new_data("data")
 
-    def _new_data(self, p, n, name):
+    def _new_data(self, name):
         """ allocates new data_i and data_o
         """
-        self.p.data_i = _spec(p.stage.ispec, "%s_i" % name)
-        self.n.data_o = _spec(n.stage.ospec, "%s_o" % name)
+        self.p.data_i, self.n.data_o = self.new_specs(name)
 
     @property
     def data_r(self):
@@ -280,7 +279,8 @@ class ControlBase(StageHelper, Elaboratable):
         # connect front and back of chain to ourselves
         front = pipechain[0]                # first in chain
         end = pipechain[-1]                 # last in chain
-        self._new_data(front, end, "chain") # NOTE: REPLACES existing data
+        self.set_specs(front, end) # NOTE: REPLACES existing data
+        self._new_data("chain") # NOTE: REPLACES existing data
         eqs += front._connect_in(self)      # front p to our p
         eqs += end._connect_out(self)       # end n   to out n
 
diff --git a/src/add/stageapi.py b/src/add/stageapi.py
index 7dc3659a..e819fefc 100644
--- a/src/add/stageapi.py
+++ b/src/add/stageapi.py
@@ -221,14 +221,27 @@ class StageHelper(Stage):
     """
     def __init__(self, stage):
         self.stage = stage
+        self._ispecfn = None
+        self._ospecfn = None
+        if stage is not None:
+            self.set_specs(self, self)
 
     def ospec(self, name):
-        assert self.stage is not None
-        return _spec(self.stage.ospec, name)
+        assert self._ospecfn is not None
+        return _spec(self._ospecfn, name)
 
     def ispec(self, name):
-        assert self.stage is not None
-        return _spec(self.stage.ispec, name)
+        assert self._ispecfn is not None
+        return _spec(self._ispecfn, name)
+
+    def set_specs(self, p, n):
+        self._ispecfn = p.stage.ispec
+        self._ospecfn = n.stage.ospec
+
+    def new_specs(self, name):
+        """ allocates new ispec and ospec pair
+        """
+        return self.ispec("%s_i" % name), self.ospec("%s_o" % name)
 
     def process(self, i):
         if self.stage and hasattr(self.stage, "process"):
-- 
2.30.2