9 files changed, 125 insertions, 41 deletions

diff --git a/pyecsca/ec/countermeasures.py b/pyecsca/ec/countermeasures.py
index d6af799..5c52f35 100644
--- a/pyecsca/ec/countermeasures.py
+++ b/pyecsca/ec/countermeasures.py
@@ -31,6 +31,7 @@ class ScalarMultiplierCountermeasure(ABC):
     def init(self, params: DomainParameters, point: Point):
         self.params = params
         self.point = point
+        self.mult.init(self.params, self.point)
 
     @abstractmethod
     def multiply(self, scalar: int) -> Point:
@@ -48,7 +49,6 @@ class GroupScalarRandomization(ScalarMultiplierCountermeasure):
     def multiply(self, scalar: int) -> Point:
         if self.params is None or self.point is None:
             raise ValueError("Not initialized.")
-        self.mult.init(self.params, self.point)
         order = self.params.order
         mask = int(Mod.random(1 << self.rand_bits))
         masked_scalar = scalar + mask * order
@@ -66,7 +66,6 @@ class AdditiveSplitting(ScalarMultiplierCountermeasure):
     def multiply(self, scalar: int) -> Point:
         if self.params is None or self.point is None:
             raise ValueError("Not initialized.")
-        self.mult.init(self.params, self.point)
 
         order = self.params.order
         r = Mod.random(order)
@@ -92,7 +91,6 @@ class MultiplicativeSplitting(ScalarMultiplierCountermeasure):
     def multiply(self, scalar: int) -> Point:
         if self.params is None or self.point is None:
             raise ValueError("Not initialized.")
-        self.mult.init(self.params, self.point)
         r = Mod.random(1 << self.rand_bits)
         R = self.mult.multiply(int(r))
 
@@ -112,19 +110,18 @@ class EuclideanSplitting(ScalarMultiplierCountermeasure):
     def multiply(self, scalar: int) -> Point:
         if self.params is None or self.point is None:
             raise ValueError("Not initialized.")
+
         order = self.params.order
         half_bits = order.bit_length() // 2
         r = Mod.random(1 << half_bits)
-        self.mult.init(self.params, self.point)
         R = self.mult.multiply(int(r))
 
         k1 = scalar % int(r)
         k2 = scalar // int(r)
+        T = self.mult.multiply(k1)
+
         self.mult.init(self.params, R)
         S = self.mult.multiply(k2)
-
-        self.mult.init(self.params, self.point)
-        T = self.mult.multiply(k1)
         if self.add is None:
             return self.mult._add(S, T)  # noqa: This is OK.
         else:
diff --git a/pyecsca/ec/mult/base.py b/pyecsca/ec/mult/base.py
index 469be39..1c51271 100644
--- a/pyecsca/ec/mult/base.py
+++ b/pyecsca/ec/mult/base.py
@@ -48,7 +48,7 @@ class ScalarMultiplicationAction(ResultAction):
 
 
 @public
-class PrecomputationAction(Action):
+class PrecomputationAction(ResultAction):
     """A precomputation of a point in scalar multiplication."""
 
     params: DomainParameters
@@ -239,6 +239,11 @@ class ScalarMultiplier(ABC):
 
 
 @public
+class PrecompMultiplier(ScalarMultiplier, ABC):
+    pass
+
+
+@public
 class AccumulatorMultiplier(ScalarMultiplier, ABC):
     """
     A scalar multiplication algorithm mix-in class for a multiplier that accumulates.
diff --git a/pyecsca/ec/mult/comb.py b/pyecsca/ec/mult/comb.py
index 9557fe6..35d1678 100644
--- a/pyecsca/ec/mult/comb.py
+++ b/pyecsca/ec/mult/comb.py
@@ -12,7 +12,7 @@ from pyecsca.ec.mult import (
     ProcessingDirection,
     AccumulationOrder,
     PrecomputationAction,
-    ScalarMultiplicationAction,
+    ScalarMultiplicationAction, PrecompMultiplier,
 )
 from pyecsca.ec.params import DomainParameters
 from pyecsca.ec.point import Point
@@ -20,7 +20,7 @@ from pyecsca.ec.scalar import convert_base
 
 
 @public
-class BGMWMultiplier(AccumulatorMultiplier, ScalarMultiplier):
+class BGMWMultiplier(AccumulatorMultiplier, PrecompMultiplier, ScalarMultiplier):
     """
     Brickell, Gordon, McCurley and Wilson (BGMW) scalar multiplier,
     or rather, its one parametrization.
@@ -86,7 +86,7 @@ class BGMWMultiplier(AccumulatorMultiplier, ScalarMultiplier):
         return f"{self.__class__.__name__}({', '.join(map(str, self.formulas.values()))}, short_circuit={self.short_circuit}, width={self.width}, direction={self.direction.name}, accumulation_order={self.accumulation_order.name})"
 
     def init(self, params: DomainParameters, point: Point):
-        with PrecomputationAction(params, point):
+        with PrecomputationAction(params, point) as action:
             super().init(params, point)
             d = ceil(params.order.bit_length() / self.width)
             self._points = {}
@@ -96,6 +96,7 @@ class BGMWMultiplier(AccumulatorMultiplier, ScalarMultiplier):
                 if i != d - 1:
                     for _ in range(self.width):
                         current_point = self._dbl(current_point)
+            action.exit(self._points)
 
     def multiply(self, scalar: int) -> Point:
         if not self._initialized:
@@ -126,7 +127,7 @@ class BGMWMultiplier(AccumulatorMultiplier, ScalarMultiplier):
 
 
 @public
-class CombMultiplier(AccumulatorMultiplier, ScalarMultiplier):
+class CombMultiplier(AccumulatorMultiplier, PrecompMultiplier, ScalarMultiplier):
     """
     Comb multiplier.
 
@@ -179,7 +180,7 @@ class CombMultiplier(AccumulatorMultiplier, ScalarMultiplier):
         return f"{self.__class__.__name__}({', '.join(map(str, self.formulas.values()))}, short_circuit={self.short_circuit}, width={self.width}, accumulation_order={self.accumulation_order.name})"
 
     def init(self, params: DomainParameters, point: Point):
-        with PrecomputationAction(params, point):
+        with PrecomputationAction(params, point) as action:
             super().init(params, point)
             d = ceil(params.order.bit_length() / self.width)
             base_points = {}
@@ -198,6 +199,7 @@ class CombMultiplier(AccumulatorMultiplier, ScalarMultiplier):
                 self._points[j] = points[0]
                 for other in points[1:]:
                     self._points[j] = self._accumulate(self._points[j], other)
+            action.exit(self._points)
 
     def multiply(self, scalar: int) -> Point:
         if not self._initialized:
diff --git a/pyecsca/ec/mult/fixed.py b/pyecsca/ec/mult/fixed.py
index fdf129d..2c01648 100644
--- a/pyecsca/ec/mult/fixed.py
+++ b/pyecsca/ec/mult/fixed.py
@@ -11,14 +11,14 @@ from pyecsca.ec.mult.base import (
     ProcessingDirection,
     AccumulationOrder,
     PrecomputationAction,
-    ScalarMultiplicationAction,
+    ScalarMultiplicationAction, PrecompMultiplier,
 )
 from pyecsca.ec.params import DomainParameters
 from pyecsca.ec.point import Point
 
 
 @public
-class FullPrecompMultiplier(AccumulatorMultiplier, ScalarMultiplier):
+class FullPrecompMultiplier(AccumulatorMultiplier, PrecompMultiplier, ScalarMultiplier):
     """
     See page 104 of [GECC]_:
 
@@ -91,7 +91,7 @@ class FullPrecompMultiplier(AccumulatorMultiplier, ScalarMultiplier):
         return f"{self.__class__.__name__}({', '.join(map(str, self.formulas.values()))}, short_circuit={self.short_circuit}, accumulation_order={self.accumulation_order.name}, always={self.always}, complete={self.complete})"
 
     def init(self, params: DomainParameters, point: Point):
-        with PrecomputationAction(params, point):
+        with PrecomputationAction(params, point) as action:
             super().init(params, point)
             self._points = {}
             current_point = point
@@ -99,6 +99,7 @@ class FullPrecompMultiplier(AccumulatorMultiplier, ScalarMultiplier):
                 self._points[i] = current_point
                 if i != params.order.bit_length():
                     current_point = self._dbl(current_point)
+            action.exit(self._points)
 
     def _ltr(self, scalar: int) -> Point:
         if self.complete:
diff --git a/pyecsca/ec/mult/naf.py b/pyecsca/ec/mult/naf.py
index 6fc9131..83dc0dc 100644
--- a/pyecsca/ec/mult/naf.py
+++ b/pyecsca/ec/mult/naf.py
@@ -9,7 +9,7 @@ from pyecsca.ec.mult.base import (
     ProcessingDirection,
     AccumulationOrder,
     PrecomputationAction,
-    AccumulatorMultiplier,
+    AccumulatorMultiplier, PrecompMultiplier,
 )
 from pyecsca.ec.formula import (
     AdditionFormula,
@@ -23,7 +23,7 @@ from pyecsca.ec.scalar import naf, wnaf
 
 
 @public
-class BinaryNAFMultiplier(AccumulatorMultiplier, ScalarMultiplier):
+class BinaryNAFMultiplier(AccumulatorMultiplier, PrecompMultiplier, ScalarMultiplier):
     """
     Binary NAF (Non Adjacent Form) multiplier.
 
@@ -82,9 +82,10 @@ class BinaryNAFMultiplier(AccumulatorMultiplier, ScalarMultiplier):
         return f"{self.__class__.__name__}({', '.join(map(str, self.formulas.values()))}, short_circuit={self.short_circuit}, direction={self.direction.name}, accumulation_order={self.accumulation_order.name})"
 
     def init(self, params: DomainParameters, point: Point):
-        with PrecomputationAction(params, point):
+        with PrecomputationAction(params, point) as action:
             super().init(params, point)
             self._point_neg = self._neg(point)
+            action.exit({-1: self._point_neg})
 
     def _ltr(self, scalar_naf: List[int]) -> Point:
         q = copy(self._params.curve.neutral)
@@ -126,7 +127,7 @@ class BinaryNAFMultiplier(AccumulatorMultiplier, ScalarMultiplier):
 
 
 @public
-class WindowNAFMultiplier(AccumulatorMultiplier, ScalarMultiplier):
+class WindowNAFMultiplier(AccumulatorMultiplier, PrecompMultiplier, ScalarMultiplier):
     """
     Window NAF (Non Adjacent Form) multiplier, left-to-right.
 
@@ -195,7 +196,7 @@ class WindowNAFMultiplier(AccumulatorMultiplier, ScalarMultiplier):
         return f"{self.__class__.__name__}({', '.join(map(str, self.formulas.values()))}, short_circuit={self.short_circuit}, width={self.width}, precompute_negation={self.precompute_negation}, accumulation_order={self.accumulation_order.name})"
 
     def init(self, params: DomainParameters, point: Point):
-        with PrecomputationAction(params, point):
+        with PrecomputationAction(params, point) as action:
             super().init(params, point)
             self._points = {}
             self._points_neg = {}
@@ -206,6 +207,7 @@ class WindowNAFMultiplier(AccumulatorMultiplier, ScalarMultiplier):
                 if self.precompute_negation:
                     self._points_neg[2 * i + 1] = self._neg(current_point)
                 current_point = self._add(current_point, double_point)
+            action.exit({**self._points, **self._points_neg})
 
     def multiply(self, scalar: int) -> Point:
         if not self._initialized:
diff --git a/pyecsca/ec/mult/window.py b/pyecsca/ec/mult/window.py
index c200cc5..6fbee24 100644
--- a/pyecsca/ec/mult/window.py
+++ b/pyecsca/ec/mult/window.py
@@ -10,7 +10,7 @@ from pyecsca.ec.mult.base import (
     ScalarMultiplicationAction,
     PrecomputationAction,
     ProcessingDirection,
-    AccumulatorMultiplier,
+    AccumulatorMultiplier, PrecompMultiplier,
 )
 from pyecsca.ec.formula import (
     AdditionFormula,
@@ -28,7 +28,7 @@ from pyecsca.ec.scalar import (
 
 
 @public
-class SlidingWindowMultiplier(AccumulatorMultiplier, ScalarMultiplier):
+class SlidingWindowMultiplier(AccumulatorMultiplier, PrecompMultiplier, ScalarMultiplier):
     """
     Sliding window scalar multiplier.
 
@@ -91,7 +91,7 @@ class SlidingWindowMultiplier(AccumulatorMultiplier, ScalarMultiplier):
         return f"{self.__class__.__name__}({', '.join(map(str, self.formulas.values()))}, short_circuit={self.short_circuit}, width={self.width}, recoding_direction={self.recoding_direction.name}, accumulation_order={self.accumulation_order.name})"
 
     def init(self, params: DomainParameters, point: Point):
-        with PrecomputationAction(params, point):
+        with PrecomputationAction(params, point) as action:
             super().init(params, point)
             self._points = {}
             current_point = point
@@ -99,6 +99,7 @@ class SlidingWindowMultiplier(AccumulatorMultiplier, ScalarMultiplier):
             for i in range(0, 2 ** (self.width - 1)):
                 self._points[2 * i + 1] = current_point
                 current_point = self._add(current_point, double_point)
+            action.exit(self._points)
 
     def multiply(self, scalar: int) -> Point:
         if not self._initialized:
@@ -121,7 +122,7 @@ class SlidingWindowMultiplier(AccumulatorMultiplier, ScalarMultiplier):
 
 
 @public
-class FixedWindowLTRMultiplier(AccumulatorMultiplier, ScalarMultiplier):
+class FixedWindowLTRMultiplier(AccumulatorMultiplier, PrecompMultiplier, ScalarMultiplier):
     """
     Like LTRMultiplier, but m-ary, not binary.
 
@@ -186,7 +187,7 @@ class FixedWindowLTRMultiplier(AccumulatorMultiplier, ScalarMultiplier):
         return f"{self.__class__.__name__}({', '.join(map(str, self.formulas.values()))}, short_circuit={self.short_circuit}, m={self.m}, accumulation_order={self.accumulation_order.name})"
 
     def init(self, params: DomainParameters, point: Point):
-        with PrecomputationAction(params, point):
+        with PrecomputationAction(params, point) as action:
             super().init(params, point)
             double_point = self._dbl(point)
             self._points = {1: point, 2: double_point}
@@ -194,6 +195,7 @@ class FixedWindowLTRMultiplier(AccumulatorMultiplier, ScalarMultiplier):
             for i in range(3, self.m):
                 current_point = self._add(current_point, point)
                 self._points[i] = current_point
+            action.exit(self._points)
 
     def _mult_m(self, point: Point) -> Point:
         if self.m & (self.m - 1) == 0:
@@ -229,7 +231,7 @@ class FixedWindowLTRMultiplier(AccumulatorMultiplier, ScalarMultiplier):
 
 
 @public
-class WindowBoothMultiplier(AccumulatorMultiplier, ScalarMultiplier):
+class WindowBoothMultiplier(AccumulatorMultiplier, PrecompMultiplier, ScalarMultiplier):
     """
 
     :param short_circuit: Whether the use of formulas will be guarded by short-circuit on inputs
@@ -297,7 +299,7 @@ class WindowBoothMultiplier(AccumulatorMultiplier, ScalarMultiplier):
         return f"{self.__class__.__name__}({', '.join(map(str, self.formulas.values()))}, short_circuit={self.short_circuit}, width={self.width}, precompute_negation={self.precompute_negation}, accumulation_order={self.accumulation_order.name})"
 
     def init(self, params: DomainParameters, point: Point):
-        with PrecomputationAction(params, point):
+        with PrecomputationAction(params, point) as actions:
             super().init(params, point)
             double_point = self._dbl(point)
             self._points = {1: point, 2: double_point}
@@ -309,6 +311,10 @@ class WindowBoothMultiplier(AccumulatorMultiplier, ScalarMultiplier):
                 self._points[i] = current_point
                 if self.precompute_negation:
                     self._points_neg[i] = self._neg(current_point)
+            if self.precompute_negation:
+                actions.exit({**self._points, **self._points_neg})
+            else:
+                actions.exit(self._points)
 
     def multiply(self, scalar: int) -> Point:
         if not self._initialized:
diff --git a/pyecsca/sca/re/rpa.py b/pyecsca/sca/re/rpa.py
index 2ed4339..66b4d1e 100644
--- a/pyecsca/sca/re/rpa.py
+++ b/pyecsca/sca/re/rpa.py
@@ -52,37 +52,54 @@ class MultipleContext(Context):
     """The mapping of points to their parent they were computed from."""
     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."""
     inside: bool
+    """Whether we are inside a scalarmult/precomp action."""
+    keep_base: bool
+    """Whether to keep the base point when building upon it."""
 
-    def __init__(self):
+    def __init__(self, keep_base: bool = False):
         self.base = None
         self.points = {}
         self.parents = {}
         self.formulas = {}
+        self.precomp = {}
         self.inside = False
+        self.keep_base = keep_base
 
     def enter_action(self, action: Action) -> None:
         if isinstance(action, (ScalarMultiplicationAction, PrecomputationAction)):
+            self.inside = True
             if self.base:
                 # If we already did some computation with this context try to see if we are building on top of it.
                 if self.base != action.point:
-                    # If we are not building on top of it we have to forget stuff and set a new base and mapping.
-                    self.base = action.point
-                    self.neutral = action.params.curve.neutral
-                    self.points = {self.base: 1, self.neutral: 0}
-                    self.parents = {self.base: [], self.neutral: []}
-                    self.formulas = {self.base: "", self.neutral: ""}
+                    # We are not doing the same point, but maybe we are doing a multiple of it.
+                    if action.point in self.points and self.keep_base:
+                        # We are doing a multiple of the base and were asked to keep it.
+                        pass
+                    else:
+                        # We are not building on top of it (or were asked to forget).
+                        # We have to forget stuff and set a new base and mapping.
+                        self.base = action.point
+                        self.neutral = action.params.curve.neutral
+                        self.points = {self.base: 1, self.neutral: 0}
+                        self.parents = {self.base: [], self.neutral: []}
+                        self.formulas = {self.base: "", self.neutral: ""}
+                        self.precomp = {}
             else:
                 self.base = action.point
                 self.neutral = action.params.curve.neutral
                 self.points = {self.base: 1, self.neutral: 0}
                 self.parents = {self.base: []}
                 self.formulas = {self.base: "", self.neutral: ""}
-            self.inside = True
+                self.precomp = {}
 
     def exit_action(self, action: Action) -> None:
         if isinstance(action, (ScalarMultiplicationAction, PrecomputationAction)):
             self.inside = False
+            if isinstance(action, PrecomputationAction):
+                self.precomp = action.result
         if isinstance(action, FormulaAction) and self.inside:
             action = cast(FormulaAction, action)
             if isinstance(action.formula, DoublingFormula):
@@ -399,7 +416,7 @@ def multiples_computed(
     mult_factory: Callable,
     use_init: bool = False,
     use_multiply: bool = True,
-    kind: Union[Literal["all"], Literal["input"], Literal["necessary"]] = "all",
+    kind: Union[Literal["all"], Literal["input"], Literal["necessary"], Literal["precomp+necessary"]] = "all",
 ) -> set[int]:
     """
     Compute the multiples computed for a given scalar and multiplier (quickly).
@@ -410,12 +427,12 @@ def multiples_computed(
     :param mult_factory: A callable that takes the formulas and instantiates the multiplier.
     :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).
-    :param kind: The kind of multiples to return. Can be one of "all", "input", "necessary".
+    :param kind: The kind of multiples to return. Can be one of "all", "input", "necessary", or "precomp+necessary".
     :return: A list of tuples, where the first element is the formula shortname (e.g. "add") and the second is a tuple of the dlog
     relationships to the input of the input points to the formula.
     """
     mult = _cached_fake_mult(mult_class, mult_factory, params)
-    ctx = MultipleContext()
+    ctx = MultipleContext(keep_base=True)
     if use_init:
         with local(ctx, copy=False):
             mult.init(params, FakePoint(params.curve.coordinate_model))
@@ -444,6 +461,14 @@ def multiples_computed(
             for parent in ctx.parents[point]:
                 res.add(ctx.points[parent])
                 queue.add(parent)
+    elif kind == "precomp+necessary":
+        res = {ctx.points[out]}
+        queue = {out, *ctx.precomp.values()}
+        while queue:
+            point = queue.pop()
+            for parent in ctx.parents[point]:
+                res.add(ctx.points[parent])
+                queue.add(parent)
     else:
         raise ValueError(f"Invalid kind {kind}")
     return res - {0}
diff --git a/test/sca/test_rpa.py b/test/sca/test_rpa.py
index 6d58597..8d843a4 100644
--- a/test/sca/test_rpa.py
+++ b/test/sca/test_rpa.py
@@ -1,3 +1,5 @@
+from functools import partial
+
 import pytest
 from math import isqrt
 
@@ -94,9 +96,36 @@ def test_multiples_kind(rpa_params):
         17, rpa_params, RTLMultiplier, RTLMultiplier, True, True,
         kind="necessary"
     )
+    multiples_precomp = multiples_computed(
+        17, rpa_params, RTLMultiplier, RTLMultiplier, True, True,
+        kind="precomp+necessary"
+    )
+    assert multiples_all != multiples_input
+    assert multiples_all != multiples_necessary
+    assert multiples_input != multiples_necessary
+    assert multiples_precomp == multiples_necessary
+
+    wnaf = partial(WindowNAFMultiplier, width=4)
+    multiples_all = multiples_computed(
+        0xff, rpa_params, WindowNAFMultiplier, wnaf, True, True,
+        kind="all"
+    )
+    multiples_input = multiples_computed(
+        0xff, rpa_params, WindowNAFMultiplier, wnaf, True, True,
+        kind="input"
+    )
+    multiples_necessary = multiples_computed(
+        0xff, rpa_params, WindowNAFMultiplier, wnaf, True, True,
+        kind="necessary"
+    )
+    multiples_precomp = multiples_computed(
+        0xff, rpa_params, WindowNAFMultiplier, wnaf, True, True,
+        kind="precomp+necessary"
+    )
     assert multiples_all != multiples_input
     assert multiples_all != multiples_necessary
     assert multiples_input != multiples_necessary
+    assert multiples_precomp != multiples_necessary
 
 
 def test_x0_point(rpa_params):
diff --git a/test/sca/test_rpa_context.py b/test/sca/test_rpa_context.py
index 7520516..44eea5b 100644
--- a/test/sca/test_rpa_context.py
+++ b/test/sca/test_rpa_context.py
@@ -9,7 +9,6 @@ from pyecsca.ec.formula import (
     DoublingFormula,
     ScalingFormula,
 )
-from pyecsca.ec.mod import Mod
 from pyecsca.ec.mult import (
     LTRMultiplier,
     BinaryNAFMultiplier,
@@ -85,9 +84,10 @@ def test_precomp(secp128r1, add, dbl, neg, scale):
 
 def test_window(secp128r1, add, dbl, neg):
     mult = WindowNAFMultiplier(add, dbl, neg, 3, precompute_negation=True)
-    with local(MultipleContext()):
+    with local(MultipleContext()) as ctx:
         mult.init(secp128r1, secp128r1.generator)
         mult.multiply(5)
+    assert ctx.precomp
 
 
 def test_ladder(curve25519):
@@ -111,3 +111,20 @@ def test_ladder(curve25519):
         dadd_mult.multiply(1339278426732672313)
     muls = list(ctx.points.values())
     assert muls[-2] == 1339278426732672313
+
+
+def test_keep_base(secp128r1, add, dbl):
+    mult = LTRMultiplier(
+        add,
+        dbl,
+        always=False,
+        complete=False,
+        short_circuit=True,
+    )
+
+    with local(MultipleContext(keep_base=True)) as ctx:
+        mult.init(secp128r1, secp128r1.generator)
+        r = mult.multiply(5)
+        mult.init(secp128r1, r)
+        mult.multiply(10)
+    assert 50 in ctx.points.values()