assign and capture return arrays, new cycle method for common use pattern

src/kyupy/logic_sim.py

@@ -65,36 +65,32 @@ class LogicSim:
 
         :param stimuli: The input data to assign. Must be in bit-parallel storage format and in a compatible shape.
         :type stimuli: :py:class:`~kyupy.logic.BPArray`
+        :returns: The given stimuli object.
         """
-        if hasattr(stimuli, 'data'):
-            stimuli = stimuli.data
-        for stim, node in zip(stimuli, self.interface):
+        for node, stim in zip(self.interface, stimuli.data if hasattr(stimuli, 'data') else stimuli):
             if len(node.outs) == 0: continue
-            outputs = [self.state[line.index] if line else self.tmp[3] for line in node.outs]
-            self.node_fct[node.index]([stim], outputs)
+            outputs = [self.state[line] if line else self.tmp[3] for line in node.outs]
+            self.node_fct[node]([stim], outputs)
             for line in node.outs:
-                if line:
-                    self.state_epoch[line.reader.index] = self.epoch
+                if line is not None: self.state_epoch[line.reader] = self.epoch
         for n in self.circuit.nodes:
-            if (n.kind == '__const1__') or (n.kind == '__const0__'):
-                outputs = [self.state[line.index] if line else self.tmp[3] for line in n.outs]
-                self.node_fct[n.index]([], outputs)
-                # print('assign const')
+            if n.kind in ('__const1__', '__const0__'):
+                outputs = [self.state[line] if line else self.tmp[3] for line in n.outs]
+                self.node_fct[n]([], outputs)
                 for line in n.outs:
-                    if line:
-                        self.state_epoch[line.reader.index] = self.epoch
+                    if line is not None: self.state_epoch[line.reader] = self.epoch
+        return stimuli
 
     def capture(self, responses):
         """Capture the current values at the primary outputs and in the state-elements (flip-flops).
 
         :param responses: A bit-parallel storage target for the responses in a compatible shape.
         :type responses: :py:class:`~kyupy.logic.BPArray`
+        :returns: The given responses object.
         """
-        if hasattr(responses, 'data'):
-            responses = responses.data
-        for resp, node in zip(responses, self.interface):
-            if len(node.ins) == 0: continue
-            resp[...] = self.state[node.ins[0].index]
+        for node, resp in zip(self.interface, responses.data if hasattr(responses, 'data') else responses):
+            if len(node.ins) > 0: resp[...] = self.state[node.ins[0]]
+        return responses
 
     def propagate(self, inject_cb=None):
         """Propagate the input values towards the outputs (Perform all logic operations in topological order).
@@ -117,16 +113,30 @@ class LogicSim:
         :type inject_cb: ``f(int, ndarray)``
         """
         for node in self.circuit.topological_order():
-            if self.state_epoch[node.index] != self.epoch: continue
-            inputs = [self.state[line.index] if line else self.zero for line in node.ins]
-            outputs = [self.state[line.index] if line else self.tmp[3] for line in node.outs]
+            if self.state_epoch[node] != self.epoch: continue
+            inputs = [self.state[line] if line else self.zero for line in node.ins]
+            outputs = [self.state[line] if line else self.tmp[3] for line in node.outs]
             # print('sim', node)
-            self.node_fct[node.index](inputs, outputs)
+            self.node_fct[node](inputs, outputs)
             for line in node.outs:
-                if inject_cb is not None: inject_cb(line.index, self.state[line.index])
-                self.state_epoch[line.reader.index] = self.epoch
+                if inject_cb is not None: inject_cb(line, self.state[line])
+                self.state_epoch[line.reader] = self.epoch
         self.epoch = (self.epoch + 1) % 128
 
+    def cycle(self, state, inject_cb=None):
+        """Assigns the given state, propagates it and captures the new state.
+
+        :param responses: A bit-parallel array in a compatible shape holding the current circuit state.
+            The contained data is assigned to the PI and PPI and overwritten by data at the PO and PPO after
+            propagation.
+        :type responses: :py:class:`~kyupy.logic.BPArray`
+        :param inject_cb: A callback function for manipulating intermediate signal values. See :py:func:`propagate`.
+        :returns: The given state object.
+        """
+        self.assign(state)
+        self.propagate(inject_cb)
+        return self.capture(state)
+
     @staticmethod
     def fork_fct(inputs, outputs):
         for o in outputs: o[...] = inputs[0]