Fixup necessary computation in EPA.

3 files changed, 9 insertions, 4 deletions

diff --git a/pyecsca/sca/re/epa.py b/pyecsca/sca/re/epa.py
index 4ef1b74..37c0b98 100644
--- a/pyecsca/sca/re/epa.py
+++ b/pyecsca/sca/re/epa.py
@@ -57,7 +57,7 @@ def graph_to_check_inputs(
         affine_points = {out}
 
     def _necessary(ctx, for_what):
-        res = {out}
+        res = {*for_what}
         queue = {*for_what}
         while queue:
             point = queue.pop()
diff --git a/pyecsca/sca/re/rpa.py b/pyecsca/sca/re/rpa.py
index 2fd5973..12b8371 100644
--- a/pyecsca/sca/re/rpa.py
+++ b/pyecsca/sca/re/rpa.py
@@ -59,7 +59,12 @@ class MultipleContext(Context):
     formulas: MutableMapping[Point, str]
     """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."""
+    """
+    The mapping of precomputed points the scalar multiplier stored.
+
+    .. note::
+        The keys are the indices used by the scalar multiplier, not (necessarily) the multiples.
+    """
     result: Optional[Point]
     """The resulting point of the computation."""
     inside: List[Action]
diff --git a/test/sca/test_epa.py b/test/sca/test_epa.py
index 0975e35..89745cc 100644
--- a/test/sca/test_epa.py
+++ b/test/sca/test_epa.py
@@ -146,7 +146,7 @@ def test_errors_out_precomp(secp128r1):
         use_init=True,
         use_multiply=False,
     )
-    assert set(affine_multiples) == set(precomp_ctx.precomp.keys())
+    assert set(affine_multiples) == set(precomp_ctx.points[pt] for pt in precomp_ctx.precomp.values())
     assert set(add_multiples) == {(1, 2)}
 
     # Here we check all, during both precomp and final multiply.
@@ -164,7 +164,7 @@ def test_errors_out_precomp(secp128r1):
         use_multiply=True,
     )
     # There should be all of the results of the precomp, plus the final multiply result.
-    assert set(affine_multiples) == set(precomp_ctx.precomp.keys()) | {
+    assert set(affine_multiples) == set(precomp_ctx.points[pt] for pt in precomp_ctx.precomp.values()) | {
         full_ctx.points[out]
     }
     # The add multiples should be the same as before, plus any inputs to add that happened