Split to_affine map to factor_set computation, fix mypy.

3 files changed, 170 insertions, 87 deletions

diff --git a/pyecsca/ec/coordinates.py b/pyecsca/ec/coordinates.py
index 25c4d09..9b63c75 100644
--- a/pyecsca/ec/coordinates.py
+++ b/pyecsca/ec/coordinates.py
@@ -37,6 +37,8 @@ class CoordinateModel:
     """Map to affine coordinates from system coordinates."""
     tosystem: List[Module]
     """Map from coordinate system to affine coordinates."""
+    homogweights: MutableMapping[str, int]
+    """Weights that homogenize the coordinates."""
     parameters: List[str]
     """Coordinate system parameters."""
     assumptions: List[Module]
diff --git a/pyecsca/sca/re/zvp.py b/pyecsca/sca/re/zvp.py
index 8aae076..0f04fe5 100644
--- a/pyecsca/sca/re/zvp.py
+++ b/pyecsca/sca/re/zvp.py
@@ -6,7 +6,7 @@ Provides functionality inspired by the Zero-value point attack.
 
 Implements ZVP point construction from [FFD]_.
 """
-from typing import List, Set, Tuple
+from typing import List, Set, Tuple, Dict
 from public import public
 from astunparse import unparse
 
@@ -20,15 +20,11 @@ from ...ec.point import Point
 
 
 @public
-def unroll_formula(formula: Formula, affine: bool = False) -> List[Tuple[str, Poly]]:
+def unroll_formula(formula: Formula) -> List[Tuple[str, Poly]]:
     """
     Unroll a given formula symbolically to obtain symbolic expressions for its intermediate values.
 
-    If :paramref:`~.compute_factor_set.affine` is set, the polynomials are transformed
-    to affine form, using some assumptions along the way (e.g. `Z = 1`).
-
     :param formula: Formula to unroll.
-    :param affine: Whether to transform the unrolled polynomials (and thus the resulting factors) into affine form.
     :return: List of symbolic intermediate values, with associated variable names.
     """
     params = {
@@ -48,50 +44,67 @@ def unroll_formula(formula: Formula, affine: bool = False) -> List[Tuple[str, Po
             for curve_param, value in params.items():
                 expr = expr.subs(curve_param, value)
             params[lhs] = expr
-    subs_map = {}
-    if affine:
-        # tosystem_map is the mapping of system variables (without indices) in affine variables (without indices)
-        tosystem_map = {}
-        for code in formula.coordinate_model.tosystem:
-            un = unparse(code).strip()
-            lhs, rhs = un.split(" = ")
-            tosystem_map[lhs] = sympify(rhs, evaluate=False)
-        # subs_map specializes the tosystem_map by adding appropriate indices
-        for i in range(1, formula.num_inputs + 1):
-            for lhs, rhs in tosystem_map.items():
-                subs_lhs = lhs + str(i)
-                subs_rhs = rhs.subs("x", f"x{i}").subs("y", f"y{i}")
-                subs_map[subs_lhs] = subs_rhs
 
     locls = {**params, **inputs}
-    values = []
+    values: List[Tuple[str, Poly]] = []
     for op in formula.code:
-        result = op(**locls)
+        result: Expr = op(**locls)  # type: ignore
         locls[op.result] = result
-        values.append((op.result, result))
+        if result.free_symbols:
+            gens = list(result.free_symbols)
+            gens.sort(key=str)
+            poly = Poly(result, *gens)
+            values.append((op.result, poly))
+        else:
+            # TODO: We cannot create a Poly here, because the result does not have free symbols (i.e. it is a constant)
+            pass
+
+    return values
+
+
+@public
+def map_to_affine(formula: Formula, polys: List[Tuple[str, Poly]]) -> List[Tuple[str, Poly]]:
+    """
+    Map unrolled polynomials of a formula to affine form, using some assumptions along the way (e.g. `Z = 1`).
 
-    values = filter_out_nonhomogenous_polynomials(formula, values)
+    :param formula: The formula the polynomials belong to.
+    :param polys: The polynomials (intermediate values) to map.
+    :return: The mapped intermediate values, with associated variable names.
+    """
+    # tosystem_map is the mapping of system variables (without indices) in affine variables (without indices)
+    tosystem_map = {}
+    for code in formula.coordinate_model.tosystem:
+        un = unparse(code).strip()
+        lhs, rhs = un.split(" = ")
+        tosystem_map[lhs] = sympify(rhs, evaluate=False)
+    subs_map = {}
+    # subs_map specializes the tosystem_map by adding appropriate indices
+    for i in range(1, formula.num_inputs + 1):
+        for lhs, rhs in tosystem_map.items():
+            subs_lhs = lhs + str(i)
+            subs_rhs = rhs.subs("x", f"x{i}").subs("y", f"y{i}")
+            subs_map[subs_lhs] = subs_rhs
 
     results = []
-    for result_var, value in values:
-        if affine:
-            expr = value
-            for lhs, rhs in subs_map.items():
-                expr = expr.subs(lhs, rhs)
-            if expr.free_symbols:
-                gens = list(expr.free_symbols)
-                gens.sort(key=str)
-                poly = Poly(expr, *gens)
-                results.append((result_var, poly))
-            else:
-                # Skip if no variables remain (constant poly)
-                continue
+    for result_var, value in polys:
+        expr = value
+        for lhs, rhs in subs_map.items():
+            expr = expr.subs(lhs, rhs)
+        if expr.free_symbols:
+            gens = list(expr.free_symbols)
+            gens.sort(key=str)
+            poly = Poly(expr, *gens)
+            results.append((result_var, poly))
         else:
-            results.append((result_var, Poly(value)))
+            # TODO: We cannot create a Poly here, because the result does not have free symbols (i.e. it is a constant)
+            #       Though here we do not care.
+            pass
     return results
 
 
-def filter_out_nonhomogenous_polynomials(formula: Formula, unrolled: List[Tuple[str, Poly]]) -> List[Tuple[str, Poly]]:
+def filter_out_nonhomogenous_polynomials(
+    formula: Formula, unrolled: List[Tuple[str, Poly]]
+) -> List[Tuple[str, Poly]]:
     """
     Remove unrolled polynomials from unrolled formula that are not homogenous.
 
@@ -102,12 +115,12 @@ def filter_out_nonhomogenous_polynomials(formula: Formula, unrolled: List[Tuple[
     if "mmadd" in formula.name:
         return unrolled
     homogenity_weights = formula.coordinate_model.homogweights
-    
+
     # we have to group variables by points and check homogenity for each group
-    input_variables_grouped = {}
+    input_variables_grouped: Dict[int, List[str]] = {}
     for var in formula.inputs:
         # here we assume that the index of the variable is <10 and on the last position
-        group = input_variables_grouped.setdefault(var[-1],[])
+        group = input_variables_grouped.setdefault(int(var[-1]), [])
         group.append(var)
 
     # zadd formulas have Z1=Z2 and so we put all variables in the same group
@@ -118,7 +131,9 @@ def filter_out_nonhomogenous_polynomials(formula: Formula, unrolled: List[Tuple[
     for name, polynomial in unrolled:
         homogenous = True
         for point_index, variables in input_variables_grouped.items():
-            weighted_variables = [(var, homogenity_weights[var[:-1]]) for var in variables]
+            weighted_variables = [
+                (var, homogenity_weights[var[:-1]]) for var in variables
+            ]
 
             # we dont check homogenity for the second point in madd formulas (which is affine)
             if "madd" in formula.name and point_index == 2:
@@ -126,8 +141,8 @@ def filter_out_nonhomogenous_polynomials(formula: Formula, unrolled: List[Tuple[
             homogenous &= is_homogeneous(Poly(polynomial), weighted_variables)
         if homogenous:
             filtered_unroll.append((name, polynomial))
-    return filtered_unroll   
-            
+    return filtered_unroll
+
 
 def is_homogeneous(polynomial: Poly, weighted_variables: List[Tuple[str, int]]) -> bool:
     """
@@ -137,28 +152,29 @@ def is_homogeneous(polynomial: Poly, weighted_variables: List[Tuple[str, int]])
     :param weighted_variables: The variables and their weights.
     :return: True if the polynomial is homogenous, otherwise False.
     """
-    hom = symbols('hom')
-    new_gens = polynomial.gens+(hom,)
-    univariate_poly = polynomial.subs({var: hom**weight for var, weight in weighted_variables})
-    univariate_poly = Poly(univariate_poly, *new_gens, domain = polynomial.domain)
+    hom = symbols("hom")
+    new_gens = polynomial.gens + (hom,)  # type: ignore[attr-defined]
+    univariate_poly = polynomial.subs(
+        {var: hom**weight for var, weight in weighted_variables}
+    )
+    univariate_poly = Poly(univariate_poly, *new_gens, domain=polynomial.domain)
     hom_index = univariate_poly.gens.index(hom)
     degrees = set(monom[hom_index] for monom in univariate_poly.monoms())
-    return len(degrees)<=1
+    return len(degrees) <= 1
 
 
 @public
-def compute_factor_set(formula: Formula, affine: bool = False) -> Set[Poly]:
+def compute_factor_set(formula: Formula) -> Set[Poly]:
     """
     Compute a set of factors present in the :paramref:`~.compute_factor_set.formula`.
 
-    If :paramref:`~.compute_factor_set.affine` is set, the polynomials are transformed
-    to affine form, using some assumptions along the way (e.g. `Z = 1`).
-
     :param formula: Formula to compute the factor set of.
-    :param affine: Whether to transform the unrolled polynomials (and thus the resulting factors) into affine form.
     :return: The set of factors present in the formula.
     """
-    unrolled = unroll_formula(formula, affine=affine)
+    unrolled = unroll_formula(formula)
+    unrolled = filter_out_nonhomogenous_polynomials(formula, unrolled)
+    unrolled = map_to_affine(formula, unrolled)
+
     factors = set()
     # Go over all the unrolled intermediates
     for name, poly in unrolled:
diff --git a/test/sca/test_zvp.py b/test/sca/test_zvp.py
index d7e3a42..32132fd 100644
--- a/test/sca/test_zvp.py
+++ b/test/sca/test_zvp.py
@@ -3,8 +3,17 @@ import pytest
 from pyecsca.ec.coordinates import AffineCoordinateModel
 from pyecsca.ec.mod import Mod
 from pyecsca.ec.point import Point
-from pyecsca.sca.re.zvp import unroll_formula, subs_curve_equation, remove_z, eliminate_y, subs_dlog, subs_curve_params, \
-    zvp_points, compute_factor_set
+from pyecsca.sca.re.zvp import (
+    unroll_formula,
+    map_to_affine,
+    subs_curve_equation,
+    remove_z,
+    eliminate_y,
+    subs_dlog,
+    subs_curve_params,
+    zvp_points,
+    compute_factor_set,
+)
 from pyecsca.ec.context import local, DefaultContext
 from sympy import symbols, Poly, sympify, FF
 
@@ -14,16 +23,23 @@ def formula(secp128r1, request):
     return secp128r1.curve.coordinate_model.formulas[request.param]
 
 
-@pytest.mark.parametrize("affine", [True, False])
-def test_unroll(formula, affine):
-    results = unroll_formula(formula, affine=affine)
+def test_unroll(formula):
+    results = unroll_formula(formula)
     assert results is not None
     for name, res in results:
         assert isinstance(res, Poly)
 
 
+def test_map_to_affine(formula):
+    results = unroll_formula(formula)
+    mapped = map_to_affine(formula, results)
+    assert mapped is not None
+    for name, res in mapped:
+        assert isinstance(res, Poly)
+
+
 def test_factor_set(formula):
-    factor_set = compute_factor_set(formula, affine=True)
+    factor_set = compute_factor_set(formula)
     assert factor_set is not None
     assert isinstance(factor_set, set)
     expr_set = set(map(lambda poly: poly.as_expr(), factor_set))
@@ -36,8 +52,8 @@ def test_factor_set(formula):
             # "x2", RPA
             # "x1", RPA
             "x1 + x2",
-            #"y1^2 + 2*y1*y2 + y2^2 + x1 + x2", Non-homogenous
-            #"y1^2 + 2*y1*y2 + y2^2 + 2*x1 + 2*x2", Non-homogenous
+            # "y1^2 + 2*y1*y2 + y2^2 + x1 + x2", Non-homogenous
+            # "y1^2 + 2*y1*y2 + y2^2 + 2*x1 + 2*x2", Non-homogenous
             "x1^2 + x1*x2 + x2^2",
             "a + x1^2 + x1*x2 + x2^2",
             # "a^2 + x1^4 + 2*x1^3*x2 + 3*x1^2*x2^2 + 2*x1*x2^3 + x2^4 - x1*y1^2 - x2*y1^2 - 2*x1*y1*y2 - 2*x2*y1*y2 - x1*y2^2 - x2*y2^2 + 2*x1^2*a + 2*x1*x2*a + 2*x2^2*a", RPA
@@ -71,16 +87,18 @@ def test_factor_set(formula):
             # "3*x1*x2^2*y1 + 3*x1^2*x2*y2 + y1^2*y2 + y1*y2^2 + x1*y1*a + 2*x2*y1*a + 2*x1*y2*a + x2*y2*a + 3*y1*b + 3*y2*b", RPA
             # "-3*x1^2*x2^2*a - y1^2*y2^2 + x1^2*a^2 + 4*x1*x2*a^2 + x2^2*a^2 - 9*x1^2*x2*b - 9*x1*x2^2*b + a^3 + 3*x1*a*b + 3*x2*a*b + 9*b^2" RPA
         },
-        "dbl-2007-bl": {"a + 3*x1^2", "a^2 + 6*x1^2*a + 9*x1^4 - 12*x1*y1^2"}
-
+        "dbl-2007-bl": {"a + 3*x1^2", "a^2 + 6*x1^2*a + 9*x1^4 - 12*x1*y1^2"},
     }
     if formula.name in expected_factors:
-        expected_set = set(map(lambda s: Poly(s).as_expr(), expected_factors[formula.name]))
+        expected_set = set(
+            map(lambda s: Poly(s).as_expr(), expected_factors[formula.name])
+        )
         assert expr_set == expected_set
 
 
 def test_curve_elimination(secp128r1, formula):
-    unrolled = unroll_formula(formula, affine=True)
+    unrolled = unroll_formula(formula)
+    unrolled = map_to_affine(formula, unrolled)
     subbed = subs_curve_equation(unrolled[-1][1], secp128r1.curve)
     assert subbed is not None
     Y1, Y2 = symbols("Y1,Y2")
@@ -93,14 +111,16 @@ def test_curve_elimination(secp128r1, formula):
 
 
 def test_remove_z(secp128r1, formula):
-    unrolled = unroll_formula(formula, affine=True)
+    unrolled = unroll_formula(formula)
+    unrolled = map_to_affine(formula, unrolled)
     removed = remove_z(unrolled[-1][1])
     for gen in removed.gens:
         assert not str(gen).startswith("Z")
 
 
 def test_eliminate_y(secp128r1, formula):
-    unrolled = unroll_formula(formula, affine=True)
+    unrolled = unroll_formula(formula)
+    unrolled = map_to_affine(formula, unrolled)
     subbed = subs_curve_equation(unrolled[-1][1], secp128r1.curve)
     eliminated = eliminate_y(subbed, secp128r1.curve)
     assert eliminated is not None
@@ -112,7 +132,8 @@ def test_eliminate_y(secp128r1, formula):
 
 
 def test_full(secp128r1, formula):
-    unrolled = unroll_formula(formula, affine=True)
+    unrolled = unroll_formula(formula)
+    unrolled = map_to_affine(formula, unrolled)
     subbed = subs_curve_equation(unrolled[-1][1], secp128r1.curve)
     removed = remove_z(subbed)
     eliminated = eliminate_y(removed, secp128r1.curve)
@@ -130,7 +151,8 @@ def test_full(secp128r1, formula):
 
 @pytest.mark.slow
 def test_zvp(secp128r1, formula):
-    unrolled = unroll_formula(formula, affine=True)
+    unrolled = unroll_formula(formula)
+    unrolled = map_to_affine(formula, unrolled)
     # Try all intermediates, zvp_point should return empty set if ZVP points do not exist
     for name, poly in unrolled:
         points = zvp_points(poly, secp128r1.curve, 5, secp128r1.order)
@@ -138,28 +160,71 @@ def test_zvp(secp128r1, formula):
 
         # If points are produced, try them all.
         for point in points:
-            second_point = secp128r1.curve.affine_multiply(point, 5)
-            p = point.to_model(formula.coordinate_model, secp128r1.curve)
-            q = second_point.to_model(formula.coordinate_model, secp128r1.curve)
+            inputs = [point.to_model(formula.coordinate_model, secp128r1.curve)]
+            if formula.num_inputs > 1:
+                second_point = secp128r1.curve.affine_multiply(point, 5)
+                inputs.append(
+                    second_point.to_model(formula.coordinate_model, secp128r1.curve)
+                )
             with local(DefaultContext()) as ctx:
-                formula(secp128r1.curve.prime, p, q, **secp128r1.curve.parameters)
+                formula(secp128r1.curve.prime, *inputs, **secp128r1.curve.parameters)
             action = next(iter(ctx.actions.keys()))
             results = list(map(lambda o: int(o.value), action.op_results))
             assert 0 in results
 
 
-@pytest.mark.parametrize("poly_str,point,k", [
-    ("y1 + y2", (54027047743185503031379008986257148598, 42633567686060343012155773792291852040), 4),
-    ("x1 + x2", (285130337309757533508049972949147801522, 55463852278545391044040942536845640298), 3),
-    ("x1*x2 + y1*y2", (155681799415564546404955983367992137717, 227436010604106449719780498844151836756), 5),
-    ("y1*y2 - x1*a - x2*a - 3*b", (169722400242675158455680894146658513260, 33263376472545436059176357032150610796), 4),
-    ("x1", (0, 594107526960909229279178399525926007), 3),
-    ("x2", (234937379492809870217296988280059595814, 101935882302108071650074851009662355573), 4),
-])
+@pytest.mark.parametrize(
+    "poly_str,point,k",
+    [
+        (
+            "y1 + y2",
+            (
+                54027047743185503031379008986257148598,
+                42633567686060343012155773792291852040,
+            ),
+            4,
+        ),
+        (
+            "x1 + x2",
+            (
+                285130337309757533508049972949147801522,
+                55463852278545391044040942536845640298,
+            ),
+            3,
+        ),
+        (
+            "x1*x2 + y1*y2",
+            (
+                155681799415564546404955983367992137717,
+                227436010604106449719780498844151836756,
+            ),
+            5,
+        ),
+        (
+            "y1*y2 - x1*a - x2*a - 3*b",
+            (
+                169722400242675158455680894146658513260,
+                33263376472545436059176357032150610796,
+            ),
+            4,
+        ),
+        ("x1", (0, 594107526960909229279178399525926007), 3),
+        (
+            "x2",
+            (
+                234937379492809870217296988280059595814,
+                101935882302108071650074851009662355573,
+            ),
+            4,
+        ),
+    ],
+)
 def test_points(secp128r1, poly_str, point, k):
-    pt = Point(AffineCoordinateModel(secp128r1.curve.model),
-               x=Mod(point[0], secp128r1.curve.prime),
-               y=Mod(point[1], secp128r1.curve.prime))
+    pt = Point(
+        AffineCoordinateModel(secp128r1.curve.model),
+        x=Mod(point[0], secp128r1.curve.prime),
+        y=Mod(point[1], secp128r1.curve.prime),
+    )
     poly_expr = sympify(poly_str)
     poly = Poly(poly_expr, domain=FF(secp128r1.curve.prime))
     res = zvp_points(poly, secp128r1.curve, k, secp128r1.order)