Process OpResults better.

2 files changed, 24 insertions, 8 deletions

diff --git a/pyecsca/ec/formula.py b/pyecsca/ec/formula.py
index 226734d..317788b 100644
--- a/pyecsca/ec/formula.py
+++ b/pyecsca/ec/formula.py
@@ -29,6 +29,8 @@ class OpResult:
     value: Mod
 
     def __init__(self, name: str, value: Mod, op: OpType, *parents: Any):
+        if len(parents) != op.num_inputs:
+            raise ValueError(f"Wrong number of parents ({len(parents)}) to OpResult: {op} ({op.num_inputs}).")
         self.parents = tuple(parents)
         self.name = name
         self.value = value
@@ -56,7 +58,7 @@ class FormulaAction(ResultAction):
     intermediates: MutableMapping[str, List[OpResult]]
     """Intermediates computed during execution."""
     op_results: List[OpResult]
-    """"""
+    """The intermediates but ordered as they were computed."""
     outputs: MutableMapping[str, OpResult]
     """The output variables."""
     output_points: List[Any]
@@ -73,12 +75,15 @@ class FormulaAction(ResultAction):
         self.output_points = []
 
     def add_operation(self, op: CodeOp, value: Mod):
-        parents: List[Union[Mod, OpResult]] = []
-        for parent in {*op.variables, *op.parameters}:
-            if parent in self.intermediates:
-                parents.append(self.intermediates[parent][-1])
-            elif parent in self.inputs:
-                parents.append(self.inputs[parent])
+        parents: List[Union[int, Mod, OpResult]] = []
+        for parent in op.parents:
+            if isinstance(parent, str):
+                if parent in self.intermediates:
+                    parents.append(self.intermediates[parent][-1])
+                elif parent in self.inputs:
+                    parents.append(self.inputs[parent])
+            else:
+                parents.append(parent)
         result = OpResult(op.result, value, op.operator, *parents)
         li = self.intermediates.setdefault(op.result, [])
         li.append(result)
diff --git a/pyecsca/ec/op.py b/pyecsca/ec/op.py
index d656bc3..b83562c 100644
--- a/pyecsca/ec/op.py
+++ b/pyecsca/ec/op.py
@@ -18,7 +18,7 @@ from ast import (
 )
 from enum import Enum
 from types import CodeType
-from typing import FrozenSet, cast, Any, Optional, Union
+from typing import FrozenSet, cast, Any, Optional, Union, Tuple
 
 from public import public
 
@@ -54,6 +54,8 @@ class CodeOp:
     """The parameters used in the operation (e.g. `a`, `b`)."""
     variables: FrozenSet[str]
     """The variables used in the operation (e.g. `X1`, `Z2`)."""
+    constants: FrozenSet[int]  # TODO: Might not be only int? See issue in Formula eval.
+    """The constants used in the operation."""
     code: Module
     """The code of the operation."""
     operator: OpType
@@ -126,6 +128,15 @@ class CodeOp:
             return OpType.Pow
         return OpType.Id
 
+    @property
+    def parents(self) -> Tuple[Union[str, int]]:
+        if self.operator == OpType.Inv or self.operator == OpType.Neg:
+            return self.right,  # type: ignore
+        elif self.operator == OpType.Sqr or self.operator == OpType.Id:
+            return self.left,  # type: ignore
+        else:
+            return self.left, self.right  # type: ignore
+
     def __str__(self):
         return f"{self.result} = {self.left if self.left is not None else ''}{self.operator.op_str}{self.right if self.right is not None else ''}"