Fix affine formula unrolling.

1 files changed, 95 insertions, 84 deletions

diff --git a/pyecsca/sca/re/zvp.py b/pyecsca/sca/re/zvp.py
index b1d5d71..a92e860 100644
--- a/pyecsca/sca/re/zvp.py
+++ b/pyecsca/sca/re/zvp.py
@@ -10,7 +10,7 @@ from typing import List, Set
 from public import public
 from astunparse import unparse
 
-from sympy import symbols, FF, Poly, Monomial, Symbol, Expr, sympify
+from sympy import FF, Poly, Monomial, Symbol, Expr, sympify, symbols
 
 from ...ec.curve import EllipticCurve
 from ...ec.divpoly import mult_by_n
@@ -20,17 +20,21 @@ from ...ec.point import Point
 
 
 @public
-def unroll_formula(formula: Formula) -> List[Poly]:
+def unroll_formula(formula: Formula, affine: bool = False) -> List[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.
-    :return: List of symbolic intermediate values.
+    :param affine: Whether to transform the unrolled polynomials (and thus the resulting factors) into affine form.
+    :return: List of symbolic intermediate values, in formula coordinate model.
     """
     params = {var: symbols(var) for var in formula.coordinate_model.curve_model.parameter_names}
     inputs = {f"{var}{i}": symbols(f"{var}{i}") for var in formula.coordinate_model.variables for i in
               range(1, formula.num_inputs + 1)}
-    for assumption, assumption_string in zip(formula.assumptions, formula.assumptions_str):
+    for assumption_string in formula.assumptions_str:
         lhs, rhs = assumption_string.split(" == ")
         if lhs in formula.parameters:
             # Handle a symbolic assignment to a new parameter.
@@ -38,28 +42,6 @@ def unroll_formula(formula: Formula) -> List[Poly]:
             for curve_param, value in params.items():
                 expr = expr.subs(curve_param, value)
             params[lhs] = expr
-
-    locals = {**params, **inputs}
-    values = []
-    for op in formula.code:
-        result = op(**locals)
-        locals[op.result] = result
-        values.append(result)
-    return [Poly(value) for value in values]
-
-
-def compute_factor_set(formula: Formula, affine: bool = False) -> 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:
-    :param affine:
-    :return:
-    """
-    unrolled = unroll_formula(formula)
     subs_map = {}
     if affine:
         # tosystem_map is the mapping of system variables (without indices) in affine variables (without indices)
@@ -74,20 +56,48 @@ def compute_factor_set(formula: Formula, affine: bool = False) -> Set[Poly]:
                 subs_lhs = lhs + str(i)
                 subs_rhs = rhs.subs("x", f"x{i}").subs("y", f"y{i}")
                 subs_map[subs_lhs] = subs_rhs
-    factors = set()
-    # Go over all of the unrolled intermediates
-    for poly in unrolled:
+
+    locls = {**params, **inputs}
+    values = []
+    for op in formula.code:
+        result = op(**locls)
+        locls[op.result] = result
+        values.append(result)
+
+    results = []
+    for value in values:
         if affine:
-            expr = poly
+            expr = value
             for lhs, rhs in subs_map.items():
                 expr = expr.subs(lhs, rhs)
             if expr.free_symbols:
-                symbols = list(expr.free_symbols)
-                symbols.sort(key=str)
-                poly = Poly(expr, *symbols, domain=poly.domain)
+                gens = list(expr.free_symbols)
+                gens.sort(key=str)
+                poly = Poly(expr, *gens)
+                results.append(poly)
             else:
                 # Skip if no variables remain (constant poly)
                 continue
+        else:
+            results.append(Poly(value))
+    return results
+
+
+def compute_factor_set(formula: Formula, affine: bool = False) -> 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)
+    factors = set()
+    # Go over all of the unrolled intermediates
+    for poly in unrolled:
         # Factor the intermediate, don't worry about the coeff
         coeff, factor_list = poly.factor_list()
         # Go over all the factors of the intermediate, forget the power
@@ -109,10 +119,10 @@ def curve_equation(x: Symbol, curve: EllipticCurve, symbolic: bool = True) -> Ex
     Get the "ysquared" curve polynomial in :paramref:`~.x` for the :paramref:`~.curve`,
     either symbolically or with concrete parameter values.
 
-    :param x:
-    :param curve:
-    :param symbolic:
-    :return:
+    :param x: The sympy symbol to use in place of x.
+    :param curve: The elliptic curve to use.
+    :param symbolic: Whether to get the symbolic equation for the curve (with symbolic parameters) or actual curve parameter values.
+    :return: The sympy expression of the "ysquared" curve polynomial.
     """
     parameters = {name: symbols(name) if symbolic else curve.parameters[name] for name in curve.model.parameter_names}
     return eval(compile(curve.model.ysquared, "", mode="eval"), {"x": x, **parameters})
@@ -120,12 +130,12 @@ def curve_equation(x: Symbol, curve: EllipticCurve, symbolic: bool = True) -> Ex
 
 def subs_curve_equation(poly: Poly, curve: EllipticCurve) -> Poly:
     """
-    Substitute in the curve equation "ysquared" for Y repeatedly to
-    eliminate all but singular powers of Y.
+    Substitute in the curve equation "ysquared" for `y{1,2}` repeatedly to
+    eliminate all but singular powers of y.
 
-    :param poly:
-    :param curve:
-    :return:
+    :param poly: The sympy polynomial to substitute in.
+    :param curve: The elliptic curve to use.
+    :return: A polynomial with eliminated all but singular powers of y.
     """
     poly = Poly(poly, domain=FF(curve.prime))
     gens = poly.gens  # type: ignore[attr-defined]
@@ -134,9 +144,9 @@ def subs_curve_equation(poly: Poly, curve: EllipticCurve) -> Poly:
         sub = 1
         new_term = []
         for power, gen in zip(term[0], gens):
-            if str(gen).startswith("Y"):
-                X = symbols("X" + str(gen)[1:])
-                ysquared = curve_equation(X, curve)
+            if str(gen).startswith("y"):
+                x = symbols("x" + str(gen)[1:])
+                ysquared = curve_equation(x, curve)
                 sub *= ysquared ** (power // 2)
                 power %= 2
             new_term.append(power)
@@ -149,9 +159,9 @@ def subs_curve_params(poly: Poly, curve: EllipticCurve) -> Poly:
     """
     Substitute in the concrete curve parameters.
 
-    :param poly:
-    :param curve:
-    :return:
+    :param poly: The sympy polynomial to substitute in.
+    :param curve: The elliptic curve to use.
+    :return: A polynomial with substituted in concrete curve parameters.
     """
     poly = Poly(poly, domain=FF(curve.prime))
     for name, value in curve.parameters.items():
@@ -163,25 +173,26 @@ def subs_curve_params(poly: Poly, curve: EllipticCurve) -> Poly:
 
 def subs_dlog(poly: Poly, k: int, curve: EllipticCurve):
     """
-    Substitute in the multiplication-by-k-map(X1) in place of X2.
+    Substitute in the multiplication-by-k-map(x1) in place of x2.
 
-    :param poly:
-    :param k:
-    :param curve:
-    :return:
+    :param poly: The sympy polynomial to substitute in.
+    :param k: The dlog between the points.
+    :param curve: The elliptic curve to use.
+    :return: A polynomial with the map substituted (with only x1 coords remaining).
     """
     poly = Poly(poly, domain=FF(curve.prime))
-    X1, X2 = symbols("X1,X2")
+    x1, x2 = symbols("x1,x2")
     gens = poly.gens  # type: ignore[attr-defined]
-    if X2 not in gens or X1 not in gens:
+    if x2 not in gens or x1 not in gens:
         return poly
-    max_degree = poly.degree(X2)
-    X2i = gens.index(X2)
+    max_degree = poly.degree(x2)
+    x2i = gens.index(x2)
     new_gens = set(gens)
-    new_gens.remove(X2)
+    new_gens.remove(x2)
 
     mx, my = mult_by_n(curve, k)
-    u, v = mx[0].subs("x", X1), mx[1].subs("x", X1)
+    # TODO: my is unnecessary here so maybe add a function to not compute it (speedup).
+    u, v = mx[0].subs("x", x1), mx[1].subs("x", x1)
 
     # The polynomials are quite dense, hence it makes sense
     # to compute all of the u and v powers in advance and
@@ -195,7 +206,7 @@ def subs_dlog(poly: Poly, k: int, curve: EllipticCurve):
 
     uv_factors = {}
     for term in poly.terms():
-        u_power = term[0][X2i]
+        u_power = term[0][x2i]
         v_power = max_degree - u_power
         if (u_power, v_power) in uv_factors:
             continue
@@ -204,9 +215,9 @@ def subs_dlog(poly: Poly, k: int, curve: EllipticCurve):
     res = 0
     for term in poly.terms():
         powers = list(term[0])
-        u_power = powers[X2i]
+        u_power = powers[x2i]
         v_power = max_degree - u_power
-        powers[X2i] = 0
+        powers[x2i] = 0
         monom = Monomial(powers, gens).as_expr() * term[1]
         res += Poly(monom, *new_gens, domain=poly.domain) * uv_factors[(u_power, v_power)]
     return Poly(res, domain=poly.domain)
@@ -216,8 +227,8 @@ def remove_z(poly: Poly) -> Poly:
     """
     Substitute in 1 for all Zs (because we can do that ;).
 
-    :param poly:
-    :return:
+    :param poly: The sympy polynomial to substitute in.
+    :return: A polynomial with Zs eliminated.
     """
     for gen in poly.gens:  # type: ignore[attr-defined]
         if str(gen).startswith("Z"):
@@ -227,19 +238,19 @@ def remove_z(poly: Poly) -> Poly:
 
 def eliminate_y(poly: Poly, curve: EllipticCurve) -> Poly:
     """
-    Eliminate the remaining Ys (only power 1).
+    Eliminate the remaining ys (only power 1).
 
     See [FFD]_ page 11.
 
-    :param poly:
-    :param curve:
+    :param poly: The sympy polynomial to eliminate in.
+    :param curve: The elliptic curve to use.
     :return:
     """
     poly = Poly(poly, domain=FF(curve.prime))
-    Y1, Y2 = symbols("Y1,Y2")
+    x1, x2, y1, y2 = symbols("x1,x2,y1,y2")
     gens = poly.gens  # type: ignore[attr-defined]
-    Y1i = gens.index(Y1) if Y1 in gens else None
-    Y2i = gens.index(Y2) if Y2 in gens else None
+    y1i = gens.index(y1) if y1 in gens else None
+    y2i = gens.index(y2) if y2 in gens else None
     f0 = 0
     f1 = 0
     f2 = 0
@@ -247,37 +258,37 @@ def eliminate_y(poly: Poly, curve: EllipticCurve) -> Poly:
     for term in poly.terms():
         monom = term[0]
         monom_expr = Monomial(term[0], gens).as_expr() * term[1]
-        y1_present = not (Y1i is None or monom[Y1i] == 0)
-        y2_present = not (Y2i is None or monom[Y2i] == 0)
+        y1_present = not (y1i is None or monom[y1i] == 0)
+        y2_present = not (y2i is None or monom[y2i] == 0)
         if y1_present and y2_present:
-            f12 += monom_expr.subs(Y1, 1).subs(Y2, 1)
+            f12 += monom_expr.subs(y1, 1).subs(y2, 1)
         elif y1_present and not y2_present:
-            f1 += monom_expr.subs(Y1, 1)
+            f1 += monom_expr.subs(y1, 1)
         elif not y1_present and y2_present:
-            f2 += monom_expr.subs(Y2, 1)
+            f2 += monom_expr.subs(y2, 1)
         elif not y1_present and not y2_present:
             f0 += monom_expr
-    fe_X1 = curve_equation(symbols("X1"), curve)
-    fe_X2 = curve_equation(symbols("X2"), curve)
+    fe_x1 = curve_equation(x1, curve)
+    fe_x2 = curve_equation(x2, curve)
 
     # [FFD] page 11
-    f_prime = fe_X2 * (fe_X1 * 2 * f1 * f12 - 2 * f0 * f2) ** 2 - \
-        (f0 ** 2 + f2 ** 2 * fe_X2 - fe_X1 * (f1 ** 2 + f12 ** 2 * fe_X2)) ** 2
+    f_prime = fe_x2 * (fe_x1 * 2 * f1 * f12 - 2 * f0 * f2) ** 2 - \
+        (f0 ** 2 + f2 ** 2 * fe_x2 - fe_x1 * (f1 ** 2 + f12 ** 2 * fe_x2)) ** 2
     return Poly(f_prime, domain=poly.domain)
 
 
 @public
 def zvp_points(poly: Poly, curve: EllipticCurve, k: int) -> Set[Point]:
     """
-    Find a set of ZVP points for a given intermediate value and dlog relationship.
+    Find a set of (affine) ZVP points for a given intermediate value and dlog relationship.
 
     :param poly: The polynomial to zero out, obtained as a result of :py:meth:`.unroll_formula` (or its factor).
     :param curve: The curve to compute over.
-    :param k: The discrete-log relationship between the two points, i.e. (X2, Y2) = [k](X1, Y1)
-    :return: The set of points (X1, Y1).
+    :param k: The discrete-log relationship between the two points, i.e. (x2, x2) = [k](x1, x1)
+    :return: The set of points (x1, x1).
     """
     # If input poly is trivial (only in params), abort early
-    if not set(symbols("X1,X2,Y1,Y2")).intersection(poly.gens):  # type: ignore[attr-defined]
+    if not set(symbols("x1,x2,y1,y2")).intersection(poly.gens):  # type: ignore[attr-defined]
         return set()
     poly = Poly(poly, domain=FF(curve.prime))
     # Start with removing all squares of Y1, Y2
@@ -299,7 +310,7 @@ def zvp_points(poly: Poly, curve: EllipticCurve, k: int) -> Set[Point]:
         # Check that the points zero out the original polynomial to filter out erroneous candidates
         for point in pt:
             other = curve.affine_multiply(point, k)
-            inputs = {"X1": point.x, "Y1": point.y, "X2": other.x, "Y2": other.y, "Z1": 1, "Z2": 1, **curve.parameters}
+            inputs = {"x1": point.x, "y1": point.y, "x2": other.x, "y2": other.y, **curve.parameters}
             res = poly.eval([inputs[str(gen)] for gen in poly.gens])  # type: ignore[attr-defined]
             if res == 0:
                 points.add(point)