Add test for formula expansion.

2 files changed, 38 insertions, 31 deletions

diff --git a/pyecsca/ec/formula/expand.py b/pyecsca/ec/formula/expand.py
index 3ec4665..3265131 100644
--- a/pyecsca/ec/formula/expand.py
+++ b/pyecsca/ec/formula/expand.py
@@ -1,5 +1,7 @@
-from typing import List
+from typing import List, Callable, Any
+from public import public
 
+from . import Formula
 from .efd import EFDFormula
 from .fliparoo import recursive_fliparoo
 from .graph import ModifiedEFDFormula
@@ -21,38 +23,29 @@ def reduce_with_similarity(formulas: List[EFDFormula], norm):
     return reduced_efd
 
 
-def expand_formula_list(formulas: List[EFDFormula]):
-    extended_efd = reduce_with_similarity(formulas, ivs_norm)
-    print(f"Reduced to {len(extended_efd)} formulas")
+@public
+def expand_formula_list(
+    formulas: List[EFDFormula], norm: Callable[[Formula], Any] = ivs_norm
+) -> List[EFDFormula]:
+    extended = reduce_with_similarity(formulas, norm)
 
-    fliparood: List[EFDFormula] = sum(list(map(recursive_fliparoo, extended_efd)), [])
-    extended_efd.extend(fliparood)
-    print(f"Fliparoo: {len(extended_efd)} formulas")
-    extended_efd = reduce_with_similarity(extended_efd, ivs_norm)
-    print(f"Reduced to {len(extended_efd)} formulas")
-    # list(map(test_formula, extended_efd))
+    fliparood: List[EFDFormula] = sum(list(map(recursive_fliparoo, extended)), [])
+    extended.extend(fliparood)
+    extended = reduce_with_similarity(extended, norm)
 
-    switch_signs: List[EFDFormula] = sum([list(generate_switched_formulas(f)) for f in extended_efd], [])
-    extended_efd.extend(switch_signs)
-    print(f"Switch signs: {len(extended_efd)} formulas")
-    extended_efd = reduce_with_similarity(extended_efd, ivs_norm)
-    print(f"Reduced to {len(extended_efd)} formulas")
-    # list(map(test_formula, extended_efd))
+    switch_signs: List[EFDFormula] = sum(
+        [list(generate_switched_formulas(f)) for f in extended], []
+    )
+    extended.extend(switch_signs)
+    extended = reduce_with_similarity(extended, norm)
 
-    extended_efd.extend(list(map(reduce_all_adds, extended_efd)))
-    print(f"Compress adds: {len(extended_efd)} formulas")
-    extended_efd = reduce_with_similarity(extended_efd, ivs_norm)
-    print(f"Reduced to {len(extended_efd)} formulas")
-    # list(map(test_formula, extended_efd))
+    extended.extend(list(map(reduce_all_adds, extended)))
+    extended = reduce_with_similarity(extended, norm)
 
-    extended_efd.extend(list(map(expand_all_muls, extended_efd)))
-    print(f"Expand muls: {len(extended_efd)} formulas")
-    extended_efd = reduce_with_similarity(extended_efd, ivs_norm)
-    print(f"Reduced to {len(extended_efd)} formulas")
-    # list(map(test_formula, extended_efd))
+    extended.extend(list(map(expand_all_muls, extended)))
+    extended = reduce_with_similarity(extended, norm)
 
-    extended_efd.extend(list(map(expand_all_nopower2_muls, extended_efd)))
-    print(f"Expand muls(!=2^):{len(extended_efd)} formulas")
-    extended_efd = reduce_with_similarity(extended_efd, ivs_norm)
-    print(f"Reduced to {len(extended_efd)} formulas")
-    # list(map(test_formula, extended_efd))
+    extended.extend(list(map(expand_all_nopower2_muls, extended)))
+    extended = reduce_with_similarity(extended, norm)
+
+    return extended
diff --git a/test/ec/test_formula.py b/test/ec/test_formula.py
index 7852453..853b47b 100644
--- a/test/ec/test_formula.py
+++ b/test/ec/test_formula.py
@@ -6,6 +6,7 @@ import pytest
 from sympy import FF, symbols
 from importlib_resources import files, as_file
 import test.data.formulas
+from pyecsca.ec.formula.expand import expand_formula_list
 from pyecsca.ec.formula.fliparoo import generate_fliparood_formulas
 from pyecsca.ec.formula.graph import rename_ivs
 from pyecsca.ec.formula.metrics import formula_similarity
@@ -13,6 +14,7 @@ from pyecsca.ec.formula.partitions import (
     reduce_all_adds,
     expand_all_muls,
     expand_all_nopower2_muls,
+    generate_partitioned_formulas,
 )
 from pyecsca.ec.formula.switch_sign import generate_switched_formulas
 from pyecsca.ec.mod import SymbolicMod, Mod
@@ -409,6 +411,13 @@ def test_fliparood_formula(library_formula_params):
         do_test_formula(fliparood, params)
 
 
+@pytest.mark.slow
+def test_partition_formula(library_formula_params):
+    formula, params = library_formula_params
+    for partitioned in generate_partitioned_formulas(formula):
+        do_test_formula(partitioned, params)
+
+
 def test_reductions(library_formula_params):
     formula, params = library_formula_params
     do_test_formula(reduce_all_adds(formula), params)
@@ -487,3 +496,8 @@ def do_test_formula(formula, params):
 def test_formula_correctness(library_formula_params):
     formula, params = library_formula_params
     do_test_formula(formula, params)
+
+
+def test_formula_expand(add):
+    res = expand_formula_list([add])
+    assert len(res) > 1