Fix OpResult repr.

2 files changed, 16 insertions, 3 deletions

diff --git a/pyecsca/ec/formula/base.py b/pyecsca/ec/formula/base.py
index a9402f1..429f4c1 100644
--- a/pyecsca/ec/formula/base.py
+++ b/pyecsca/ec/formula/base.py
@@ -43,9 +43,13 @@ class OpResult:
         return self.name
 
     def __repr__(self):
-        # TODO: This repr is broken for square and neg and inv.
         char = self.op.op_str
-        parents = char.join(str(parent) for parent in self.parents)
+        if self.op == OpType.Inv:
+            parents = "1" + char + str(self.parents[0])
+        elif self.op == OpType.Neg:
+            parents = char + str(self.parents[0])
+        else:
+            parents = char.join(str(parent) for parent in self.parents)
         return f"{self.name} = {parents}"
 
 
diff --git a/test/ec/test_op.py b/test/ec/test_op.py
index 94edcc3..ece8cd9 100644
--- a/test/ec/test_op.py
+++ b/test/ec/test_op.py
@@ -1,6 +1,8 @@
 from ast import parse
 
 import pytest
+
+from pyecsca.ec.formula import OpResult
 from pyecsca.ec.mod import Mod
 from pyecsca.ec.op import CodeOp, OpType
 
@@ -27,4 +29,11 @@ def test_call(name, module, locals, result):
     res = op(**locals)
     assert res == result
 
-# TODO: Add op_type tests
+
+def test_opresult_repr():
+    res = OpResult("a", Mod(7, 11), OpType.Neg, "b")
+    assert repr(res) == "a = -b"
+    res = OpResult("a", Mod(5, 7), OpType.Add, "c", 3)
+    assert repr(res) == "a = c+3"
+    res = OpResult("a", Mod(3, 11), OpType.Inv, "d")
+    assert repr(res) == "a = 1/d"