2 files changed, 134 insertions, 17 deletions

diff --git a/test/sca/test_epa.py b/test/sca/test_epa.py
index 6f4afbb..e564177 100644
--- a/test/sca/test_epa.py
+++ b/test/sca/test_epa.py
@@ -3,20 +3,21 @@ from functools import partial
 import pytest
 
 from pyecsca.ec.coordinates import EFDCoordinateModel
-from pyecsca.ec.mult import LTRMultiplier, CombMultiplier
+from pyecsca.ec.mult import LTRMultiplier, CombMultiplier, WindowNAFMultiplier
 from pyecsca.sca.re.rpa import multiple_graph
 from pyecsca.sca.re.epa import errors_out
 
 
 def test_errors_out(secp128r1):
-    ctx, out = multiple_graph(
+    precomp_ctx, full_ctx, out = multiple_graph(
         scalar=15,
         params=secp128r1,
         mult_class=LTRMultiplier,
         mult_factory=LTRMultiplier,
     )
     res_empty_checks = errors_out(
-        ctx,
+        precomp_ctx,
+        full_ctx,
         out,
         check_funcs={},
         check_condition="all",
@@ -28,7 +29,8 @@ def test_errors_out(secp128r1):
         return k == 6
 
     res_check_k_add = errors_out(
-        ctx,
+        precomp_ctx,
+        full_ctx,
         out,
         check_funcs={"add": add_check},
         check_condition="all",
@@ -40,7 +42,8 @@ def test_errors_out(secp128r1):
         return k == 15
 
     res_check_k_affine = errors_out(
-        ctx,
+        precomp_ctx,
+        full_ctx,
         out,
         check_funcs={"affine": affine_check},
         check_condition="all",
@@ -50,7 +53,7 @@ def test_errors_out(secp128r1):
 
 
 def test_errors_out_comb(secp128r1):
-    ctx, out = multiple_graph(
+    precomp_ctx, full_ctx, out = multiple_graph(
         scalar=15,
         params=secp128r1,
         mult_class=CombMultiplier,
@@ -61,7 +64,8 @@ def test_errors_out_comb(secp128r1):
         return k == 2**64
 
     res_check_k_affine_precomp = errors_out(
-        ctx,
+        precomp_ctx,
+        full_ctx,
         out,
         check_funcs={"affine": affine_check_comb},
         check_condition="all",
@@ -70,7 +74,8 @@ def test_errors_out_comb(secp128r1):
     assert res_check_k_affine_precomp
 
     res_check_k_no_affine_precomp = errors_out(
-        ctx,
+        precomp_ctx,
+        full_ctx,
         out,
         check_funcs={"affine": affine_check_comb},
         check_condition="all",
@@ -81,7 +86,7 @@ def test_errors_out_comb(secp128r1):
 
 @pytest.mark.skip(reason="Debug only")
 def test_memory_consumption(secp128r1):
-    ctx, out = multiple_graph(
+    precomp_ctx, full_ctx, out = multiple_graph(
         scalar=2**127 + 12127486321,
         params=secp128r1,
         mult_class=LTRMultiplier,
@@ -92,7 +97,114 @@ def test_memory_consumption(secp128r1):
 
         sizer = Asizer()
         sizer.exclude_types(EFDCoordinateModel)
-        print(sizer.asized(ctx, detail=2).format())
+        print(sizer.asized(precomp_ctx, detail=2).format())
+        print(sizer.asized(full_ctx, detail=2).format())
         print(sizer.asized(out, detail=2).format())
     except ImportError:
         pass
+
+
+def test_errors_out_precomp(secp128r1):
+    precomp_ctx, full_ctx, out = multiple_graph(
+        scalar=15,
+        params=secp128r1,
+        mult_class=WindowNAFMultiplier,
+        mult_factory=partial(WindowNAFMultiplier, width=3, complete=False, precompute_negation=True),
+    )
+
+    affine_multiples = []
+
+    def affine_check(k):
+        affine_multiples.append(k)
+        return False
+
+    add_multiples = []
+
+    def add_check(k, l):  # noqa
+        add_multiples.append((k, l))
+        return False
+
+    # Here we check "all" but only during the precomp.
+    errors_out(
+        precomp_ctx,
+        full_ctx,
+        out,
+        check_funcs={"affine": affine_check, "add": add_check},
+        check_condition="all",
+        precomp_to_affine=True,
+        use_init=True,
+        use_multiply=False,
+    )
+    assert set(affine_multiples) == set(precomp_ctx.precomp.keys())
+    assert set(add_multiples) == {(1, 2), (3, 2)}
+
+    # Here we check all, during both precomp and final multiply.
+    affine_multiples = []
+    add_multiples = []
+
+    errors_out(
+        precomp_ctx,
+        full_ctx,
+        out,
+        check_funcs={"affine": affine_check, "add": add_check},
+        check_condition="all",
+        precomp_to_affine=True,
+        use_init=True,
+        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()) | {full_ctx.points[out]}
+    # 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)}
+
+    # Now check just the multiply with all.
+    affine_multiples = []
+    add_multiples = []
+
+    errors_out(
+        precomp_ctx,
+        full_ctx,
+        out,
+        check_funcs={"affine": affine_check, "add": add_check},
+        check_condition="all",
+        precomp_to_affine=True,
+        use_init=False,
+        use_multiply=True,
+    )
+    # Only the final result should be converted to affine, because we ignore precomp.
+    assert set(affine_multiples) == {full_ctx.points[out]}
+    # Only the single add in the multiply should be checked.
+    assert set(add_multiples) == {(16, -1)}
+
+    # Now check just the multiply with necessary only.
+    affine_multiples = []
+    add_multiples = []
+
+    errors_out(
+        precomp_ctx,
+        full_ctx,
+        out,
+        check_funcs={"affine": affine_check, "add": add_check},
+        check_condition="necessary",
+        precomp_to_affine=True,
+        use_init=False,
+        use_multiply=True,
+    )
+    # Only the final result should be converted to affine, because we ignore precomp.
+    assert set(affine_multiples) == {full_ctx.points[out]}
+    # Only the single add in the multiply should be checked.
+    assert set(add_multiples) == {(16, -1)}
+
+    # Doing use_init = False and use_multiply = False does not make sense.
+    with pytest.raises(ValueError):
+        errors_out(
+            precomp_ctx,
+            full_ctx,
+            out,
+            check_funcs={"affine": affine_check, "add": add_check},
+            check_condition="all",
+            precomp_to_affine=True,
+            use_init=False,
+            use_multiply=False,
+        )
diff --git a/test/sca/test_rpa.py b/test/sca/test_rpa.py
index 8ba719c..01293af 100644
--- a/test/sca/test_rpa.py
+++ b/test/sca/test_rpa.py
@@ -83,13 +83,18 @@ def test_multiples(rpa_params):
 
 
 def test_multiples_no_init(rpa_params):
-    multiples = multiples_computed(
-        78699, rpa_params, LTRMultiplier,
-        lambda add, dbl, *args, **kwargs: LTRMultiplier(
-            add, dbl, None, False, AccumulationOrder.PeqPR, True, False
-        ), False, True
-    )
-    assert multiples
+    wnaf = partial(WindowNAFMultiplier, width=4)
+
+    multiples = multiples_computed(78699, rpa_params, WindowNAFMultiplier, wnaf, False, True)
+    assert multiples & {1, 2, 3, 5, 7, 9} == set()
+    assert 78699 in multiples
+
+
+def test_multiples_no_multiply(rpa_params):
+    wnaf = partial(WindowNAFMultiplier, width=4)
+
+    multiples = multiples_computed(78699, rpa_params, WindowNAFMultiplier, wnaf, True, False)
+    assert multiples == {1, 2, 3, 5, 7, 9}
 
 
 def test_multiples_bnaf(rpa_params):