5 files changed, 182 insertions, 57 deletions

diff --git a/pyecsca/ec/formula/expand.py b/pyecsca/ec/formula/expand.py
index 2f631f7..7e9b5ee 100644
--- a/pyecsca/ec/formula/expand.py
+++ b/pyecsca/ec/formula/expand.py
@@ -1,29 +1,48 @@
 """Provides a formula expansion function."""
-from typing import Set, Callable, Any
+
+from typing import Set, Callable, Any, List, Iterable, Optional
 from public import public
+from operator import attrgetter
+from functools import lru_cache
 
 from pyecsca.ec.formula.base import Formula
 from pyecsca.ec.formula.efd import EFDFormula
 from pyecsca.ec.formula.fliparoo import recursive_fliparoo
 from pyecsca.ec.formula.metrics import ivs_norm
-from pyecsca.ec.formula.partitions import reduce_all_adds, expand_all_muls, expand_all_nopower2_muls
-from pyecsca.ec.formula.switch_sign import generate_switched_formulas
+from pyecsca.ec.formula.partitions import (
+    reduce_all_adds,
+    expand_all_muls,
+    expand_all_nopower2_muls,
+)
+from pyecsca.ec.formula.switch_sign import generate_switched_formulas, switch_signs
+from pyecsca.misc.utils import TaskExecutor
 
 
 def reduce_with_similarity(
-    formulas: Set[Formula], norm: Callable[[Formula], Any]
-) -> Set[Formula]:
-    reduced = set(filter(lambda x: isinstance(x, EFDFormula), formulas))
+    formulas: List[Formula], norm: Callable[[Formula], Any]
+) -> List[Formula]:
+    formulas = reduce_by_eq(formulas)
+    reduced = list(filter(lambda x: isinstance(x, EFDFormula), formulas))
     similarities = list(map(norm, reduced))
     for formula in formulas:
         n = norm(formula)
         if n in similarities:
             continue
         similarities.append(n)
-        reduced.add(formula)
+        reduced.append(formula)
     return reduced
 
 
+def reduce_by_eq(formulas: Iterable[Formula]) -> List[Formula]:
+    unique = set()
+    result = []
+    for formula in formulas:
+        if formula not in unique:
+            unique.add(formula)
+            result.append(formula)
+    return result
+
+
 @public
 def expand_formula_set(
     formulas: Set[Formula], norm: Callable[[Formula], Any] = ivs_norm
@@ -34,29 +53,103 @@ def expand_formula_set(
      - Sign switching
      - Associativity and Commutativity
 
-    :param formulas:
-    :param norm:
-    :return:
+    :param formulas: The set of formulas to expand.
+    :param norm: The norm to use while reducing.
+    :return: The expanded set of formulas.
     """
-    extended = reduce_with_similarity(formulas, norm)
 
-    fliparood: Set[Formula] = set().union(*map(recursive_fliparoo, extended))
-    extended.update(fliparood)
-    extended = reduce_with_similarity(extended, norm)
+    @lru_cache(maxsize=1000)
+    def cached(formula):
+        return norm(formula)
+
+    extended = sorted(formulas, key=attrgetter("name"))
+    extended = reduce_with_similarity(extended, cached)
+
+    fliparood: List[Formula] = reduce_by_eq(
+        sum(map(lambda f: reduce_by_eq(recursive_fliparoo(f)), extended), [])
+    )
+    extended.extend(fliparood)
+    extended = reduce_with_similarity(extended, cached)
 
-    switch_signs: Set[Formula] = set().union(
-        *(set(generate_switched_formulas(f)) for f in extended)
+    switch_signs: List[Formula] = reduce_by_eq(
+        sum(map(lambda f: reduce_by_eq(generate_switched_formulas(f)), extended), [])
     )
-    extended.update(switch_signs)
-    extended = reduce_with_similarity(extended, norm)
+    extended.extend(switch_signs)
+    extended = reduce_with_similarity(extended, cached)
+
+    add_reduced: List[Formula] = reduce_by_eq(list(map(reduce_all_adds, extended)))
+    extended.extend(add_reduced)
+    extended = reduce_with_similarity(extended, cached)
+
+    mul_expanded: List[Formula] = reduce_by_eq(list(map(expand_all_muls, extended)))
+    extended.extend(mul_expanded)
+    extended = reduce_with_similarity(extended, cached)
+
+    np2_expanded: List[Formula] = reduce_by_eq(
+        list(map(expand_all_nopower2_muls, extended))
+    )
+    extended.extend(np2_expanded)
+    extended = reduce_with_similarity(extended, cached)
+
+    return set(reduce_by_eq(extended))
+
+
+@public
+def expand_formula_set_parallel(
+    formulas: Set[Formula],
+    norm: Callable[[Formula], Any] = ivs_norm,
+    num_workers: int = 1,
+) -> Set[Formula]:
+    """
+    Expand a set of formulas by using transformations (parallelized):
+     - Fliparoos
+     - Sign switching
+     - Associativity and Commutativity
+
+    :param formulas: The set of formulas to expand.
+    :param norm: The norm to use while reducing.
+    :param num_workers: The amount of workers to use.
+    :return: The expanded set of formulas.
+    """
+
+    @lru_cache(maxsize=1000)
+    def cached(formula):
+        return norm(formula)
+
+    def map_multiple(pool, formulas, fn):
+        results: List[List[Formula]] = []
+        for f in formulas:
+            pool.submit_task(f, fn, f)
+            results.append([])
+        for f, future in pool.as_completed():
+            results[formulas.index(f)] = reduce_by_eq(future.result())
+        return reduce_by_eq(sum(results, []))
+
+    def map_single(pool, formulas, fn):
+        return reduce_by_eq(pool.map(fn, formulas))
+
+    extended = sorted(formulas, key=attrgetter("name"))
+    extended = reduce_with_similarity(extended, cached)
+
+    with TaskExecutor(max_workers=num_workers) as pool:
+        fliparood = map_multiple(pool, extended, recursive_fliparoo)
+        extended.extend(fliparood)
+        extended = reduce_with_similarity(extended, cached)
+
+        switched = map_multiple(pool, extended, switch_signs)
+        extended.extend(switched)
+        extended = reduce_with_similarity(extended, cached)
 
-    extended.update(set(map(reduce_all_adds, extended)))
-    extended = reduce_with_similarity(extended, norm)
+        add_reduced = map_single(pool, extended, reduce_all_adds)
+        extended.extend(add_reduced)
+        extended = reduce_with_similarity(extended, cached)
 
-    extended.update(set(map(expand_all_muls, extended)))
-    extended = reduce_with_similarity(extended, norm)
+        mul_expanded = map_single(pool, extended, expand_all_muls)
+        extended.extend(mul_expanded)
+        extended = reduce_with_similarity(extended, cached)
 
-    extended.update(set(map(expand_all_nopower2_muls, extended)))
-    extended = reduce_with_similarity(extended, norm)
+        np2_expanded = map_single(pool, extended, expand_all_nopower2_muls)
+        extended.extend(np2_expanded)
+        extended = reduce_with_similarity(extended, cached)
 
-    return extended
+    return set(reduce_by_eq(extended))
diff --git a/pyecsca/ec/formula/fliparoo.py b/pyecsca/ec/formula/fliparoo.py
index 4e1d221..4d14f80 100644
--- a/pyecsca/ec/formula/fliparoo.py
+++ b/pyecsca/ec/formula/fliparoo.py
@@ -1,6 +1,6 @@
 """Provides a way to Fliparoo formulas."""
 from ast import parse
-from typing import Iterator, List, Type, Optional, Set
+from typing import Iterator, List, Type, Optional
 from public import public
 
 from pyecsca.ec.op import OpType
@@ -254,7 +254,7 @@ def generate_fliparood_formulas(
     fliparoos = find_fliparoos(graph)
     for i, fliparoo in enumerate(fliparoos):
         for j, flip_graph in enumerate(generate_fliparood_graphs(fliparoo)):
-            yield flip_graph.to_formula(f"fliparoo[{i},{j}]")
+            yield flip_graph.to_formula(f"fliparoo[{i},{j}]")  # noqa
 
 
 def generate_fliparood_graphs(fliparoo: Fliparoo) -> Iterator[FormulaGraph]:
@@ -380,19 +380,19 @@ def combine_signed_nodes(
 
 
 @public
-def recursive_fliparoo(formula: Formula, depth: int = 2) -> Set[Formula]:
-    all_fliparoos = {0: {formula}}
+def recursive_fliparoo(formula: Formula, depth: int = 2) -> List[Formula]:
+    all_fliparoos = {0: [formula]}
     counter = 0
     while depth > counter:
         prev_level = all_fliparoos[counter]
-        fliparoo_level: Set[Formula] = set()
+        fliparoo_level: List[Formula] = []
         for flipparood_formula in prev_level:
             rename = not counter  # rename ivs before the first fliparoo
             for newly_fliparood in generate_fliparood_formulas(
                 flipparood_formula, rename
             ):
-                fliparoo_level.add(newly_fliparood)
+                fliparoo_level.append(newly_fliparood)
         counter += 1
         all_fliparoos[counter] = fliparoo_level
 
-    return set().union(*all_fliparoos.values())
+    return sum(all_fliparoos.values(), [])
diff --git a/pyecsca/ec/formula/graph.py b/pyecsca/ec/formula/graph.py
index d4a79b1..e2fb1ee 100644
--- a/pyecsca/ec/formula/graph.py
+++ b/pyecsca/ec/formula/graph.py
@@ -125,7 +125,7 @@ class CodeOpNode(Node):
 
     @property
     def label(self) -> str:
-        return f"{self.op.result}:{self.op.operator.op_str}"
+        return f"{self.op.result}:{self.op.operator.op_str}"  # noqa
 
     @property
     def result(self) -> str:
@@ -198,7 +198,7 @@ class InputNode(Node):
 
 def formula_input_variables(formula: Formula) -> List[str]:
     return (
-        list(formula.inputs)
+        sorted(formula.inputs)
         + formula.parameters
         + formula.coordinate_model.curve_model.parameter_names
     )
@@ -282,13 +282,10 @@ class FormulaGraph:
 
     def networkx_graph(self) -> nx.DiGraph:
         graph = nx.DiGraph()
-        vertices = list(range(len(self.nodes)))
-        graph.add_nodes_from(vertices)
-        stack = self.roots.copy()
-        while stack:
-            node = stack.pop()
+        for i, node in enumerate(self.nodes):
+            graph.add_node(i, result=node.result, label=node.label, op=getattr(node, "op", None))
+        for node in self.nodes:
             for out in node.outgoing_nodes:
-                stack.append(out)
                 graph.add_edge(self.node_index(node), self.node_index(out))
         return graph
 
@@ -348,7 +345,7 @@ class FormulaGraph:
             iu = self.node_index(u)
             for v in self.output_nodes():
                 iv = self.node_index(v)
-                index_paths.extend(nx.all_simple_paths(gnx, iu, iv))
+                index_paths.extend(sorted(nx.all_simple_paths(gnx, iu, iv)))
         paths = []
         for p in index_paths:
             paths.append([self.nodes[v] for v in p])
diff --git a/pyecsca/ec/formula/switch_sign.py b/pyecsca/ec/formula/switch_sign.py
index e6397a0..95c18be 100644
--- a/pyecsca/ec/formula/switch_sign.py
+++ b/pyecsca/ec/formula/switch_sign.py
@@ -1,5 +1,6 @@
-from typing import Dict, Iterator, List, Any
+from typing import Dict, Iterator, List, Any, Iterable
 from ast import parse
+
 from public import public
 from itertools import chain, combinations
 
@@ -7,7 +8,7 @@ from pyecsca.ec.op import OpType, CodeOp
 from pyecsca.ec.formula.base import Formula
 from pyecsca.ec.formula.graph import FormulaGraph, ConstantNode, CodeOpNode, CodeFormula
 from pyecsca.ec.point import Point
-from pyecsca.ec.mod import Mod, mod
+from pyecsca.ec.mod import mod
 
 
 @public
@@ -20,14 +21,20 @@ def generate_switched_formulas(formula: Formula, rename=True) -> Iterator[CodeFo
             continue
 
 
-def subnode_lists(graph: FormulaGraph):
+def switch_signs(formula: Formula, rename=True) -> List[CodeFormula]:
+    return list(generate_switched_formulas(formula, rename))
+
+
+def subnode_lists(graph: FormulaGraph) -> List[List[CodeOpNode]]:
     return powerlist(filter(lambda x: x not in graph.roots and x.is_sub, graph.nodes))
 
 
-def switch_sign(graph: FormulaGraph, node_combination) -> FormulaGraph:
+def switch_sign(
+    graph: FormulaGraph, node_combination: List[CodeOpNode]
+) -> FormulaGraph:
     nodes_i = [graph.node_index(node) for node in node_combination]
     graph = graph.deepcopy()
-    node_combination = {graph.nodes[node_i] for node_i in nodes_i}
+    node_combination = [graph.nodes[node_i] for node_i in nodes_i]  # type: ignore
     output_signs = {out: 1 for out in graph.output_names}
 
     queue = []
@@ -131,6 +138,6 @@ def change_sides(node):
     )
 
 
-def powerlist(iterable: Iterator) -> List:
+def powerlist(iterable: Iterable) -> List:
     s = list(iterable)
     return list(chain.from_iterable(combinations(s, r) for r in range(len(s) + 1)))
diff --git a/test/ec/test_formula.py b/test/ec/test_formula.py
index e5fce6e..137f243 100644
--- a/test/ec/test_formula.py
+++ b/test/ec/test_formula.py
@@ -1,12 +1,13 @@
 import pickle
 from operator import itemgetter
 from typing import Tuple
+import multiprocessing as mp
 
 import pytest
 from sympy import FF, symbols
 from importlib_resources import files, as_file
 import pyecsca.ec
-from pyecsca.ec.formula.expand import expand_formula_set
+from pyecsca.ec.formula.expand import expand_formula_set, expand_formula_set_parallel
 from pyecsca.ec.formula.fliparoo import generate_fliparood_formulas
 from pyecsca.ec.formula.graph import rename_ivs
 from pyecsca.ec.formula.metrics import (
@@ -111,18 +112,27 @@ def test_assumptions(secp128r1, mdbl):
 
 @pytest.mark.parametrize(
     "formula,category,curve,coords",
-    [("dbl-1998-hnm", "secg", "secp128r1", "jacobian"),
-     ("add-2015-rcb", "secg", "secp128r1", "projective"),
-     ("dbl-1987-m-2", "other", "Curve25519", "xz"),
-     ("add-20090311-hwcd", "other", "E-222", "projective")]
+    [
+        ("dbl-1998-hnm", "secg", "secp128r1", "jacobian"),
+        ("add-2015-rcb", "secg", "secp128r1", "projective"),
+        ("dbl-1987-m-2", "other", "Curve25519", "xz"),
+        ("add-20090311-hwcd", "other", "E-222", "projective"),
+    ],
 )
 def test_eval(formula, category, curve, coords):
     params = get_params(category, curve, coords)
     f = params.curve.coordinate_model.formulas[formula]
 
     points_aff = [params.curve.affine_random() for _ in range(f.num_inputs)]
-    points = [point.to_model(params.curve.coordinate_model, params.curve) for point in points_aff]
-    expected = params.curve.affine_double(*points_aff) if f.shortname == "dbl" else params.curve.affine_add(*points_aff)
+    points = [
+        point.to_model(params.curve.coordinate_model, params.curve)
+        for point in points_aff
+    ]
+    expected = (
+        params.curve.affine_double(*points_aff)
+        if f.shortname == "dbl"
+        else params.curve.affine_add(*points_aff)
+    )
 
     res = f(
         params.curve.prime,
@@ -418,9 +428,12 @@ def library_formula_params(request) -> Tuple[CodeFormula, DomainParameters]:
     name, model, coords_name, param_spec, formula_type = request.param
     model = model()
     coordinate_model = model.coordinates[coords_name]
-    with as_file(files(pyecsca.ec).joinpath("data", "formulas", name)) as meta_path, as_file(
-        files(pyecsca.ec).joinpath("data", "formulas", name + ".op3")
-    ) as op3_path:
+    with (
+        as_file(files(pyecsca.ec).joinpath("data", "formulas", name)) as meta_path,
+        as_file(
+            files(pyecsca.ec).joinpath("data", "formulas", name + ".op3")
+        ) as op3_path,
+    ):
         formula = formula_type(meta_path, op3_path, name, coordinate_model).to_code()
     params = get_params(*param_spec, coords_name)
     return formula, params
@@ -543,3 +556,18 @@ def test_formula_correctness(library_formula_params):
 def test_formula_expand(add):
     res = expand_formula_set({add})
     assert len(res) > 1
+
+
+def test_formula_expand_parallel(add):
+    res = expand_formula_set_parallel({add})
+    assert len(res) > 1
+    other = expand_formula_set({add})
+    assert res == other
+
+
+def test_formula_expand_deterministic(add):
+    ctx = mp.get_context("spawn")
+    with ctx.Pool(4) as p:
+        results = p.map(expand_formula_set, [{add}] * 10)
+    for result in results:
+        assert result == results[0]