5 files changed, 62 insertions, 35 deletions

diff --git a/.travis.yml b/.travis.yml
index a65f1a1..5ad51c1 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -13,6 +13,7 @@ addons:
     packages:
       - libps4000
       - libps6000
+      - swig
 
 before_install:
   - git clone https://github.com/colinoflynn/pico-python
diff --git a/pyecsca/ec/context.py b/pyecsca/ec/context.py
index b5ae389..8f797a8 100644
--- a/pyecsca/ec/context.py
+++ b/pyecsca/ec/context.py
@@ -8,7 +8,7 @@ from public import public
 
 from .formula import Formula
 from .mod import Mod
-from .op import CodeOp
+from .op import CodeOp, OpType
 from .point import Point
 
 
@@ -16,11 +16,11 @@ from .point import Point
 class OpResult(object):
     """A result of an operation."""
     parents: Tuple
-    op: Optional[ast.operator]
+    op: OpType
     name: str
     value: Mod
 
-    def __init__(self, name: str, value: Mod, op: Optional[ast.operator], *parents: Any):
+    def __init__(self, name: str, value: Mod, op: OpType, *parents: Any):
         self.parents = tuple(parents)
         self.name = name
         self.value = value
@@ -30,15 +30,7 @@ class OpResult(object):
         return self.name
 
     def __repr__(self):
-        char = ""
-        if isinstance(self.op, ast.Mult):
-            char = "*"
-        elif isinstance(self.op, ast.Add):
-            char = "+"
-        elif isinstance(self.op, ast.Sub):
-            char = "-"
-        elif isinstance(self.op, ast.Div):
-            char = "/"
+        char = self.op.op_str
         parents = char.join(str(parent) for parent in self.parents)
         return f"{self.name} = {parents}"
 
diff --git a/pyecsca/ec/curves.py b/pyecsca/ec/curves.py
index 496edf2..1bb2161 100644
--- a/pyecsca/ec/curves.py
+++ b/pyecsca/ec/curves.py
@@ -6,7 +6,7 @@ from .curve import EllipticCurve
 from .params import DomainParameters
 from .mod import Mod
 from .model import (ShortWeierstrassModel, MontgomeryModel, TwistedEdwardsModel,
-                              EdwardsModel)
+                    EdwardsModel, CurveModel)
 from .point import Point, InfinityPoint
 
 
@@ -199,6 +199,7 @@ def get_curve(name: str, coords: str) -> DomainParameters:
     :param coords: The name of the coordinate system to use.
     :return: The curve.
     """
+    model: CurveModel
     if name in SHORT_WEIERSTRASS:
         params = SHORT_WEIERSTRASS[name]
         model = ShortWeierstrassModel()
diff --git a/pyecsca/ec/op.py b/pyecsca/ec/op.py
index 72b06f3..3bf502e 100644
--- a/pyecsca/ec/op.py
+++ b/pyecsca/ec/op.py
@@ -1,16 +1,37 @@
-from ast import Module, walk, Name, BinOp, Constant, operator, Mult, Div, Add, Sub, Pow, Assign
+from ast import (Module, walk, Name, BinOp, Constant, Mult, Div, Add, Sub, Pow, Assign,
+                 operator as ast_operator)
+from enum import Enum
 from types import CodeType
-from typing import FrozenSet, Optional, cast
+from typing import FrozenSet, cast, Any, Optional
+
+from public import public
 
 from .mod import Mod
 
 
+@public
+class OpType(Enum):
+    Add = (2, "+")
+    Sub = (2, "-")
+    Mult = (2, "*")
+    Div = (2, "/")
+    Inv = (1, "/")
+    Sqr = (1, "^")
+    Pow = (2, "^")
+    Id = (1, "")
+
+    def __init__(self, num_inputs: int, op_str: str):
+        self.num_inputs = num_inputs
+        self.op_str = op_str
+
+
+@public
 class CodeOp(object):
     result: str
     parameters: FrozenSet[str]
     variables: FrozenSet[str]
     code: Module
-    operator: Optional[operator]
+    operator: OpType
     compiled: CodeType
 
     def __init__(self, code: Module):
@@ -21,6 +42,8 @@ class CodeOp(object):
         variables = set()
         constants = set()
         op = None
+        self.left = None
+        self.right = None
         for node in walk(assign.value):
             if isinstance(node, Name):
                 name = node.id
@@ -34,10 +57,11 @@ class CodeOp(object):
                 op = node.op
                 self.left = self.__to_name(node.left)
                 self.right = self.__to_name(node.right)
-        if op is None and len(constants) == 1:
-            self.left = next(iter(constants))
-            self.right = None
-        self.operator = op
+        if isinstance(assign.value, Name):
+            self.left = assign.value.id
+        elif isinstance(assign.value, Constant):
+            self.left = assign.value.value
+        self.operator = self.__to_op(op, self.left, self.right)
         self.parameters = frozenset(params)
         self.variables = frozenset(variables)
         self.constants = frozenset(constants)
@@ -51,21 +75,25 @@ class CodeOp(object):
         else:
             return None
 
-    def __to_opsymbol(self, op):
+    def __to_op(self, op: Optional[ast_operator], left: Any, right: Any) -> OpType:
         if isinstance(op, Mult):
-            return "*"
+            return OpType.Mult
         elif isinstance(op, Div):
-            return "/"
+            if left == 1:
+                return OpType.Inv
+            return OpType.Div
         elif isinstance(op, Add):
-            return "+"
+            return OpType.Add
         elif isinstance(op, Sub):
-            return "-"
+            return OpType.Sub
         elif isinstance(op, Pow):
-            return "^"
-        return ""
+            if right == 2:
+                return OpType.Sqr
+            return OpType.Pow
+        return OpType.Id
 
     def __str__(self):
-        return f"{self.result} = {self.left}{self.__to_opsymbol(self.operator)}{self.right}"
+        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 ''}"
 
     def __repr__(self):
         return f"CodeOp({self.result} = f(params={self.parameters}, vars={self.variables}, consts={self.constants}))"
diff --git a/test/ec/test_op.py b/test/ec/test_op.py
index 13efc53..a7f9cd0 100644
--- a/test/ec/test_op.py
+++ b/test/ec/test_op.py
@@ -4,22 +4,27 @@ from unittest import TestCase
 from parameterized import parameterized
 
 from pyecsca.ec.mod import Mod
-from pyecsca.ec.op import CodeOp
+from pyecsca.ec.op import CodeOp, OpType
 
 
 class OpTests(TestCase):
 
     @parameterized.expand([
-        ("add", "x = a+b", "x = a+b"),
-        ("sub", "x = a-b", "x = a-b"),
-        ("mul", "y = a*b", "y = a*b"),
-        ("div", "z = a/b", "z = a/b"),
-        ("pow", "b = a**d", "b = a^d")
+        ("add", "x = a+b", "x = a+b", OpType.Add),
+        ("sub", "x = a-b", "x = a-b", OpType.Sub),
+        ("mul", "y = a*b", "y = a*b", OpType.Mult),
+        ("div", "z = a/b", "z = a/b", OpType.Div),
+        ("inv", "z = 1/b", "z = 1/b", OpType.Inv),
+        ("pow", "b = a**d", "b = a^d", OpType.Pow),
+        ("sqr", "b = a**2", "b = a^2", OpType.Sqr),
+        ("id1", "b = 7", "b = 7", OpType.Id),
+        ("id2", "b = a", "b = a", OpType.Id),
     ])
-    def test_str(self, name, module, result):
+    def test_str(self, name, module, result, op_type):
         code = parse(module, mode="exec")
         op = CodeOp(code)
         self.assertEqual(str(op), result)
+        self.assertEqual(op.operator, op_type)
 
     @parameterized.expand([
         ("add", "x = a+b", {"a": Mod(5, 21), "b": Mod(7, 21)}, Mod(12, 21)),