Cleanup ZVP notebook.

1 files changed, 487 insertions, 111 deletions

diff --git a/re/zvp.ipynb b/re/zvp.ipynb
index 072b0aa..4beff1c 100644
--- a/re/zvp.ipynb
+++ b/re/zvp.ipynb
@@ -17,16 +17,25 @@
    "source": [
     "import io\n",
     "import numpy as np\n",
-    "from sympy import FF, sympify, symbols, Poly\n",
+    "import pandas as pd\n",
     "import random\n",
     "import tabulate\n",
     "import pickle\n",
+    "import multiprocessing\n",
+    "import inspect\n",
+    "import tempfile\n",
+    "import sys\n",
+    "import re\n",
+    "from sympy import FF, ZZ, sympify, symbols, Poly\n",
+    "from contextlib import contextmanager\n",
+    "from importlib import import_module, invalidate_caches\n",
+    "from pathlib import Path\n",
     "from functools import partial\n",
     "from itertools import product\n",
     "from IPython.display import HTML, display\n",
-    "from tqdm.notebook import tqdm\n",
+    "from tqdm.notebook import tqdm, trange\n",
     "from anytree import RenderTree, PreOrderIter\n",
-    "from concurrent.futures import ProcessPoolExecutor, as_completed\n",
+    "\n",
     "\n",
     "from pyecsca.ec.model import ShortWeierstrassModel\n",
     "from pyecsca.ec.coordinates import AffineCoordinateModel\n",
@@ -35,14 +44,33 @@
     "from pyecsca.ec.formula import FormulaAction, AdditionFormula, DoublingFormula\n",
     "from pyecsca.ec.point import Point\n",
     "from pyecsca.ec.mod import Mod, gcd, SymbolicMod\n",
-    "from pyecsca.sca.re.tree import Map, Tree\n",
-    "from pyecsca.sca.re.rpa import MultipleContext\n",
-    "from pyecsca.sca.re.zvp import zvp_points, compute_factor_set\n",
     "from pyecsca.ec.context import DefaultContext, local\n",
     "from pyecsca.ec.mult import LTRMultiplier, AccumulationOrder\n",
-    "from pyecsca.misc.cfg import getconfig\n",
     "from pyecsca.ec.error import NonInvertibleError, UnsatisfiedAssumptionError\n",
-    "from pyecsca.sca.re.zvp import unroll_formula, compute_factor_set, zvp_points, addition_chain"
+    "from pyecsca.sca.re.tree import Map, Tree\n",
+    "from pyecsca.sca.re.rpa import MultipleContext\n",
+    "from pyecsca.sca.re.zvp import zvp_points, compute_factor_set, unroll_formula,  addition_chain, eliminate_y\n",
+    "from pyecsca.misc.cfg import getconfig\n",
+    "from pyecsca.misc.utils import TaskExecutor\n",
+    "\n",
+    "\n",
+    "# Allow to use \"spawn\" multiprocessing method for function defined in a Jupyter notebook.\n",
+    "# https://neuromancer.sk/article/35\n",
+    "@contextmanager\n",
+    "def enable_spawn(func):\n",
+    "    invalidate_caches()\n",
+    "    source = inspect.getsource(func)\n",
+    "    with tempfile.NamedTemporaryFile(suffix=\".py\", mode=\"w\") as f:\n",
+    "        f.write(source)\n",
+    "        f.flush()\n",
+    "        path = Path(f.name)\n",
+    "        directory = str(path.parent)\n",
+    "        sys.path.append(directory)\n",
+    "        module = import_module(str(path.stem))\n",
+    "        yield getattr(module, func.__name__)\n",
+    "        sys.path.remove(directory)\n",
+    "\n",
+    "spawn_context = multiprocessing.get_context(\"spawn\")"
    ]
   },
   {
@@ -62,7 +90,7 @@
    "metadata": {},
    "source": [
     "## Exploration\n",
-    "First lets explore the behavior of addition formulas. The following two cells pick a coordinate model along with some formulas and symbolically unroll a scalar multiplication."
+    "First lets explore the behavior of addition formulas. The following two cells pick a coordinate model along with some formulas and symbolically unroll a scalar multiplication (assuming a simple LTR multiplier)."
    ]
   },
   {
@@ -113,7 +141,12 @@
     "              Y=SymbolicMod(Poly(y, x, y, z, domain=field), params.curve.prime),\n",
     "              Z=SymbolicMod(Poly(z, x, y, z, domain=field), params.curve.prime))\n",
     "mult.init(params, point)\n",
-    "res = mult.multiply(5)"
+    "res = mult.multiply(5)\n",
+    "\n",
+    "x_poly = Poly(res.X.x, domain=field)\n",
+    "y_poly = Poly(res.Y.x, domain=field)\n",
+    "z_poly = Poly(res.Z.x, domain=field)\n",
+    "display(x_poly, y_poly, z_poly)"
    ]
   },
   {
@@ -121,7 +154,7 @@
    "id": "0bcc8b9e-39ad-4b53-8f45-a7552b05baa2",
    "metadata": {},
    "source": [
-    "The result is a Point with coordinates that are polynomials in the input coordinates and curve parameters."
+    "The result is a Point with coordinates that are polynomials in the input coordinates and curve parameters. We now switch back to concrete representation."
    ]
   },
   {
@@ -141,7 +174,7 @@
    "metadata": {},
    "source": [
     "## Reverse-engineering\n",
-    "Now, lets look at using the ZVP attack for reverse-engineering. First pick 10 random curves, some with $a \\in \\{0, -1, -3 \\}$. The randomcurves are not special in any way and just serve to randomize the process, as the existence of ZVP points for a given intermediate value polynomial depends on the curve."
+    "Now, lets look at using the ZVP attack for reverse-engineering. First pick 10 curves per group, some random some with $a \\in \\{0, -1, -3 \\}$. The curves are otherwise not special in any way and just serve to randomize the process, as the existence of ZVP points for a given intermediate value polynomial depends on the curve."
    ]
   },
   {
@@ -218,7 +251,9 @@
    "id": "4276de4c-78f4-4cdb-b60e-8c24eabfa00d",
    "metadata": {},
    "source": [
-    "First lets fix some scalars, go over the curves and compute the addition chain to obtain information about which multiples of the input point will go into the formulas."
+    "### Computing addition chains\n",
+    "\n",
+    "First lets fix some scalars, go over the curves and compute the addition chain to obtain information about which multiples of the input point will go into the formulas. The i-th scalar will be used with the i-th curve as defined above. There are 10 unique scalars, so each curve group will share those."
    ]
   },
   {
@@ -232,21 +267,64 @@
     "\n",
     "chains = []\n",
     "scalar_map = {}\n",
+    "chain_map = {}\n",
     "ops = set()\n",
     "for scalar, params in zip(scalars, curves):\n",
     "    chain = addition_chain(scalar, params, LTRMultiplier, lambda add,dbl: LTRMultiplier(add, dbl, None, False, AccumulationOrder.PeqRP, True, True))\n",
     "    chains.append(chain)\n",
     "    scalar_map[params] = scalar\n",
+    "    chain_map[params] = chain\n",
     "    ops.update(chain)\n",
     "print(sorted(list(ops)))"
    ]
   },
   {
    "cell_type": "markdown",
+   "id": "8d1df0a5-ac90-45ee-b54a-1fa25dd86c74",
+   "metadata": {},
+   "source": [
+    "### Loading the formulas\n",
+    "\n",
+    "Now lets load the formulas, either just those from the EFD or also load the expanded library formulas if they are available. See the [formulas](formulas.ipynb) notebook."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "da4e2981-cf3d-40ea-9d1e-9b6c57a6e01d",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "load_expanded = False\n",
+    "\n",
+    "formula_classes = [AdditionFormula, DoublingFormula]\n",
+    "formula_groups = {}\n",
+    "for coord_name, coords in tqdm(model.coordinates.items(), desc=f\"Loading {'expanded' if load_expanded else 'EFD'} formulas\"):\n",
+    "    groups = []\n",
+    "    for formula_class in formula_classes:\n",
+    "        expanded_path = Path(f\"sw_{coord_name}_{formula_class.shortname}s.pickle\")\n",
+    "        if load_expanded:\n",
+    "            if not expanded_path.exists():\n",
+    "                raise ValueError(f\"Expanded formulas do not exist {expanded_path}.\")\n",
+    "            with expanded_path.open(\"rb\") as f:\n",
+    "                expanded = pickle.load(f)\n",
+    "            formula_group = list(expanded)\n",
+    "        else:\n",
+    "            formula_group = list(filter(lambda formula: isinstance(formula, formula_class) and (formula.name.startswith(\"add\") or formula.name.startswith(\"dbl\")), coords.formulas.values()))\n",
+    "        groups.append(formula_group)\n",
+    "    formula_groups[coords] = groups\n",
+    "\n",
+    "print(f\"Loaded {sum(sum(len(group) for group in pair) for pair in formula_groups.values())} formulas.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
    "id": "12646534-8ca5-48c1-a4ae-1ad62575821f",
    "metadata": {},
    "source": [
-    "Now, lets compute the sets of ZVP points, going over all coordinate systems and all of their formulas (that fit the scalar multiplier) and store them into the `point_chains`. These items form individual distinct entries of the distinguishing table."
+    "### Computing the factor sets\n",
+    "\n",
+    "Now, compute the factor sets of the formulas **in parallel**. There are two options available. `xonly_fsets` builds the factor sets \"x\"-only by eliminating y-coords using the curve equation. `filter_nonhomo` specifies whether to filter out non-homogenous polynomials."
    ]
   },
   {
@@ -256,30 +334,101 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "formula_classes = [AdditionFormula, DoublingFormula]\n",
+    "xonly_fsets = False\n",
+    "filter_nonhomo = False\n",
+    "\n",
+    "def compute_fsets(formula_group, formula_class, fh, xo):\n",
+    "    from pyecsca.sca.re.zvp import compute_factor_set\n",
+    "    from pyecsca.ec.formula import DoublingFormula\n",
+    "    fsets = []\n",
+    "    for formula in formula_group:\n",
+    "        fset = compute_factor_set(formula, filter_nonhomo=fh, xonly=xo)\n",
+    "\n",
+    "        # Fix the factor set polynomials for the case of doubling.\n",
+    "        # TODO: Investigate how this plays with xonly an filter_nonhomo arguments.\n",
+    "        if formula_class == DoublingFormula:\n",
+    "            new_fset = set()\n",
+    "            for poly in fset:\n",
+    "                pl = poly.copy()\n",
+    "                for symbol in poly.free_symbols:\n",
+    "                    original = str(symbol)\n",
+    "                    if original.endswith(\"1\"):\n",
+    "                        new = original.replace(\"1\", \"2\")\n",
+    "                        pl = pl.subs(original, new)\n",
+    "                new_fset.add(pl)\n",
+    "            fset = new_fset\n",
+    "        fsets.append(fset)\n",
+    "    return fsets\n",
+    "\n",
     "factor_sets = {}\n",
-    "for coord_name, coords in model.coordinates.items():\n",
-    "    formula_groups = [list(filter(lambda formula: isinstance(formula, formula_class) and (formula.name.startswith(\"add\") or formula.name.startswith(\"dbl\")), coords.formulas.values())) for formula_class in formula_classes]\n",
+    "with TaskExecutor(max_workers=22, mp_context=spawn_context) as pool, enable_spawn(compute_fsets) as target:\n",
+    "    for coord_name, coords in model.coordinates.items():\n",
+    "        for formula_group, formula_class in zip(formula_groups[coords], formula_classes):\n",
+    "            pool.submit_task((coords, formula_group, formula_class),\n",
+    "                             target, formula_group, formula_class, filter_nonhomo, xonly_fsets)\n",
     "\n",
-    "    for formula_group, formula_class in zip(formula_groups, formula_classes):\n",
-    "        for formula in formula_group:\n",
-    "            fset = compute_factor_set(formula, filter_nonhomo=False)\n",
-    "            # Fix the factor set polynomials for the case of doubling.\n",
-    "            if formula_class == DoublingFormula:\n",
-    "                new_fset = set()\n",
-    "                for poly in fset:\n",
-    "                    pl = poly.copy()\n",
-    "                    for symbol in poly.free_symbols:\n",
-    "                        original = str(symbol)\n",
-    "                        if original.endswith(\"1\"):\n",
-    "                            new = original.replace(\"1\", \"2\")\n",
-    "                            pl = pl.subs(original, new)\n",
-    "                    new_fset.add(pl)\n",
-    "                fset = new_fset\n",
-    "            factor_sets[formula] = fset\n",
+    "    for (coords, formula_group, formula_class), future in tqdm(pool.as_completed(), desc=\"Computing factor sets\", total=len(pool.tasks)):\n",
+    "        if error := future.exception():\n",
+    "            print(coords, formula_class.shortname, error)\n",
+    "            raise error\n",
+    "        else:\n",
+    "            fsets = future.result()\n",
+    "            print(f\"Got {coords.name} {formula_class.shortname}s {len(fsets)}.\")\n",
+    "            for formula, fset in zip(formula_group, fsets):\n",
+    "                factor_sets[formula] = fset"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "c7a753be-741e-4687-b1a9-4688281d2c71",
+   "metadata": {},
+   "source": [
+    "You can now store the (or load the previously computed) factor sets."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "aef30807-e9eb-44f2-8837-f027c54eb724",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(\"factor_sets.pickle\", \"wb\") as f:\n",
+    "    pickle.dump(factor_sets, f)\n",
+    "    print(f\"Stored {len(factor_sets)} factor sets.\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "86df44e6-aa3c-48eb-8561-36b08eaa781c",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(\"factor_sets.pickle\", \"rb\") as f:\n",
+    "    factor_sets = pickle.load(f)\n",
+    "    print(f\"Loaded {len(factor_sets)} factor sets.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "5781ae2b-28ce-4926-a2ab-b3d4689df29c",
+   "metadata": {},
+   "source": [
+    "### Accumulating polynomials\n",
     "\n",
+    "We now accumulate all of the polynomials for the adds and the dbls. We do so for all the compatible curves from our curves list. We will be looking for ZVP points on all of these curves, to make sure at least some lead to solutions."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "564b2c16-6655-4ddb-8657-611d1d489e6a",
+   "metadata": {},
+   "outputs": [],
+   "source": [
     "polynomials = {}\n",
-    "for coord_name, coords in model.coordinates.items():\n",
+    "for coord_name, coords in tqdm(model.coordinates.items(), desc=\"Accumulating\"):\n",
     "    coord_adds = 0\n",
     "    coord_dbls = 0\n",
     "    for chain, affine_params in zip(chains, curves):\n",
@@ -289,20 +438,28 @@
     "            continue\n",
     "        add_polynomials = set()\n",
     "        dbl_polynomials = set()\n",
-    "        formula_groups = [list(filter(lambda formula: isinstance(formula, formula_class) and (formula.name.startswith(\"add\") or formula.name.startswith(\"dbl\")), coords.formulas.values())) for formula_class in formula_classes]\n",
-    "        for formula_group, formula_class in zip(formula_groups, formula_classes):\n",
+    "        for formula_group, formula_class in zip(formula_groups[coords], formula_classes):\n",
     "            for formula in formula_group:\n",
     "                if formula_class == AdditionFormula:\n",
-    "                    add_polynomials.update(factor_sets[formula])\n",
+    "                    add_polynomials.update(factor_sets.get(formula, []))\n",
     "                else:\n",
-    "                    dbl_polynomials.update(factor_sets[formula])\n",
+    "                    dbl_polynomials.update(factor_sets.get(formula, []))\n",
     "        polynomials[params] = {\n",
     "            \"add\": add_polynomials,\n",
     "            \"dbl\": dbl_polynomials\n",
     "        }\n",
     "        coord_adds += len(add_polynomials)\n",
     "        coord_dbls += len(dbl_polynomials)\n",
-    "    print(f\"Got {coord_adds} add polys and {coord_dbls} dbl polys for {coord_name}\")"
+    "    print(f\"Got {coord_adds} add polys and {coord_dbls} dbl polys for {coord_name}.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "04e7245b-a36a-43dc-9e73-626241112923",
+   "metadata": {},
+   "source": [
+    "### Computing ZVP points\n",
+    "Now, lets compute the sets of ZVP points, going over all of the polynomials. Also filter the points such that for each \"polynomial, curve category, k\" we have only one point, as more are unneccessary."
    ]
   },
   {
@@ -313,18 +470,21 @@
    "outputs": [],
    "source": [
     "# bound is the maximal dlog in the hard case of the DCP to be solved\n",
-    "bound = 100\n",
+    "bound = 50\n",
+    "# Note that if you do not have the \"pari\" extra dependency installed (\"cysignals\", \"cypari2\") this bound\n",
+    "# will have to be limited very low and the memory usage will be significant.\n",
     "\n",
     "all_points = set()\n",
-    "with ProcessPoolExecutor(max_workers=20) as pool:\n",
-    "    futures = []\n",
-    "    args = []\n",
-    "    for coord_name, coords in model.coordinates.items():\n",
+    "all_points_filtered = {}\n",
+    "dk = set()\n",
+    "with TaskExecutor(max_workers=20, mp_context=spawn_context) as pool:\n",
+    "    for coord_name, coords in tqdm(model.coordinates.items(), desc=\"Submitting\"):\n",
     "        for chain, affine_params in zip(chains, curves):\n",
     "            try:\n",
     "                params = affine_params.to_coords(coords)\n",
     "            except UnsatisfiedAssumptionError:\n",
     "                continue\n",
+    "            unique = set()\n",
     "            for op, ks in chain:\n",
     "                if len(ks) == 1:\n",
     "                    k = ks[0]\n",
@@ -339,15 +499,52 @@
     "                    # This is the hard case where a dlog needs to be substituted, so bound it.\n",
     "                    if not (only_1 or only_2) and k > bound:\n",
     "                        continue\n",
-    "                    futures.append(pool.submit(zvp_points, poly, params.curve, k, params.order))\n",
-    "                    args.append((poly, affine_params, k))\n",
-    "    for future in tqdm(as_completed(futures), desc=\"Computing\", total=len(futures), smoothing=0):\n",
-    "        j = futures.index(future)\n",
-    "        poly, affine_params, k = args[j]\n",
-    "        result = future.result()\n",
-    "        for point in result:\n",
-    "            all_points.add((point, affine_params))\n",
-    "print(f\"Got {len(all_points)} points\")"
+    "                    unique.add((poly, k))\n",
+    "            for poly, k in unique:\n",
+    "                pool.submit_task((poly, affine_params, k),\n",
+    "                                 zvp_points, poly, params.curve, k, params.order)\n",
+    "    for (poly, affine_params, k), future in tqdm(pool.as_completed(), desc=\"Computing\", total=len(pool.tasks), smoothing=0):\n",
+    "        params_name_match = re.match(\"(.+)\\[([0-9]+)\\]\", affine_params.name)\n",
+    "        params_category = params_name_match.group(1)\n",
+    "        if error := future.exception():\n",
+    "            print(error)\n",
+    "        elif result := future.result():\n",
+    "            for point in result:\n",
+    "                all_points.add((point, affine_params))\n",
+    "            all_points_filtered[(poly, params_category, k)] = (result, affine_params)\n",
+    "\n",
+    "print(f\"Got {len(all_points)} points.\")\n",
+    "print(f\"Got {len(all_points_filtered)} filtered points\")#, but just {len(set(all_points_filtered.values()))} unique.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "8d571e64-5261-4cda-b490-5b85dd6663d2",
+   "metadata": {},
+   "source": [
+    "You can now store the (or load the previously computed) point sets."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "82bc0305-b302-4240-a34c-81e13fe646b5",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(\"all_points.pickle\", \"wb\") as f:\n",
+    "    pickle.dump((all_points, all_points_filtered), f)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "411832e9-26a8-46a4-82e6-5f6afe879188",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "with open(\"all_points.pickle\", \"rb\") as f:\n",
+    "    all_points, all_points_filtered = pickle.load(f)"
    ]
   },
   {
@@ -355,6 +552,8 @@
    "id": "ed2d4b51-573f-4fef-b1a4-d30b3be07423",
    "metadata": {},
    "source": [
+    "### Remapping\n",
+    "\n",
     "Our ZVP points might (due to the bounds above thus the incompleteness of our analysis) lead to more zeros than we attribute to them (in more configurations), i.e. \"false negatives\". They might also be erroneous and not lead to zeros if the argument `filter_nonhomo` is False, as non-homogenous intermediate polynomials are not filtered out of the analysis. They introduce \"false positives\" but also some true positives.\n",
     "\n",
     "Thus we perform a remapping step where we execute the scalar multiplication with given points and trace whether they introduce the zeros. This gives us a new distinguishing map `remapped_hit_point_map`, now without \"false negatives\" or \"false positives\".\n",
@@ -369,62 +568,99 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def remap(coords, formulas, points, scalar_map):\n",
-    "    mult = LTRMultiplier(*formulas, None, False, AccumulationOrder.PeqRP, True, True)\n",
-    "    hit_points = set()\n",
-    "    count_points = {}\n",
-    "    position_points = {}\n",
-    "    \n",
-    "    param_map = {}\n",
-    "    for point, params in points:\n",
-    "        if params not in param_map:\n",
-    "            try:\n",
-    "                param_map[params] = params.to_coords(coords)\n",
-    "            except UnsatisfiedAssumptionError:\n",
-    "                param_map[params] = None\n",
-    "                continue\n",
-    "        elif param_map[params] is None:\n",
-    "            continue\n",
-    "        mult.init(param_map[params], point.to_model(param_map[params].curve.coordinate_model, param_map[params].curve))\n",
-    "        scalar = scalar_map[params]\n",
-    "        with local(DefaultContext()) as ctx:\n",
-    "            try:\n",
-    "                mult.multiply(scalar)\n",
-    "            except UnsatisfiedAssumptionError:\n",
-    "                continue\n",
-    "        trace = []\n",
+    "def remap(coords, chunk, points, scalar_map):\n",
+    "    import numpy as np\n",
+    "    from pyecsca.ec.mult import LTRMultiplier, AccumulationOrder\n",
+    "    from pyecsca.ec.context import local, DefaultContext\n",
+    "    from pyecsca.ec.error import UnsatisfiedAssumptionError\n",
+    "    from pyecsca.ec.formula import FormulaAction\n",
+    "    lp = len(points)\n",
+    "    lc = len(chunk)\n",
+    "    counts = np.full((lc, lp), -1, dtype=np.int16)\n",
+    "    positions = np.full((lc, lp), None, dtype=object)\n",
+    "    for i, formulas in enumerate(chunk):\n",
+    "        mult = LTRMultiplier(*formulas, None, False, AccumulationOrder.PeqRP, True, True)\n",
     "        \n",
-    "        def callback(action):\n",
-    "            if isinstance(action, FormulaAction):\n",
-    "                for intermediate in action.op_results:\n",
-    "                    trace.append(intermediate.value)\n",
-    "        ctx.actions.walk(callback)\n",
-    "        zeros = tuple(map(lambda x: int(x) == 0, trace))\n",
-    "        if any(zeros):\n",
-    "            hit_points.add((point, params))\n",
-    "        count_points[(point, params)] = sum(zeros)\n",
-    "        position_points[(point, params)] = zeros\n",
-    "    return hit_points, count_points, position_points\n",
+    "        for j, entry in enumerate(points):\n",
+    "            if entry is None:\n",
+    "                continue\n",
+    "            point, params = entry\n",
+    "            mult.init(params, point)\n",
+    "            scalar = scalar_map[params]\n",
+    "            with local(DefaultContext()) as ctx:\n",
+    "                try:\n",
+    "                    mult.multiply(scalar)\n",
+    "                except UnsatisfiedAssumptionError:\n",
+    "                    continue\n",
+    "\n",
+    "            zeros = []\n",
+    "            \n",
+    "            def callback(action):\n",
+    "                if isinstance(action, FormulaAction):\n",
+    "                    for intermediate in action.op_results:\n",
+    "                        zeros.append(int(intermediate.value) == 0)\n",
+    "            ctx.actions.walk(callback)\n",
+    "            count = sum(zeros)\n",
+    "            counts[i, j] = count\n",
+    "            positions[i, j] = tuple(zeros)        \n",
+    "    return counts, positions\n",
     "\n",
     "remapped_hit_point_map = {}\n",
     "remapped_count_point_map = {}\n",
     "remapped_position_point_map = {}\n",
-    "with ProcessPoolExecutor(max_workers=30) as pool:\n",
-    "    futures = []\n",
-    "    pairs = []\n",
-    "    for coord_name, coords in model.coordinates.items():\n",
-    "        formula_groups = [list(filter(lambda formula: isinstance(formula, formula_class) and (formula.name.startswith(\"add\") or formula.name.startswith(\"dbl\")), coords.formulas.values())) for formula_class in formula_classes]\n",
-    "        formula_combinations = list(product(*formula_groups))\n",
-    "        for formulas in formula_combinations:\n",
-    "            futures.append(pool.submit(remap, coords, formulas, all_points, scalar_map))\n",
-    "            pairs.append(formulas)\n",
-    "    for future in tqdm(as_completed(futures), total=len(futures), desc=\"Remapping\", smoothing=0):\n",
-    "        j = futures.index(future)\n",
-    "        cfg = pairs[j]\n",
-    "        h, c, p = future.result()\n",
-    "        remapped_hit_point_map[cfg] = h\n",
-    "        remapped_count_point_map[cfg] = c\n",
-    "        remapped_position_point_map[cfg] = p"
+    "all_points_list = list(all_points) #list(set(all_points_filtered.values()))\n",
+    "\n",
+    "with TaskExecutor(max_workers=30, mp_context=spawn_context) as pool, enable_spawn(remap) as remap_spawn:\n",
+    "    for coord_name, coords in tqdm(model.coordinates.items()):\n",
+    "        param_map = {}\n",
+    "        points_mapped = []\n",
+    "        scalars_mapped = {}\n",
+    "        mapped = 0\n",
+    "        for point, params in tqdm(all_points_list, desc=f\"Map points to {coord_name}\", leave=False):\n",
+    "            if params not in param_map:\n",
+    "                try:\n",
+    "                    param_map[params] = params.to_coords(coords)\n",
+    "                except UnsatisfiedAssumptionError:\n",
+    "                    param_map[params] = None\n",
+    "            if param_map[params] is None:\n",
+    "                points_mapped.append(None)\n",
+    "            else:\n",
+    "                mapped += 1\n",
+    "                points_mapped.append((point.to_model(param_map[params].curve.coordinate_model, param_map[params].curve), param_map[params]))\n",
+    "        print(f\"{coord_name}: {mapped} are compatible. Remapping...\")\n",
+    "        for params, scalar in scalar_map.items():\n",
+    "            if params not in param_map or param_map[params] is None:\n",
+    "                continue\n",
+    "            scalars_mapped[param_map[params]] = scalar\n",
+    "        \n",
+    "        pairs = list(product(*formula_groups[coords]))\n",
+    "        chunk_size = 10\n",
+    "        chunks = 0\n",
+    "        for i in trange(0, len(pairs), chunk_size, desc=f\"Chunking {coord_name}\", smoothing=0, leave=False):\n",
+    "            chunk = pairs[i:i+chunk_size]\n",
+    "            chunks += 1\n",
+    "            pool.submit_task(chunk,\n",
+    "                             remap_spawn, coords, chunk, points_mapped, scalars_mapped)\n",
+    "\n",
+    "        for chunk, future in tqdm(pool.as_completed(), total=chunks, unit_scale=chunk_size, desc=f\"Remapping {coord_name} ({len(pairs)} formula pairs)\", leave=False, smoothing=0):\n",
+    "            error = future.exception()\n",
+    "            if error:\n",
+    "                print(j, error)\n",
+    "            else:\n",
+    "                counts, positions = future.result()\n",
+    "                for cfg, counts_row, positions_row in zip(chunk, counts, positions):\n",
+    "                    hit_set = set()\n",
+    "                    hit_map = {}\n",
+    "                    count_map = {}\n",
+    "                    position_map = {}\n",
+    "                    for pp_tuple, count, position in zip(all_points_list, counts_row, positions_row):\n",
+    "                        hit_map[pp_tuple] = None if count == -1 else (count > 0)\n",
+    "                        count_map[pp_tuple] = count\n",
+    "                        position_map[pp_tuple] = position\n",
+    "                    remapped_hit_point_map[cfg] = hit_map\n",
+    "                    remapped_count_point_map[cfg] = count_map\n",
+    "                    remapped_position_point_map[cfg] = position_map\n",
+    "        print(f\"Remapped {mapped} points for {coord_name}.\")"
    ]
   },
   {
@@ -432,7 +668,8 @@
    "id": "3c0d710b-19fd-4ff7-a39f-f38b1b8856f8",
    "metadata": {},
    "source": [
-    "Finally, we can build a tree using the remapped distinguishing map."
+    "### Distinguishing map and distinguishing tree building\n",
+    "Finally, we can build a tree using the remapped distinguishing map. Let's first wrap the raw mapping with a distinguishing Map object."
    ]
   },
   {
@@ -442,7 +679,28 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dmap_remapped = Map.from_binary_sets(set(remapped_hit_point_map.keys()), remapped_hit_point_map)"
+    "dmap_remapped = Map.from_io_maps(set(remapped_hit_point_map.keys()), remapped_hit_point_map) "
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "48baa442-549f-4ffa-b296-4f964c8efc77",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from copy import deepcopy\n",
+    "dmap_copy = deepcopy(dmap_remapped)\n",
+    "dmap_copy.deduplicate()\n",
+    "print(f\"Points before: {len(dmap_remapped.mapping)} points after deduplication: {len(dmap_copy.mapping)}.\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "ce87cea8-ea88-49ee-a7dd-ead62dff5255",
+   "metadata": {},
+   "source": [
+    "Let's now watch the tree getting built."
    ]
   },
   {
@@ -454,7 +712,20 @@
    },
    "outputs": [],
    "source": [
-    "tree_remapped = Tree.build(set(remapped_hit_point_map.keys()), dmap_remapped)"
+    "tree_remapped = Tree.build(dmap_copy.cfgs, dmap_copy)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "bdb02213-4743-488c-b6ab-a934a409df3b",
+   "metadata": {},
+   "source": [
+    "The tree is built, we can examine its properties, such as:\n",
+    " - the total number of configurations (formula pairs),\n",
+    " - its depth,\n",
+    " - number of leaves (configurations that cannot be further distinguished),\n",
+    " - average leaf size\n",
+    " - mean result size (if this tree were to be used for reverse-enginnering)."
    ]
   },
   {
@@ -482,7 +753,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dmap_count = Map.from_io_map(set(remapped_count_point_map.keys()), remapped_count_point_map)"
+    "dmap_count = Map.from_io_maps(set(remapped_count_point_map.keys()), remapped_count_point_map)"
    ]
   },
   {
@@ -504,7 +775,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dmap_position = Map.from_io_map(set(remapped_position_point_map.keys()), remapped_position_point_map)"
+    "dmap_position = Map.from_io_maps(set(remapped_position_point_map.keys()), remapped_position_point_map)"
    ]
   },
   {
@@ -583,8 +854,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "dmap_fset = Map.from_binary_sets(set(fset_map.keys()), fset_map)\n",
-    "dmap_fset_nonhomo = Map.from_binary_sets(set(fset_nonhomo_map.keys()), fset_nonhomo_map)"
+    "dmap_fset = Map.from_sets(set(fset_map.keys()), fset_map)\n",
+    "dmap_fset_nonhomo = Map.from_sets(set(fset_nonhomo_map.keys()), fset_nonhomo_map)"
    ]
   },
   {
@@ -651,6 +922,14 @@
    ]
   },
   {
+   "cell_type": "markdown",
+   "id": "07151071-33ed-4f24-8b47-73fc1feda2d7",
+   "metadata": {},
+   "source": [
+    "## Miscellaneous analysis"
+   ]
+  },
+  {
    "cell_type": "code",
    "execution_count": null,
    "id": "b78f99d6-46b9-44bb-b9b4-5df7fc4fa990",
@@ -676,7 +955,104 @@
   {
    "cell_type": "code",
    "execution_count": null,
+   "id": "1bd71766-4633-4136-b00a-e42fa304ff92",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "rev_point_map = {}\n",
+    "for (poly, params_cat, k), (points, affine_params) in all_points_filtered.items():\n",
+    "    for point in points:\n",
+    "        poly_set = rev_point_map.setdefault((point, affine_params), set())\n",
+    "        poly_set.add((poly, params_cat, k))\n",
+    "\n",
+    "for (point, affine_params), poly_set in rev_point_map.items():\n",
+    "    if len(poly_set) > 1:\n",
+    "        print(point, affine_params.curve.parameters, f\"p={affine_params.curve.prime}\") \n",
+    "        cond = affine_params.name.split(\"=\")[1].split(\"[\")[0] if \"=\" in affine_params.name else \"\"\n",
+    "        polys_mapped = set()\n",
+    "        for poly, params_cat, k in poly_set:\n",
+    "            mapd = eliminate_y(poly, affine_params.curve.model)\n",
+    "            polys_mapped.add(mapd)\n",
+    "            print(poly.as_expr(), \"|\", params_cat, \"|\", k)\n",
+    "            #for formula, fset in factor_sets.items():\n",
+    "            #    if poly in fset and (cond in formula.coordinate_model.name or \"-\" not in formula.coordinate_model.name):\n",
+    "            #        print(\"\\t\", formula)\n",
+    "        print(\"->\")\n",
+    "        for poly in polys_mapped:\n",
+    "            p = Poly(poly, domain=ZZ)\n",
+    "            print(p.factor_list())\n",
+    "        print(\"------\")\n",
+    "    else:\n",
+    "        print(\".\", end=\"\") #poly_set, point)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d38370e9-c222-4f7e-9a08-d4e299083f87",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "for node in PreOrderIter(tree_remapped.root):\n",
+    "    if node.dmap_input:\n",
+    "        pad = \" \" * node.depth\n",
+    "        poly_set = rev_point_map[node.dmap_input]\n",
+    "        point, affine_params = node.dmap_input\n",
+    "        print(pad, point, affine_params.name, affine_params.curve.parameters, f\"p={affine_params.curve.prime}\") \n",
+    "        chain = chain_map[affine_params]\n",
+    "        print(pad, chain)\n",
+    "        cond = affine_params.name.split(\"=\")[1].split(\"[\")[0] if \"=\" in affine_params.name else \"\"\n",
+    "        true_formulas = set()\n",
+    "        for poly, params_cat, k in poly_set:\n",
+    "            print(pad, poly.as_expr(), \"|\", params_cat, \"|\", k)\n",
+    "            for formula, fset in factor_sets.items():\n",
+    "                if poly in fset and (cond in formula.coordinate_model.name or \"-\" not in formula.coordinate_model.name):\n",
+    "                    kvar = (\"add\", (1, k)) if formula.shortname == \"add\" else (\"dbl\", (k, ))\n",
+    "                    print(pad, \"\\t\", formula, \"*\" if kvar in chain else \"not\")\n",
+    "                    if kvar in chain:\n",
+    "                        true_formulas.add(formula)\n",
+    "        for child in node.children:\n",
+    "            print(pad, child.response)\n",
+    "            for cfg in child.cfgs:\n",
+    "                if child.response == True:\n",
+    "                    print(pad, \"ok\" if true_formulas.intersection(cfg) else \"nok\", cfg)\n",
+    "                elif child.response == False:\n",
+    "                    print(pad, \"ok\" if not true_formulas.intersection(cfg) else \"nok\", cfg)\n",
+    "                else:\n",
+    "                    print(pad, cfg)\n",
+    "            print(\"\")\n",
+    "        print(\"------\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
    "id": "d327e346-9c82-4e50-9357-ba66b3c511ed",
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "for coord_name, coords in model.coordinates.items():\n",
+    "    total = 1\n",
+    "    print(coord_name)\n",
+    "    for formula_group, formula_class in zip(formula_groups[coords], formula_classes):\n",
+    "        lf = len(formula_group)\n",
+    "        print(\"\\t\", formula_class.shortname, lf)\n",
+    "        for formula in formula_group:\n",
+    "            print(\"\\t\", formula)\n",
+    "        total *= lf\n",
+    "    print(\"\\t\", total)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "cbabc614-f571-4c92-a3f3-f201388a5606",
    "metadata": {},
    "outputs": [],
    "source": []