Unify unroll and unroll with names.

2 files changed, 38 insertions, 43 deletions

diff --git a/pyecsca/sca/re/zvp.py b/pyecsca/sca/re/zvp.py
index 16f5a82..9189ae7 100644
--- a/pyecsca/sca/re/zvp.py
+++ b/pyecsca/sca/re/zvp.py
@@ -20,11 +20,7 @@ from ...ec.point import Point
 
 
 @public
-def unroll_formula(formula: Formula, affine: bool = False) -> List[Poly]:
-    return [poly for _,poly in unroll_formula_with_names(formula, affine)]
-
-
-def unroll_formula_with_names(formula: Formula, affine: bool = False) -> List[Tuple[str, Poly]]:
+def unroll_formula(formula: Formula, affine: bool = False) -> List[Tuple[str, Poly]]:
     """
     Unroll a given formula symbolically to obtain symbolic expressions for its intermediate values.
 
@@ -33,7 +29,7 @@ def unroll_formula_with_names(formula: Formula, affine: bool = False) -> List[Tu
 
     :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, in formula coordinate model.
+    :return: List of symbolic intermediate values, with associated variable names.
     """
     params = {
         var: symbols(var)
@@ -72,7 +68,7 @@ def unroll_formula_with_names(formula: Formula, affine: bool = False) -> List[Tu
     for op in formula.code:
         result = op(**locls)
         locls[op.result] = result
-        values.append((op.result,result))
+        values.append((op.result, result))
 
     results = []
     for result_var, value in values:
@@ -105,9 +101,9 @@ def compute_factor_set(formula: Formula, affine: bool = False) -> Set[Poly]:
     :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_with_names(formula, affine=affine)
+    unrolled = unroll_formula(formula, affine=affine)
     factors = set()
-    # Go over all of the unrolled intermediates
+    # Go over all the unrolled intermediates
     for name, poly in unrolled:
         # Factor the intermediate, don't worry about the coeff
         coeff, factor_list = poly.factor_list()
@@ -129,13 +125,20 @@ def compute_factor_set(formula: Formula, affine: bool = False) -> Set[Poly]:
     return factors
 
 
-def filter_out_rpa_polynomials(factor_set: Set[Poly], formula: Formula, unrolled: List[Tuple[str, Poly]]) -> Set[Poly]:
+def filter_out_rpa_polynomials(
+    factor_set: Set[Poly], formula: Formula, unrolled: List[Tuple[str, Poly]]
+) -> Set[Poly]:
     """
-    Remove polynomials from factorset that imply RPA
+    Remove polynomials from factorset that imply RPA points (on input or output).
+
+    :param factor_set: The factor set to filter.
+    :param formula: The formula that the factor set belongs to.
+    :param unrolled: The unrolled formula.
+    :return: The filtered factor set.
     """
-    
+
     # Find polynomials that define the output variables
-    # We save the latest occurence of output variable in the list of ops
+    # We save the latest occurrence of output variable in the list of ops
     output_polynomials = {}
     for name, polynomial in unrolled:
         if name in formula.outputs:
@@ -153,8 +156,8 @@ def filter_out_rpa_polynomials(factor_set: Set[Poly], formula: Formula, unrolled
         if not divisible:
             filtered_factorset.add(poly)
     return filtered_factorset
-                
-    
+
+
 def curve_equation(x: Symbol, curve: EllipticCurve, symbolic: bool = True) -> Expr:
     """
     Get the "ysquared" curve polynomial in :paramref:`~.x` for the :paramref:`~.curve`,
@@ -356,11 +359,7 @@ def zvp_points(poly: Poly, curve: EllipticCurve, k: int, n: int) -> Set[Point]:
         for root, multiplicity in roots.items():
             pt = curve.affine_lift_x(Mod(int(root), curve.prime))
             for point in pt:
-                inputs = {
-                    "x1": point.x,
-                    "y1": point.y,
-                    **curve.parameters
-                }
+                inputs = {"x1": point.x, "y1": point.y, **curve.parameters}
                 res = poly.eval([inputs[str(gen)] for gen in poly.gens])  # type: ignore[attr-defined]
                 if res == 0:
                     points.add(point)
@@ -372,11 +371,7 @@ def zvp_points(poly: Poly, curve: EllipticCurve, k: int, n: int) -> Set[Point]:
         for root, multiplicity in roots.items():
             pt = curve.affine_lift_x(Mod(int(root), curve.prime))
             for point in pt:
-                inputs = {
-                    "x2": point.x,
-                    "y2": point.y,
-                    **curve.parameters
-                }
+                inputs = {"x2": point.x, "y2": point.y, **curve.parameters}
                 res = poly.eval([inputs[str(gen)] for gen in poly.gens])  # type: ignore[attr-defined]
                 if res == 0:
                     one = curve.affine_multiply(point, int(k_inv))
diff --git a/test/sca/test_zvp.py b/test/sca/test_zvp.py
index 7548963..f54f72a 100644
--- a/test/sca/test_zvp.py
+++ b/test/sca/test_zvp.py
@@ -18,7 +18,7 @@ def formula(secp128r1, request):
 def test_unroll(formula, affine):
     results = unroll_formula(formula, affine=affine)
     assert results is not None
-    for res in results:
+    for name, res in results:
         assert isinstance(res, Poly)
 
 
@@ -30,11 +30,11 @@ def test_factor_set(formula):
 
     expected_factors = {
         "add-2007-bl": {
-            #"y2", RPA
-            #"y1", RPA
-            #"y1 + y2", RPA
-            #"x2", RPA
-            #"x1", RPA
+            # "y2", RPA
+            # "y1", RPA
+            # "y1 + y2", RPA
+            # "x2", RPA
+            # "x1", RPA
             "x1 + x2",
             "y1^2 + 2*y1*y2 + y2^2 + x1 + x2",
             "y1^2 + 2*y1*y2 + y2^2 + 2*x1 + 2*x2",
@@ -45,15 +45,15 @@ def test_factor_set(formula):
             # "2*a^3 + 2*x1^6 + 6*x1^5*x2 + 12*x1^4*x2^2 + 14*x1^3*x2^3 + 12*x1^2*x2^4 + 6*x1*x2^5 + 2*x2^6 - 3*x1^3*y1^2 - 6*x1^2*x2*y1^2 - 6*x1*x2^2*y1^2 - 3*x2^3*y1^2 - 6*x1^3*y1*y2 - 12*x1^2*x2*y1*y2 - 12*x1*x2^2*y1*y2 - 6*x2^3*y1*y2 - 3*x1^3*y2^2 - 6*x1^2*x2*y2^2 - 6*x1*x2^2*y2^2 - 3*x2^3*y2^2 + 6*x1^4*a + 12*x1^3*x2*a + 18*x1^2*x2^2*a + 12*x1*x2^3*a + 6*x2^4*a + y1^4 + 4*y1^3*y2 + 6*y1^2*y2^2 + 4*y1*y2^3 + y2^4 - 3*x1*y1^2*a - 3*x2*y1^2*a - 6*x1*y1*y2*a - 6*x2*y1*y2*a - 3*x1*y2^2*a - 3*x2*y2^2*a + 6*x1^2*a^2 + 6*x1*x2*a^2 + 6*x2^2*a^2" RPA
         },
         "add-2015-rcb": {
-            #"y2", RPA
+            # "y2", RPA
             "y2 + 1",
-            #"y1", RPA
+            # "y1", RPA
             "y1 + 1",
             "y1 + y2",
-            #"x2", RPA
+            # "x2", RPA
             "x2 + 1",
             "x2 + y2",
-            #"x1", RPA
+            # "x1", RPA
             "x1 + 1",
             "x1 + y1",
             "x1 + x2",
@@ -67,9 +67,9 @@ def test_factor_set(formula):
             "x1*x2 + y1*y2",
             "3*x1*x2 + a",
             "a^2 - x1*x2*a - 3*x1*b - 3*x2*b",
-            #"x2*y1^2*y2 + x1*y1*y2^2 - 2*x1*x2*y1*a - x2^2*y1*a - x1^2*y2*a - 2*x1*x2*y2*a + y1*a^2 + y2*a^2 - 3*x1*y1*b - 6*x2*y1*b - 6*x1*y2*b - 3*x2*y2*b", RPA
-            #"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
+            # "x2*y1^2*y2 + x1*y1*y2^2 - 2*x1*x2*y1*a - x2^2*y1*a - x1^2*y2*a - 2*x1*x2*y2*a + y1*a^2 + y2*a^2 - 3*x1*y1*b - 6*x2*y1*b - 6*x1*y2*b - 3*x2*y2*b", RPA
+            # "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
         }
     }
     if formula.name in expected_factors:
@@ -79,7 +79,7 @@ def test_factor_set(formula):
 
 def test_curve_elimination(secp128r1, formula):
     unrolled = unroll_formula(formula, affine=True)
-    subbed = subs_curve_equation(unrolled[-1], secp128r1.curve)
+    subbed = subs_curve_equation(unrolled[-1][1], secp128r1.curve)
     assert subbed is not None
     Y1, Y2 = symbols("Y1,Y2")
 
@@ -92,14 +92,14 @@ def test_curve_elimination(secp128r1, formula):
 
 def test_remove_z(secp128r1, formula):
     unrolled = unroll_formula(formula, affine=True)
-    removed = remove_z(unrolled[-1])
+    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)
-    subbed = subs_curve_equation(unrolled[-1], secp128r1.curve)
+    subbed = subs_curve_equation(unrolled[-1][1], secp128r1.curve)
     eliminated = eliminate_y(subbed, secp128r1.curve)
     assert eliminated is not None
     assert isinstance(eliminated, Poly)
@@ -111,7 +111,7 @@ def test_eliminate_y(secp128r1, formula):
 
 def test_full(secp128r1, formula):
     unrolled = unroll_formula(formula, affine=True)
-    subbed = subs_curve_equation(unrolled[-1], secp128r1.curve)
+    subbed = subs_curve_equation(unrolled[-1][1], secp128r1.curve)
     removed = remove_z(subbed)
     eliminated = eliminate_y(removed, secp128r1.curve)
     dlog = subs_dlog(eliminated, 3, secp128r1.curve)
@@ -130,7 +130,7 @@ def test_full(secp128r1, formula):
 def test_zvp(secp128r1, formula):
     unrolled = unroll_formula(formula, affine=True)
     # Try all intermediates, zvp_point should return empty set if ZVP points do not exist
-    for poly in unrolled:
+    for name, poly in unrolled:
         points = zvp_points(poly, secp128r1.curve, 5, secp128r1.order)
         assert isinstance(points, set)