Add proper EPA.

2 files changed, 244 insertions, 139 deletions

diff --git a/re/epa.ipynb b/re/epa.ipynb
index f36d829..d27a2eb 100644
--- a/re/epa.ipynb
+++ b/re/epa.ipynb
@@ -19,10 +19,12 @@
     "import tabulate\n",
     "import secrets\n",
     "from tqdm.notebook import tqdm, trange\n",
+    "from functools import partial\n",
     "from itertools import product\n",
     "from IPython.display import HTML, display\n",
     "from sympy.ntheory import factorint\n",
     "from sympy.ntheory.modular import crt\n",
+    "from anytree import Node\n",
     "\n",
     "from pyecsca.ec.model import ShortWeierstrassModel\n",
     "from pyecsca.ec.coordinates import AffineCoordinateModel\n",
@@ -32,11 +34,23 @@
     "from pyecsca.ec.point import Point, InfinityPoint\n",
     "from pyecsca.ec.error import NonInvertibleError\n",
     "from pyecsca.ec.mult import LTRMultiplier, AccumulationOrder\n",
+    "from pyecsca.ec.formula.base import *\n",
     "from pyecsca.ec.formula.fake import FakeAdditionFormula, FakeDoublingFormula, FakePoint\n",
+    "from pyecsca.sca.re.tree import build_distinguishing_tree\n",
     "from pyecsca.sca.re.rpa import MultipleContext\n",
     "from pyecsca.sca.re.zvp import unroll_formula_expr\n",
     "from pyecsca.ec.context import local\n",
-    "from pyecsca.ec.error import UnsatisfiedAssumptionError"
+    "from pyecsca.ec.error import UnsatisfiedAssumptionError\n",
+    "from pyecsca.misc.utils import log, warn\n",
+    "from pyecsca.misc.cfg import TemporaryConfig"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "0f7d1396-3b56-4f9a-bd21-33e8906ab600",
+   "metadata": {},
+   "source": [
+    "A few curves with composite \"p\"."
    ]
   },
   {
@@ -77,32 +91,25 @@
    "outputs": [],
    "source": [
     "model = ShortWeierstrassModel()\n",
-    "affine = AffineCoordinateModel(model)\n",
-    "which = \"projective\"\n",
-    "coords = model.coordinates[which]\n",
-    "\n",
-    "params = load_params_ectester(io.BytesIO(curves[4].encode()), which)\n",
-    "curve = params.curve\n",
-    "p = params.curve.prime\n",
-    "g = params.generator\n",
-    "n = params.order"
+    "affine = AffineCoordinateModel(model)"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "703da72a-1c41-41a9-886d-453da1160932",
+   "cell_type": "markdown",
+   "id": "8e0ffde3-27e1-4d28-92af-69686e42f3e6",
    "metadata": {},
-   "outputs": [],
    "source": [
-    "adds = list(filter(lambda formula: formula.name.startswith(\"add\"), coords.formulas.values()))\n",
-    "dbls = list(filter(lambda formula: formula.name.startswith(\"dbl\"), coords.formulas.values()))\n",
-    "formula_pairs = list(product(adds, dbls))\n",
-    "\n",
-    "fake_add = FakeAdditionFormula(params.curve.coordinate_model)\n",
-    "fake_dbl = FakeDoublingFormula(params.curve.coordinate_model)\n",
-    "fake_mult = LTRMultiplier(fadd, fdbl, None, False, AccumulationOrder.PeqPR, True, True)\n",
-    "fake_mult.init(params, FakePoint(params.curve.coordinate_model))"
+    "## Exploration\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c4274c9c-5efa-4c4a-8b70-0c7ca6eb89de",
+   "metadata": {},
+   "source": [
+    "Now, let's define some functions for picking random scalars mod $n$ and random points on the curve.\n",
+    "There are several ways to do so, some guarantee that the scalars will be \"trivial\" w.r.t. the curve order $n$\n",
+    "or even that all subscalars for a given scalarmult algo will be trivial w.r.t. the curve order."
    ]
   },
   {
@@ -134,9 +141,11 @@
     "        scalar = random_scalar_trivial(n)\n",
     "\n",
     "def fixed_point(params):\n",
+    "    \"\"\"Generate a fixed point on the params.\"\"\"\n",
     "    return params.generator\n",
     "\n",
     "def random_point(splitted, top, randomized=False):\n",
+    "    \"\"\"Generate a random point on the splitted params.\"\"\"\n",
     "    results = {}\n",
     "    for factor, params in splitted.items():\n",
     "        results[factor] = params.curve.affine_random()\n",
@@ -150,6 +159,14 @@
    ]
   },
   {
+   "cell_type": "markdown",
+   "id": "48727195-4c38-4b7f-8953-7108e46db98c",
+   "metadata": {},
+   "source": [
+    "Let's also define a way to project the points down to a subcurve, a way to split the curve to subcurves and a scalarmult algo that correctly computes on the top curve by splitting over the subcurves."
+   ]
+  },
+  {
    "cell_type": "code",
    "execution_count": null,
    "id": "b6e888d3-38a0-4028-add5-a7b428b8b6cd",
@@ -157,12 +174,11 @@
    "outputs": [],
    "source": [
     "def project_down(point, subcurve):\n",
+    "    \"\"\"Project a point down onto a subcurve.\"\"\"\n",
     "    return Point(subcurve.coordinate_model, **{name: Mod(int(value), subcurve.prime) for name, value in point.coords.items()})\n",
     "\n",
-    "def lift_up(point, topcurve):\n",
-    "    return Point(topcurve.coordinate_model, **{name: Mod(int(value), topcurve.prime) for name, value in point.coords.items()})\n",
-    "\n",
     "def split_params(params):\n",
+    "    \"\"\"Split composite \"p\" params into subcurves.\"\"\"\n",
     "    factors = factorint(params.curve.prime)\n",
     "    if set(factors.values()) != {1}:\n",
     "        raise ValueError(\"Not squarefree\")\n",
@@ -189,6 +205,7 @@
     "    return results\n",
     "\n",
     "def split_scalarmult(splitted, top, point, scalar):\n",
+    "    \"\"\"Perform affine scalarmult of \"point\" by \"scalar\" on the splitted params.\"\"\"\n",
     "    results = {}\n",
     "    for factor, params in splitted.items():\n",
     "        order = params.order\n",
@@ -200,6 +217,7 @@
     "            result = params.curve.affine_multiply(projected.to_affine(), partial_scalar)\n",
     "        results[factor] = result\n",
     "    if any(map(lambda point: isinstance(point, InfinityPoint), results.values())):\n",
+    "        # This is actually undefined if only one point is the infinity point.\n",
     "        return InfinityPoint(top.curve.coordinate_model)\n",
     "    factors = list(results.keys())\n",
     "    xs = list(map(lambda factor: int(results[factor].x), factors))\n",
@@ -210,6 +228,102 @@
    ]
   },
   {
+   "cell_type": "markdown",
+   "id": "ba3d3c51-0c84-4374-a85c-019462f7a55c",
+   "metadata": {},
+   "source": [
+    "With all of that we can now explore the behavior of the formulas, focusing on projective coordinates for now."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "703da72a-1c41-41a9-886d-453da1160932",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "which = \"projective\"\n",
+    "coords = model.coordinates[which]\n",
+    "\n",
+    "params = load_params_ectester(io.BytesIO(curves[4].encode()), which)\n",
+    "curve = params.curve\n",
+    "p = params.curve.prime\n",
+    "g = params.generator\n",
+    "n = params.order\n",
+    "\n",
+    "adds = list(filter(lambda formula: formula.name.startswith(\"add\"), coords.formulas.values()))\n",
+    "dbls = list(filter(lambda formula: formula.name.startswith(\"dbl\"), coords.formulas.values()))\n",
+    "formula_pairs = list(product(adds, dbls))\n",
+    "\n",
+    "#fake_add = FakeAdditionFormula(params.curve.coordinate_model)\n",
+    "#fake_dbl = FakeDoublingFormula(params.curve.coordinate_model)\n",
+    "#fake_mult = LTRMultiplier(fake_add, fake_dbl, None, False, AccumulationOrder.PeqPR, True, True)\n",
+    "#fake_mult.init(params, FakePoint(params.curve.coordinate_model))"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "c72b0fd3-21d1-4333-9a6b-f8c0d020c87d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def simulate_table(scalars, points, split, params, formula_pairs, adds, dbls):\n",
+    "    results = []\n",
+    "    chains = []\n",
+    "    gcds = []\n",
+    "    fgcds = []\n",
+    "    for scalar, point in tqdm(zip(scalars, points), desc=\"Precomp\", total=len(scalars)):\n",
+    "        try:\n",
+    "            result = split_scalarmult(split, params, point, scalar)\n",
+    "        except NonInvertibleError:\n",
+    "            result = None\n",
+    "        results.append(result)\n",
+    "        with local(MultipleContext()) as ctx:\n",
+    "            fake_mult.multiply(scalar)\n",
+    "        chains.append(list(ctx.points.values()))\n",
+    "        scalar_trivial_gcd = gcd(scalar, n) == 1\n",
+    "        all_subscalars_trivial_gcd = all(map(lambda x: gcd(x, n) == 1, ctx.points.values()))\n",
+    "        gcds.append(scalar_trivial_gcd)\n",
+    "        fgcds.append(all_subscalars_trivial_gcd)\n",
+    "    \n",
+    "    table = [[\"Pair\", \"scalars with trivial gcd\", \"scalars with all multiples with trivial gcds\", \"scalars with invertible final zs\", \"scalars with all multiples's zs invertible\", \"scalars with correct result\"]]\n",
+    "    pair_table = [[None for _ in dbls] for _ in adds]\n",
+    "    for pair in tqdm(formula_pairs):\n",
+    "        mult = LTRMultiplier(*pair, None, False, AccumulationOrder.PeqPR, True, True)\n",
+    "        inv = []\n",
+    "        correct = []\n",
+    "        zs = []\n",
+    "        for scalar, point, result in tqdm(zip(scalars, points, results), leave=None, total=len(scalars)):\n",
+    "            mult.init(params, point)\n",
+    "            with local(MultipleContext()) as ctx:\n",
+    "                res = mult.multiply(scalar)\n",
+    "            \n",
+    "            all_submultiples_invertible_z = all(map(lambda x: gcd(int(x.Z), p) == 1, ctx.points.keys()))\n",
+    "            result_invertible_z = False\n",
+    "            result_correct = False\n",
+    "            try:\n",
+    "                res_aff = res.to_affine()\n",
+    "                result_invertible_z = True\n",
+    "                if res_aff == result:\n",
+    "                    result_correct = True\n",
+    "            except NonInvertibleError as e:\n",
+    "                pass\n",
+    "            zs.append(all_submultiples_invertible_z)\n",
+    "            inv.append(result_invertible_z)\n",
+    "            correct.append(result_correct)\n",
+    "        pair_table[adds.index(pair[0])][dbls.index(pair[1])] = sum(inv)\n",
+    "        for i in inv:\n",
+    "            print(\"x\" if i else \".\", end=\"\")\n",
+    "        print()\n",
+    "        table.append([f\"{pair[0].name}, {pair[1].name}\", sum(gcds), sum(fgcds), sum(inv), sum(zs), sum(correct)])\n",
+    "    for pl, add in zip(pair_table, adds):\n",
+    "        pl.insert(0, add.name)\n",
+    "    pair_table.insert(0, [None] + [dbl.name for dbl in dbls])\n",
+    "    return table, pair_table"
+   ]
+  },
+  {
    "cell_type": "code",
    "execution_count": null,
    "id": "ffce5b71-3029-4219-a45a-1c8d78748fee",
@@ -218,76 +332,43 @@
    "source": [
     "split = split_params(params)\n",
     "scalars = [random_scalar_trivial(n) for _ in trange(50, desc=\"Generate scalars\")]\n",
-    "points = [random_point(split, params, randomized=True) for _ in trange(50, desc=\"Generate points\")]\n",
-    "results = []\n",
-    "chains = []\n",
-    "\n",
-    "\n",
-    "gcds = []\n",
-    "fgcds = []\n",
-    "for scalar, point in tqdm(zip(scalars, points), desc=\"Precomp\", total=len(scalars)):\n",
-    "    try:\n",
-    "        result = split_scalarmult(split, params, point, scalar)\n",
-    "    except NonInvertibleError:\n",
-    "        result = None\n",
-    "    results.append(result)\n",
-    "    with local(MultipleContext()) as ctx:\n",
-    "        fake_mult.multiply(scalar)\n",
-    "    chains.append(list(ctx.points.values()))\n",
-    "    scalar_trivial_gcd = gcd(scalar, n) == 1\n",
-    "    all_subscalars_trivial_gcd = all(map(lambda x: gcd(x, n) == 1, ctx.points.values()))\n",
-    "    gcds.append(scalar_trivial_gcd)\n",
-    "    fgcds.append(all_subscalars_trivial_gcd)\n",
+    "random_points = [random_point(split, params, randomized=False) for _ in trange(50, desc=\"Generate points\")]\n",
+    "fixed_points = [fixed_point(params) for _ in trange(50, desc=\"Generate points\")]\n",
     "\n",
+    "table, pair_table = simulate_table(scalars, random_points, split, params, formula_pairs, adds, dbls)\n",
+    "display(HTML(tabulate.tabulate(table, tablefmt=\"html\", headers=\"firstrow\")))\n",
+    "display(HTML(tabulate.tabulate(pair_table, tablefmt=\"html\", headers=\"firstrow\")))\n",
     "\n",
-    "table = [[\"Pair\", \"scalars with trivial gcd\", \"scalars with all multiples with trivial gcds\", \"scalars with invertible final zs\", \"scalars with all multiples's zs invertible\", \"scalars with correct result\"]]\n",
-    "pair_table = [[None for _ in dbls] for _ in adds]\n",
-    "for pair in tqdm(formula_pairs):\n",
-    "    mult = LTRMultiplier(*pair, None, False, AccumulationOrder.PeqPR, True, True)\n",
-    "    inv = []\n",
-    "    correct = []\n",
-    "    zs = []\n",
-    "    for scalar, point, result in tqdm(zip(scalars, points, results), leave=None, total=len(scalars)):\n",
-    "        mult.init(params, point)\n",
-    "        with local(MultipleContext()) as ctx:\n",
-    "            res = mult.multiply(scalar)\n",
-    "        \n",
-    "        all_submultiples_invertible_z = all(map(lambda x: gcd(int(x.Z), p) == 1, ctx.points.keys()))\n",
-    "        result_invertible_z = False\n",
-    "        result_correct = False\n",
-    "        try:\n",
-    "            res_aff = res.to_affine()\n",
-    "            result_invertible_z = True\n",
-    "            if res_aff == result:\n",
-    "                result_correct = True\n",
-    "        except NonInvertibleError as e:\n",
-    "            pass\n",
-    "        zs.append(all_submultiples_invertible_z)\n",
-    "        inv.append(result_invertible_z)\n",
-    "        correct.append(result_correct)\n",
-    "    pair_table[adds.index(pair[0])][dbls.index(pair[1])] = sum(inv)\n",
-    "    for i in inv:\n",
-    "        print(\"x\" if i else \".\", end=\"\")\n",
-    "    print()\n",
-    "    table.append([f\"{pair[0].name}, {pair[1].name}\", sum(gcds), sum(fgcds), sum(inv), sum(zs), sum(correct)])\n",
-    "for pl, add in zip(pair_table, adds):\n",
-    "    pl.insert(0, add.name)\n",
-    "pair_table.insert(0, [None] + [dbl.name for dbl in dbls])\n",
+    "table, pair_table = simulate_table(scalars, fixed_points, split, params, formula_pairs, adds, dbls)\n",
     "display(HTML(tabulate.tabulate(table, tablefmt=\"html\", headers=\"firstrow\")))\n",
     "display(HTML(tabulate.tabulate(pair_table, tablefmt=\"html\", headers=\"firstrow\")))"
    ]
   },
   {
+   "cell_type": "markdown",
+   "id": "cf709ab1-4520-41f1-a38a-19c9c9a6ff32",
+   "metadata": {},
+   "source": [
+    "## Reverse-engineering"
+   ]
+  },
+  {
    "cell_type": "code",
    "execution_count": null,
    "id": "f47c0743-a6cf-408a-8c6e-dca6732278e2",
    "metadata": {},
    "outputs": [],
    "source": [
-    "def simulate_epa_oracle(affine_params, affine_point, scalar):\n",
-    "    real_coords = model.coordinates[\"projective\"]\n",
-    "    real_add = real_coords.formulas[\"add-2007-bl\"]\n",
-    "    real_dbl = real_coords.formulas[\"dbl-2007-bl\"]\n",
+    "def simulate_epa_oracle(affine_params, affine_point, scalar, real_coord_name=\"projective\", real_add_name=\"add-2007-bl\", real_dbl_name=\"dbl-2007-bl\"):\n",
+    "    \"\"\"\n",
+    "    Simulate an EPA oracle that computes a scalar multiplication of `affine_point` by `scalar` on `affine_params`.\n",
+    "    To select the \"real\" implementation, change the `real_coord_name`, `real_add_name` and `real_dbl_name` parameters.\n",
+    "\n",
+    "    This simulates an LTR multiplier, we assume we already know the multiplier at this point.\n",
+    "    \"\"\"\n",
+    "    real_coords = model.coordinates[real_coord_name]\n",
+    "    real_add = real_coords.formulas[real_add_name]\n",
+    "    real_dbl = real_coords.formulas[real_dbl_name]\n",
     "    real_mult = LTRMultiplier(real_add, real_dbl, None, False, AccumulationOrder.PeqPR, True, True)\n",
     "    params = affine_params.to_coords(real_coords)\n",
     "    point = affine_point.to_model(real_coords, params.curve)\n",
@@ -297,97 +378,121 @@
     "        res.to_affine()\n",
     "        return True\n",
     "    except NonInvertibleError as e:\n",
-    "        return False\n",
+    "        return False"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1a4457b4-78e8-4b82-af3d-e3563125d6d7",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def epa_precomp(affine_params, mult_factory, mult_class, model, queries=30):\n",
+    "    \"\"\"\n",
+    "    Precompute a map of (cfg) -> set of indices into inputs for which the given cfg oracle will answer True,\n",
+    "    where inputs is a list of (scalar, point) pairs.\n",
     "\n",
-    "def epa_distinguish(oracle, mult_factory):\n",
-    "    affine_params = load_params_ectester(io.BytesIO(curves[3].encode()), \"affine\")\n",
-    "    model = affine_params.curve.model\n",
-    "    scalars = [int(Mod.random(affine_params.order)) for _ in range(100)]\n",
-    "    responses = [oracle(affine_params,  affine_params.generator, scalar) for scalar in scalars]\n",
-    "    candidates = set()\n",
-    "    print(\"Got responses\")\n",
-    "    total = 0\n",
-    "    for coords in model.coordinates.values():\n",
-    "        adds = list(filter(lambda formula: formula.name.startswith(\"add\"), coords.formulas.values()))\n",
-    "        dbls = list(filter(lambda formula: formula.name.startswith(\"dbl\"), coords.formulas.values()))\n",
-    "        formula_pairs = list(product(adds, dbls))\n",
-    "        total += len(formula_pairs)\n",
+    "    Returns the list of inputs, the mapping and all of the considered cfgs.\n",
+    "    Note that the mapping might be restricted over a subset of the cfgs.\n",
+    "    \"\"\"\n",
+    "    split = split_params(affine_params)\n",
+    "    scalars = [random_scalar_trivial(n) for _ in trange(queries, desc=\"Generate scalars\")]\n",
+    "    random_points = [random_point(split, affine_params, randomized=False) for _ in trange(queries, desc=\"Generate points\")]\n",
+    "    formula_classes = list(filter(lambda klass: klass in mult_class.requires, [AdditionFormula, DifferentialAdditionFormula, DoublingFormula, LadderFormula, NegationFormula]))\n",
+    "    results = {}\n",
+    "    inputs = list(zip(scalars, random_points))\n",
+    "    configs = set()\n",
+    "    for coord_name, coords in tqdm(model.coordinates.items(), desc=\"Precompute for coord systems\"):\n",
     "        try:\n",
     "            params = affine_params.to_coords(coords)\n",
     "        except UnsatisfiedAssumptionError:\n",
-    "            print(f\"Skipping {coords.name}, does not fit\")\n",
+    "            log(f\"Skipping {coords.name}, does not fit.\")\n",
     "            continue\n",
-    "\n",
-    "        for pair in formula_pairs:\n",
-    "            mult = mult_factory(*pair)\n",
-    "            mult.init(params, params.generator)\n",
-    "            abort = False\n",
-    "            print(f\"Trying {coords.name} {pair[0].name} {pair[1].name}\", end=\"\") \n",
-    "            for scalar, target in zip(scalars, responses):\n",
-    "                res = mult.multiply(scalar)\n",
+    "        log(f\"Precomputing {coords.name}.\")\n",
+    "        mapped_inputs = [(scalar, point.to_model(coords, params.curve)) for scalar, point in inputs]\n",
+    "        \n",
+    "        formula_groups = [list(filter(lambda formula: isinstance(formula, formula_class), coords.formulas.values())) for formula_class in formula_classes]\n",
+    "        formula_combinations = list(product(*formula_groups))\n",
+    "        \n",
+    "        for formulas in tqdm(formula_combinations, desc=coord_name, leave=False):\n",
+    "            cfg = (coord_name, *[formula.name for formula in formulas])\n",
+    "            configs.add(cfg)\n",
+    "            mult = mult_factory(*formulas)\n",
+    "            result = set()\n",
+    "            for i, pair in enumerate(mapped_inputs):\n",
+    "                scalar, point = pair\n",
+    "                mult.init(params, point)\n",
+    "                try:\n",
+    "                    res = mult.multiply(scalar)\n",
+    "                except UnsatisfiedAssumptionError as e:\n",
+    "                    break\n",
     "                try:\n",
     "                    res.to_affine()\n",
-    "                    if not target:\n",
-    "                        # not this one\n",
-    "                        abort = True\n",
-    "                        break\n",
+    "                    result.add(i)\n",
     "                except NonInvertibleError as e:\n",
-    "                    if target:\n",
-    "                        # not this one\n",
-    "                        abort = True\n",
-    "                        break\n",
-    "            if abort:\n",
-    "                print(\" not\")\n",
-    "                continue\n",
+    "                    pass\n",
     "            else:\n",
-    "                print(\" candidate\")\n",
-    "                candidates.add((coords.name, pair[0].name, pair[1].name))\n",
-    "    print(f\"Got {len(candidates)} out of {total} total\")\n",
-    "    return candidates\n",
-    "                "
+    "                results[cfg] = result\n",
+    "    return inputs, results, configs"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d2c26c45-a66a-4472-830e-b22010a967e5",
+   "id": "5defdea2-0633-40fb-b53c-afb8f8eda035",
    "metadata": {},
    "outputs": [],
    "source": [
-    "c = epa_distinguish(simulate_epa_oracle, lambda add,dbl:LTRMultiplier(add, dbl, None, False, AccumulationOrder.PeqPR, True, True))"
+    "def epa_distinguish_precomp(inputs, precomp, configs, affine_params, oracle):\n",
+    "    \"\"\"\n",
+    "    Distinguish the coordinate system and formulas using EPA given the precomputation.\n",
+    "    \"\"\"\n",
+    "    tree = build_distinguishing_tree(precomp)\n",
+    "    log(\"Built distinguishing tree.\")\n",
+    "    log(RenderTree(tree).by_attr(lambda n: n.name if n.name is not None else n.cfgs))\n",
+    "\n",
+    "    current_node = tree\n",
+    "    cfgs = list(precomp.keys())\n",
+    "    while current_node.children:\n",
+    "        best_distinguishing_index = current_node.name\n",
+    "        scalar, point = inputs[best_distinguishing_index]\n",
+    "        response = oracle(affine_params, point, scalar)\n",
+    "        log(f\"Oracle response -> {response}\")\n",
+    "        for cfg in cfgs:\n",
+    "            log(cfg, best_distinguishing_index in precomp[cfg])\n",
+    "        response_map = {child.oracle_response: child for child in current_node.children}\n",
+    "        current_node = response_map[response]\n",
+    "        cfgs = current_node.cfgs\n",
+    "        log(cfgs)\n",
+    "        log()\n",
+    "    return cfgs"
    ]
   },
   {
-   "cell_type": "code",
-   "execution_count": null,
-   "id": "0c8a43dc-1a2f-4135-b83c-6ac84929789f",
+   "cell_type": "markdown",
+   "id": "b966e32e-2ed4-4437-bc68-17494a56d559",
    "metadata": {},
-   "outputs": [],
    "source": [
-    "split = split_params(params)\n",
-    "split"
+    "Now we can run the precomp and the EPA reverse-engineering."
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "d6d8bb2a-2e5d-475e-b775-5aff54eeaf47",
+   "id": "413c49eb-c096-485c-8471-18ac9c1b9523",
    "metadata": {},
    "outputs": [],
    "source": [
-    "for add in adds:\n",
-    "    r = None\n",
-    "    for iv in unroll_formula_expr(add):\n",
-    "        if iv[0] == \"Z3\":\n",
-    "            r = iv[1]\n",
-    "    print(add, r)\n",
-    "    print(\"---\")"
+    "affine_params = load_params_ectester(io.BytesIO(curves[3].encode()), \"affine\")\n",
+    "inputs, precomp, configs = epa_precomp(affine_params, lambda add,dbl:LTRMultiplier(add, dbl, None, False, AccumulationOrder.PeqPR, True, True), LTRMultiplier, model)\n",
+    "epa_distinguish_precomp(inputs, precomp, configs, affine_params, simulate_epa_oracle)"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "64202cc7-a8f8-4d88-a0b5-5cf4d483a547",
+   "id": "467d3ed8-8eb9-4075-8621-d3653d16deeb",
    "metadata": {},
    "outputs": [],
    "source": []
diff --git a/re/rpa.ipynb b/re/rpa.ipynb
index dc04a58..f328b2d 100644
--- a/re/rpa.ipynb
+++ b/re/rpa.ipynb
@@ -620,7 +620,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.4"
+   "version": "3.11.5"
   }
  },
  "nbformat": 4,