Add scalarmult strucutre plot tests.

3 files changed, 190 insertions, 28 deletions

diff --git a/pyecsca/sca/re/epa.py b/pyecsca/sca/re/epa.py
index 60398f0..ba1d520 100644
--- a/pyecsca/sca/re/epa.py
+++ b/pyecsca/sca/re/epa.py
@@ -1,8 +1,13 @@
 """
 Provides functionality inspired by the Exceptional Procedure Attack [EPA]_.
 """
+
 from typing import Callable, Literal, Union, Optional
 
+import matplotlib.pyplot as plt
+import networkx as nx
+
+
 from public import public
 
 from pyecsca.ec.point import Point
@@ -74,9 +79,13 @@ def graph_to_check_inputs(
         if use_init and use_multiply:
             points = _necessary(full_ctx, affine_points)
         elif use_init:
-            points = _necessary(full_ctx, affine_points) & set(precomp_ctx.points.keys())
+            points = _necessary(full_ctx, affine_points) & set(
+                precomp_ctx.points.keys()
+            )
         elif use_multiply:
-            points = _necessary(full_ctx, affine_points) - set(precomp_ctx.points.keys())
+            points = _necessary(full_ctx, affine_points) - set(
+                precomp_ctx.points.keys()
+            )
     else:
         raise ValueError("check_condition must be 'all' or 'necessary'")
     # Special case the "to affine" transform and checks
@@ -91,7 +100,9 @@ def graph_to_check_inputs(
 
     for point in points:
         formula = full_ctx.formulas[point]
-        if not formula or (check_formulas is not None and formula not in check_formulas):
+        if not formula or (
+            check_formulas is not None and formula not in check_formulas
+        ):
             # Skip input point or infty point (they magically appear and do not have an origin formula)
             continue
         inputs = tuple(map(get_point, full_ctx.parents[point]))
@@ -101,8 +112,77 @@ def graph_to_check_inputs(
 
 
 @public
+def graph_plot(
+    precomp_ctx: MultipleContext,
+    full_ctx: MultipleContext,
+    out: Point,
+) -> plt.Figure:
+    """
+    Plot the computation graph, highlighting necessary points and precomputed points.
+
+    :param precomp_ctx: The context containing the points and formulas (precomputation phase).
+    :param full_ctx: The context containing the points and formulas (full computation).
+    :param out: The output point of the computation.
+    :return: The matplotlib figure object representing the graph.
+    """
+    graph = full_ctx.to_networkx()
+
+    for layer, nodes in enumerate(nx.topological_generations(graph)):
+        for node in nodes:
+            graph.nodes[node]["layer"] = layer
+    for node in graph.nodes():
+        graph.nodes[node]["necessary"] = False
+    queue = {out}
+    while queue:
+        node = queue.pop()
+        graph.nodes[node]["necessary"] = True
+        for n in graph.predecessors(node):
+            queue.add(n)
+    fig, ax = plt.subplots(figsize=(60, 10))
+    pos = nx.multipartite_layout(graph, subset_key="layer")
+    for point, p in pos.items():
+        p[0] *= 0.15
+        if not graph.nodes[point]["necessary"]:
+            p[1] += 0.01
+        if point in precomp_ctx.points.keys():
+            if graph.nodes[point]["precomp"]:
+                p[1] -= 0.01
+
+    colors = []
+    for point in graph.nodes():
+        if graph.nodes[point]["necessary"]:
+            color = "#202080"
+        else:
+            color = "#802020"
+        if point in precomp_ctx.points.keys():
+            color = "#208020"
+            if graph.nodes[point]["precomp"]:
+                color = "#30a0a0"
+        colors.append(color)
+
+    nx.draw_networkx_nodes(
+        graph, pos, ax=ax, node_color=colors, node_size=500, margins=[0.1, 0.1]
+    )
+    nx.draw_networkx_edges(graph, pos, ax=ax, connectionstyle="arc3,rad=0.05")
+    nx.draw_networkx_edge_labels(
+        graph,
+        pos,
+        ax=ax,
+        edge_labels={(u, v): graph.edges[u, v]["formula"] for u, v in graph.edges()},
+    )
+    for p in pos.values():
+        p[1] += 0.003
+    nx.draw_networkx_labels(
+        graph, pos, ax=ax, labels={n: graph.nodes[n]["multiple"] for n in graph.nodes()}
+    )
+    fig.tight_layout()
+    return fig
+
+
+@public
 def evaluate_checks(
-    check_funcs: dict[str, Union[Callable[[int, int], bool], Callable[[int], bool]]], check_inputs: dict[str, set[tuple[int, ...]]]
+    check_funcs: dict[str, Union[Callable[[int, int], bool], Callable[[int], bool]]],
+    check_inputs: dict[str, set[tuple[int, ...]]],
 ) -> bool:
     """
     Evaluate the checks for each formula type based on the provided functions and inputs.
@@ -159,5 +239,13 @@ def errors_out(
     .. note::
         The scalar multiplier must not short-circuit.
     """
-    formula_checks = graph_to_check_inputs(precomp_ctx, full_ctx, out, check_condition, precomp_to_affine, use_init, use_multiply)
+    formula_checks = graph_to_check_inputs(
+        precomp_ctx,
+        full_ctx,
+        out,
+        check_condition,
+        precomp_to_affine,
+        use_init,
+        use_multiply,
+    )
     return evaluate_checks(check_funcs, formula_checks)
diff --git a/pyecsca/sca/re/rpa.py b/pyecsca/sca/re/rpa.py
index 0d7fd6c..2fd5973 100644
--- a/pyecsca/sca/re/rpa.py
+++ b/pyecsca/sca/re/rpa.py
@@ -18,6 +18,8 @@ from typing import (
     Tuple,
 )
 
+import networkx as nx
+
 from sympy import FF, sympify, Poly, symbols
 
 from pyecsca.ec.error import NonInvertibleError
@@ -58,6 +60,8 @@ class MultipleContext(Context):
     """The mapping of points to the formula types they are a result of."""
     precomp: MutableMapping[int, Point]
     """The mapping of precomputed multiples to the points they represent."""
+    result: Optional[Point]
+    """The resulting point of the computation."""
     inside: List[Action]
     """Whether we are inside a scalarmult/precomp action."""
     keep_base: bool
@@ -75,6 +79,7 @@ class MultipleContext(Context):
         self.formulas = {}
         self.precomp = {}
         self.inside = []
+        self.result = None
         self.keep_base = keep_base
         self._track_precomp = track_precomp
         self._track_scalarmult = track_scalarmult
@@ -129,6 +134,8 @@ class MultipleContext(Context):
             self.inside.remove(action)
             if isinstance(action, PrecomputationAction):
                 self.precomp.update(action.result)
+            if isinstance(action, ScalarMultiplicationAction):
+                self.result = action.result
         if isinstance(action, FormulaAction) and self.inside:
             action = cast(FormulaAction, action)
             shortname = action.formula.shortname
@@ -181,6 +188,35 @@ class MultipleContext(Context):
     def __repr__(self):
         return f"{self.__class__.__name__}({self.base!r}, multiples={self.points.values()!r})"
 
+    def to_networkx(self) -> nx.DiGraph:
+        """
+        Convert the context to a NetworkX graph.
+
+        The nodes represent points, with attributes:
+            - multiple: The multiple of the base point.
+            - formula: The formula used to compute the point.
+            - precomp: Whether the point was the result of precomputation.
+            - result: Whether the point is the final result of the computation.
+
+        The edges represent the computation steps, with attribute:
+            - formula: The formula used to compute the child point from the parent point.
+
+        :return: A NetworkX DiGraph representing the computation.
+        """
+        graph = nx.DiGraph()
+        for point, multiple in self.points.items():
+            graph.add_node(
+                point,
+                multiple=multiple,
+                formula=self.formulas[point],
+                precomp=point in self.precomp.values(),
+                result=point == self.result,
+            )
+        for point, parents in self.parents.items():
+            for parent in parents:
+                graph.add_edge(parent, point, formula=self.formulas[point])
+        return graph
+
 
 @public
 def rpa_point_0y(params: DomainParameters) -> Optional[Point]:
@@ -456,7 +492,8 @@ def multiple_graph(
     :param mult_factory: A callable that takes the formulas and instantiates the multiplier.
     :param dlog: Make an assumption that the symbolic input point is the `dlog` multiple of the base point.
                  This is necessary if the multiplier does computation with the base point.
-    :return: The context with the computed multiples and the resulting point.
+    :return: The context with the computed multiples during precomputation, the context with the computed multiples
+             during the full computation, and the resulting point.
     """
     params = fake_params(params)
     mult = cached_fake_mult(mult_class, mult_factory, params)
@@ -497,10 +534,11 @@ def multiples_from_graph(
     use_multiply: bool = True,
 ):
     """
+    Compute the multiples computed for a given scalar and multiplier from the multiple graph.
 
-    :param precomp_ctx:
-    :param full_ctx:
-    :param out:
+    :param precomp_ctx: The context containing the points and formulas (precomputation phase).
+    :param full_ctx: The context containing the points and formulas (full computation).
+    :param out: The output point of the computation.
     :param kind: The kind of multiples to return. Can be one of "all", "input", "necessary", or "precomp+necessary".
     :param use_init: Whether to consider the point multiples that happen in scalarmult initialization.
     :param use_multiply: Whether to consider the point multiples that happen in scalarmult multiply (after initialization).
diff --git a/test/sca/test_epa.py b/test/sca/test_epa.py
index cab2f03..a6ac542 100644
--- a/test/sca/test_epa.py
+++ b/test/sca/test_epa.py
@@ -4,9 +4,13 @@ from functools import partial
 import pytest
 
 from pyecsca.ec.coordinates import EFDCoordinateModel
+from pyecsca.ec.curve import EllipticCurve
+from pyecsca.ec.mod import mod
+from pyecsca.ec.model import ShortWeierstrassModel
+from pyecsca.ec.params import Point, InfinityPoint
 from pyecsca.ec.mult import *
 from pyecsca.sca.re.rpa import multiple_graph, multiples_from_graph
-from pyecsca.sca.re.epa import errors_out, graph_to_check_inputs
+from pyecsca.sca.re.epa import errors_out, graph_to_check_inputs, graph_plot
 
 
 def test_errors_out(secp128r1):
@@ -139,7 +143,7 @@ def test_errors_out_precomp(secp128r1):
         use_multiply=False,
     )
     assert set(affine_multiples) == set(precomp_ctx.precomp.keys())
-    assert set(add_multiples) == {(1, 2), (3, 2)}
+    assert set(add_multiples) == {(1, 2)}
 
     # Here we check all, during both precomp and final multiply.
     affine_multiples = []
@@ -161,7 +165,7 @@ def test_errors_out_precomp(secp128r1):
     }
     # The add multiples should be the same as before, plus any inputs to add that happened
     # during the final multiply, there is only one, rest are doubles.
-    assert set(add_multiples) == {(1, 2), (3, 2), (16, -1)}
+    assert set(add_multiples) == {(1, 2), (16, -1)}
 
     # Now check just the multiply with all.
     affine_multiples = []
@@ -320,6 +324,34 @@ def mult(secp128r1, request):
     return mult_class, partial(mult_class, **mult_kwargs)
 
 
+@pytest.fixture()
+def toy_params():
+    model = ShortWeierstrassModel()
+    coords = model.coordinates["projective"]
+    p = 0xCB5E1D94A6168511
+    a = mod(0xB166CA7D2DFBF69F, p)
+    b = mod(0x855BB40CB6937C4B, p)
+    gx = mod(0x253B2638BD13D6F4, p)
+    gy = mod(0x1E91A1A182287E71, p)
+
+    infty = InfinityPoint(coords)
+    g = Point(coords, X=gx, Y=gy, Z=mod(1, p))
+    curve = EllipticCurve(model, coords, p, infty, dict(a=a, b=b))
+    return DomainParameters(curve, g, 0xCB5E1D94601A3AC5, 1)
+
+
+def test_plot(toy_params, mult, plot_path):
+    mult_class, mult_factory = mult
+    precomp_ctx, full_ctx, out = multiple_graph(
+        scalar=15546875464546546545644687 % toy_params.order,
+        params=toy_params,
+        mult_class=mult_class,
+        mult_factory=mult_factory,
+    )
+    fig = graph_plot(precomp_ctx, full_ctx, out)
+    fig.savefig(str(plot_path()) + ".png")
+
+
 def test_independent_check_inputs(secp128r1, mult):
     """
     Check that the set of check inputs is constant if (use_init = True, use_multiply = False) for all scalars
@@ -354,13 +386,22 @@ def test_independent_check_inputs(secp128r1, mult):
                     last_check_inputs = check_inputs
 
 
-@pytest.mark.parametrize("check_condition,precomp_to_affine,multiples_kind", [
-    ("all", True, "all"),
-    ("all", False, "all"),
-    ("necessary", True, "precomp+necessary"),
-    ("necessary", False, "necessary"),
-])
-def test_consistency_multiples(secp128r1, mult, check_condition, precomp_to_affine, multiples_kind):
+@pytest.mark.parametrize(
+    "check_condition,precomp_to_affine,multiples_kind",
+    [
+        ("all", True, "all"),
+        ("all", False, "all"),
+        ("necessary", True, "precomp+necessary"),
+        ("necessary", False, "necessary"),
+    ],
+)
+def test_consistency_multiples(
+    secp128r1,
+    mult,
+    check_condition,
+    precomp_to_affine,
+    multiples_kind,
+):
     """
     Test consistency between the graph_to_check_inputs and multiples_computed functions for the same error model
     """
@@ -379,21 +420,19 @@ def test_consistency_multiples(secp128r1, mult, check_condition, precomp_to_affi
             out,
             check_condition=check_condition,
             precomp_to_affine=precomp_to_affine,
-            use_init=True,
-            use_multiply=True,
         )
         # Now map the check inputs to the set of multiples they cover
         multiples_from_check_inputs = set()
-        for k, in check_inputs.get("neg", []):
+        for (k,) in check_inputs.get("neg", []):
             multiples_from_check_inputs.add(k)
             multiples_from_check_inputs.add(-k)
-        for k, in check_inputs.get("affine", []):
+        for (k,) in check_inputs.get("affine", []):
             multiples_from_check_inputs.add(k)
         for k, l in check_inputs.get("add", []):
             multiples_from_check_inputs.add(k)
             multiples_from_check_inputs.add(l)
             multiples_from_check_inputs.add(k + l)
-        for k, in check_inputs.get("dbl", []):
+        for (k,) in check_inputs.get("dbl", []):
             multiples_from_check_inputs.add(k)
             multiples_from_check_inputs.add(2 * k)
         # Multiples computed removes the zero
@@ -401,9 +440,6 @@ def test_consistency_multiples(secp128r1, mult, check_condition, precomp_to_affi
 
         # Now compute the multiples via the other function to compare.
         multiples = multiples_from_graph(
-            precomp_ctx,
-            full_ctx,
-            out,
-            kind=multiples_kind
+            precomp_ctx, full_ctx, out, kind=multiples_kind
         )
         assert multiples_from_check_inputs == multiples