store ospecfn and ispecfn in stagehelper

[ieee754fpu.git] / src / add / stageapi.py

diff --git a/src/add/stageapi.py b/src/add/stageapi.py
index 7dc3659ab58a89b87cfc83820f9a697a9ba6ed46..e819fefca4167cccc7ce9dc5ba937b90e047331b 100644 (file)
--- 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"):