{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0577f61d-635c-436c-96ca-1471265ac1fb",
   "metadata": {},
   "source": [
    "# Exploration of formulas in open-source ECC libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "397771ed-7408-47a6-8b82-75e047e100fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "import tabulate\n",
    "import pickle\n",
    "import numpy as np\n",
    "import inspect\n",
    "import tempfile\n",
    "import sys\n",
    "import multiprocessing\n",
    "\n",
    "from copy import deepcopy\n",
    "from concurrent.futures import ProcessPoolExecutor, as_completed\n",
    "from contextlib import contextmanager\n",
    "from importlib import import_module, invalidate_caches\n",
    "from pathlib import Path\n",
    "\n",
    "from matplotlib import pyplot as plt\n",
    "from IPython.display import HTML, display\n",
    "from tqdm.notebook import tqdm\n",
    "from importlib_resources import files, as_file\n",
    "\n",
    "import pyecsca\n",
    "from pyecsca.ec.model import ShortWeierstrassModel, TwistedEdwardsModel, MontgomeryModel\n",
    "from pyecsca.ec.formula import AdditionEFDFormula, DoublingEFDFormula, LadderEFDFormula\n",
    "from pyecsca.ec.params import get_params\n",
    "from pyecsca.ec.formula.metrics import formula_similarity, formula_similarity_fuzz\n",
    "from pyecsca.ec.formula.unroll import unroll_formula_expr, unroll_formula\n",
    "from pyecsca.ec.formula.expand import expand_formula_set\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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1727bd9a-f209-4d91-9286-8cca9118c187",
   "metadata": {},
   "outputs": [],
   "source": [
    "sw = ShortWeierstrassModel()\n",
    "mont = MontgomeryModel()\n",
    "te = TwistedEdwardsModel()\n",
    "\n",
    "curve25519 = get_params(\"other\", \"Curve25519\", \"xz\")\n",
    "ed25519 = get_params(\"other\", \"Ed25519\", \"extended\")\n",
    "p256_jac3 = get_params(\"secg\", \"secp256r1\", \"jacobian-3\")\n",
    "p256_jac = get_params(\"secg\", \"secp256r1\", \"jacobian\")\n",
    "p256_mod = get_params(\"secg\", \"secp256r1\", \"modified\")\n",
    "p256_proj3 = get_params(\"secg\", \"secp256r1\", \"projective-3\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d686ad1-4b20-4327-95be-36d6e92052b3",
   "metadata": {},
   "source": [
    "## The formulas\n",
    "The following formulas were collected from open-source cryptographic libraries and are stored in the tests directory of the pyecsca toolkit."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21572b7d-380c-4570-b5bd-2043306adc3e",
   "metadata": {},
   "outputs": [],
   "source": [
    "lib_formula_defs = [\n",
    "        [\n",
    "            \"add-bc-r1rv76-jac\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"jacobian\",\n",
    "            (\"secg\", \"secp128r1\"),\n",
    "            AdditionEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"add-bc-r1rv76-mod\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"modified\",\n",
    "            (\"secg\", \"secp128r1\"),\n",
    "            AdditionEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"dbl-bc-r1rv76-jac\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"jacobian\",\n",
    "            (\"secg\", \"secp128r1\"),\n",
    "            DoublingEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"dbl-bc-r1rv76-mod\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"modified\",\n",
    "            (\"secg\", \"secp128r1\"),\n",
    "            DoublingEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"dbl-bc-r1rv76-x25519\", #ok\n",
    "            MontgomeryModel,\n",
    "            \"xz\",\n",
    "            (\"other\", \"Curve25519\"),\n",
    "            DoublingEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"ladd-bc-r1rv76-x25519\", #ok\n",
    "            MontgomeryModel,\n",
    "            \"xz\",\n",
    "            (\"other\", \"Curve25519\"),\n",
    "            LadderEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"dbl-boringssl-p224\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"jacobian-3\",\n",
    "            (\"secg\", \"secp224r1\"),\n",
    "            DoublingEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"add-boringssl-p224\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"jacobian-3\",\n",
    "            (\"secg\", \"secp224r1\"),\n",
    "            AdditionEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"add-libressl-v382\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"jacobian\",\n",
    "            (\"secg\", \"secp128r1\"),\n",
    "            AdditionEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"dbl-libressl-v382\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"jacobian\",\n",
    "            (\"secg\", \"secp128r1\"),\n",
    "            DoublingEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"dbl-secp256k1-v040\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"jacobian\",\n",
    "            (\"secg\", \"secp256k1\"),\n",
    "            DoublingEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"add-openssl-z256\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"jacobian-3\",\n",
    "            (\"secg\", \"secp256r1\"),\n",
    "            AdditionEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"add-openssl-z256a\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"jacobian-3\",\n",
    "            (\"secg\", \"secp256r1\"),\n",
    "            AdditionEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"ladd-openssl-x25519\", #ok\n",
    "            MontgomeryModel,\n",
    "            \"xz\",\n",
    "            (\"other\", \"Curve25519\"),\n",
    "            LadderEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"ladd-hacl-x25519\", #ok\n",
    "            MontgomeryModel,\n",
    "            \"xz\",\n",
    "            (\"other\", \"Curve25519\"),\n",
    "            LadderEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"dbl-hacl-x25519\", #ok\n",
    "            MontgomeryModel,\n",
    "            \"xz\",\n",
    "            (\"other\", \"Curve25519\"),\n",
    "            DoublingEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"dbl-sunec-v21\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"projective-3\",\n",
    "            (\"secg\", \"secp256r1\"),\n",
    "            DoublingEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"add-sunec-v21\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"projective-3\",\n",
    "            (\"secg\", \"secp256r1\"),\n",
    "            AdditionEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"add-sunec-v21-ed25519\", #ok\n",
    "            TwistedEdwardsModel,\n",
    "            \"extended-1\",\n",
    "            (\"other\", \"Ed25519\"),\n",
    "            AdditionEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"dbl-sunec-v21-ed25519\", #ok\n",
    "            TwistedEdwardsModel,\n",
    "            \"extended-1\",\n",
    "            (\"other\", \"Ed25519\"),\n",
    "            DoublingEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"ladd-rfc7748\", #ok\n",
    "            MontgomeryModel,\n",
    "            \"xz\",\n",
    "            (\"other\", \"Curve25519\"),\n",
    "            LadderEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"add-bearssl-v06\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"jacobian\",\n",
    "            (\"secg\", \"secp256r1\"),\n",
    "            AdditionEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"dbl-bearssl-v06\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"jacobian\",\n",
    "            (\"secg\", \"secp256r1\"),\n",
    "            DoublingEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"add-libgcrypt-v1102\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"jacobian\",\n",
    "            (\"secg\", \"secp256r1\"),\n",
    "            AdditionEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"dbl-libgcrypt-v1102\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"jacobian\",\n",
    "            (\"secg\", \"secp256r1\"),\n",
    "            DoublingEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"ladd-go-1214\", #ok\n",
    "            MontgomeryModel,\n",
    "            \"xz\",\n",
    "            (\"other\", \"Curve25519\"),\n",
    "            LadderEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"add-gecc-322\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"jacobian-3\",\n",
    "            (\"secg\", \"secp256r1\"),\n",
    "            AdditionEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"dbl-gecc-321\", #ok\n",
    "            ShortWeierstrassModel,\n",
    "            \"jacobian-3\",\n",
    "            (\"secg\", \"secp256r1\"),\n",
    "            DoublingEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"ladd-boringssl-x25519\", #ok\n",
    "            MontgomeryModel,\n",
    "            \"xz\",\n",
    "            (\"other\", \"Curve25519\"),\n",
    "            LadderEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"dbl-ipp-x25519\", #ok\n",
    "            MontgomeryModel,\n",
    "            \"xz\",\n",
    "            (\"other\", \"Curve25519\"),\n",
    "            DoublingEFDFormula,\n",
    "        ],\n",
    "        [\n",
    "            \"ladd-botan-x25519\", #ok\n",
    "            MontgomeryModel,\n",
    "            \"xz\",\n",
    "            (\"other\", \"Curve25519\"),\n",
    "            LadderEFDFormula,\n",
    "        ],\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "176fcb8c-e725-4e3b-9835-729cfb795206",
   "metadata": {},
   "source": [
    "Let's load the formulas now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0cfcbf60-7d6b-4ab0-b957-e998037ff8ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "lib_formulas = {}\n",
    "base_path = files(pyecsca).joinpath(\"../test/data/formulas/\")\n",
    "for formula_def in lib_formula_defs:\n",
    "    meta_path = base_path / formula_def[0]\n",
    "    op3_path = base_path / (formula_def[0] + \".op3\")\n",
    "    model = formula_def[1]()\n",
    "    formula = formula_def[4](meta_path, op3_path, formula_def[0], model.coordinates[formula_def[2]]).to_code()\n",
    "    lib_formulas[formula_def[0]] = formula"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74002746-f3f3-4528-9042-b9ab51791bc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "len(lib_formulas)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5800e043-85bc-4fac-8c40-75c234beec1c",
   "metadata": {},
   "source": [
    "Now we can setup some code for examining the similarities between the formulas."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0a91c61-1565-49bd-9604-5a371710f91a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_similarities(formulas, curve):\n",
    "    table = [[\"One\", \"Other\", \"Similarity (output)\", \"Similarity (IV)\"]]\n",
    "\n",
    "    im_iv = np.zeros((len(formulas), len(formulas)))\n",
    "    im_out = np.zeros((len(formulas), len(formulas)))\n",
    "    for formula in formulas:\n",
    "        if formula.assumptions:\n",
    "            for assumption_str in formula.assumptions_str:\n",
    "                lhs, rhs = assumption_str.strip().split(\" == \")\n",
    "                if lhs in formula.inputs:\n",
    "                    print(f\"Warning, formula {formula.name} has assumptions: {assumption_str}\")\n",
    "    for one, other in itertools.product(formulas, formulas):\n",
    "        i = formulas.index(one)\n",
    "        j = formulas.index(other)\n",
    "        if curve is None:\n",
    "            sim = formula_similarity(one, other)\n",
    "        else:\n",
    "            sim = formula_similarity_fuzz(one, other, curve, 100)\n",
    "        im_iv[i, j] = sim[\"ivs\"]\n",
    "        im_out[i, j] = sim[\"output\"]\n",
    "        table.append([one.name, other.name, f\"{sim['output']:.2}\", f\"{sim['ivs']:.2}\"])\n",
    "    \n",
    "    return table, im_iv, im_out\n",
    "\n",
    "def plot_similarities(formulas, im_data, name):\n",
    "    fig, ax = plt.subplots()\n",
    "    im = ax.imshow(im_data, vmin=0)\n",
    "    cbar_ax = fig.add_axes((0.90, 0.11, 0.04, 0.76))\n",
    "    cbar = fig.colorbar(im, cax=cbar_ax)\n",
    "    cbar.ax.set_ylabel(f\"Similarity ({name})\", rotation=-90, va=\"bottom\")\n",
    "    \n",
    "    ax.set_xticks(np.arange(len(formulas)), labels=list(map(lambda f: f.name, formulas)), rotation=90)\n",
    "    ax.set_yticks(np.arange(len(formulas)), labels=list(map(lambda f: f.name, formulas)), rotation=0)\n",
    "    \n",
    "    for i, one in enumerate(formulas):\n",
    "        for j, other in enumerate(formulas):\n",
    "            ax.text(j, i, f\"{im_data[i, j]:.2}\", ha=\"center\", va=\"center\", color=\"white\")\n",
    "    plt.show()\n",
    "\n",
    "def analyze_formulas(formulas, curve=None):\n",
    "    table, im_iv, im_out = compute_similarities(formulas, curve)\n",
    "    plot_similarities(formulas, im_iv, \"IV\")\n",
    "    plot_similarities(formulas, im_out, \"output\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "392d7eea-abcb-4d83-82b4-a9394d3d44aa",
   "metadata": {},
   "source": [
    "## Analysis"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "932e8ec8-77f6-4079-909c-3ca021fde717",
   "metadata": {},
   "outputs": [],
   "source": [
    "xz_ladders = [formula for formula in mont.coordinates[\"xz\"].formulas.values() if formula.name.startswith(\"ladd\") or formula.name.startswith(\"mladd\")] + [\n",
    "    lib_formulas[\"ladd-rfc7748\"], lib_formulas[\"ladd-hacl-x25519\"], lib_formulas[\"ladd-openssl-x25519\"], lib_formulas[\"ladd-bc-r1rv76-x25519\"],\n",
    "    lib_formulas[\"ladd-go-1214\"], lib_formulas[\"ladd-boringssl-x25519\"], lib_formulas[\"ladd-botan-x25519\"]]\n",
    "analyze_formulas(xz_ladders)\n",
    "analyze_formulas(xz_ladders, curve25519.curve)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2e0e3b5-b7e8-4ed7-94a1-a74976b108fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "xz_dbls = [formula for formula in mont.coordinates[\"xz\"].formulas.values() if formula.name.startswith(\"dbl\")] + [\n",
    "    lib_formulas[\"dbl-bc-r1rv76-x25519\"], lib_formulas[\"dbl-hacl-x25519\"], lib_formulas[\"dbl-ipp-x25519\"]]\n",
    "analyze_formulas(xz_dbls)\n",
    "analyze_formulas(xz_dbls, curve25519.curve)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "724791a0-4137-4d3e-abaa-a90e4e2187ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "jac3_adds = [formula for formula in sw.coordinates[\"jacobian-3\"].formulas.values() if formula.name.startswith(\"add\") or formula.name.startswith(\"madd\")] + [\n",
    "    lib_formulas[\"add-boringssl-p224\"], lib_formulas[\"add-openssl-z256\"], lib_formulas[\"add-openssl-z256a\"],\n",
    "    lib_formulas[\"add-gecc-322\"]]\n",
    "jac3_adds_fixed = []\n",
    "for formula in jac3_adds:\n",
    "    if formula.coordinate_model != sw.coordinates[\"jacobian-3\"]:\n",
    "        formula = deepcopy(formula)\n",
    "        formula.coordinate_model = sw.coordinates[\"jacobian-3\"]\n",
    "    jac3_adds_fixed.append(formula)\n",
    "analyze_formulas(jac3_adds_fixed)\n",
    "analyze_formulas(jac3_adds_fixed, p256_jac3.curve)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d12b98bc-71b5-4800-8762-ca87eb02d050",
   "metadata": {},
   "outputs": [],
   "source": [
    "jac3_dbls = [formula for formula in sw.coordinates[\"jacobian-3\"].formulas.values() if formula.name.startswith(\"dbl\")] + [\n",
    "        lib_formulas[\"dbl-boringssl-p224\"], lib_formulas[\"dbl-gecc-321\"]]\n",
    "jac3_dbls_fixed = []\n",
    "for formula in jac3_dbls:\n",
    "    if formula.coordinate_model != sw.coordinates[\"jacobian-3\"]:\n",
    "        formula = deepcopy(formula)\n",
    "        formula.coordinate_model = sw.coordinates[\"jacobian-3\"]\n",
    "    jac3_dbls_fixed.append(formula)\n",
    "analyze_formulas(jac3_dbls_fixed)\n",
    "analyze_formulas(jac3_dbls_fixed, p256_jac3.curve)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae0cde46-061c-48f1-9fb6-67a2798c5068",
   "metadata": {},
   "outputs": [],
   "source": [
    "mod_adds = [formula for formula in sw.coordinates[\"modified\"].formulas.values() if formula.name.startswith(\"add\") or formula.name.startswith(\"madd\")] + [\n",
    "    lib_formulas[\"add-bc-r1rv76-mod\"]]\n",
    "analyze_formulas(mod_adds)\n",
    "analyze_formulas(mod_adds, p256_mod.curve)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7226513-9866-405f-9f4e-ad83134b957a",
   "metadata": {},
   "outputs": [],
   "source": [
    "mod_dbls = [formula for formula in sw.coordinates[\"modified\"].formulas.values() if formula.name.startswith(\"dbl\")] + [\n",
    "    lib_formulas[\"dbl-bc-r1rv76-mod\"]]\n",
    "analyze_formulas(mod_dbls)\n",
    "analyze_formulas(mod_dbls, p256_mod.curve)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "313570ca-082e-46df-8420-ed47dc78855b",
   "metadata": {},
   "outputs": [],
   "source": [
    "jac_adds = [formula for formula in sw.coordinates[\"jacobian\"].formulas.values() if formula.name.startswith(\"add\")] + [\n",
    "    lib_formulas[\"add-bc-r1rv76-jac\"], lib_formulas[\"add-libressl-v382\"], lib_formulas[\"add-bearssl-v06\"], lib_formulas[\"add-libgcrypt-v1102\"]]\n",
    "analyze_formulas(jac_adds)\n",
    "analyze_formulas(jac_adds, p256_jac.curve)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a0e3899-872d-411b-8778-362e464390df",
   "metadata": {},
   "outputs": [],
   "source": [
    "jac_dbls = [formula for formula in sw.coordinates[\"jacobian\"].formulas.values() if formula.name.startswith(\"dbl\")] + [\n",
    "    lib_formulas[\"dbl-secp256k1-v040\"], lib_formulas[\"dbl-libressl-v382\"], lib_formulas[\"dbl-bearssl-v06\"], lib_formulas[\"dbl-libgcrypt-v1102\"]]\n",
    "analyze_formulas(jac_dbls)\n",
    "analyze_formulas(jac_dbls, p256_jac.curve)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "881c383e-3062-4f1e-98e3-adc336749a20",
   "metadata": {},
   "outputs": [],
   "source": [
    "proj3_adds = [formula for formula in sw.coordinates[\"projective-3\"].formulas.values() if formula.name.startswith(\"add\") or formula.name.startswith(\"madd\")] + [\n",
    "    lib_formulas[\"add-sunec-v21\"]]\n",
    "analyze_formulas(proj3_adds)\n",
    "analyze_formulas(proj3_adds, p256_proj3.curve)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d34e20a4-bf1b-4245-bbb2-a53fdd3446a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "proj3_dbls = [formula for formula in sw.coordinates[\"projective-3\"].formulas.values() if formula.name.startswith(\"dbl\")] + [\n",
    "    lib_formulas[\"dbl-sunec-v21\"]]\n",
    "analyze_formulas(proj3_dbls)\n",
    "analyze_formulas(proj3_dbls, p256_proj3.curve)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78f7ff19-9f97-4459-b31e-0d7efe3b7dfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "ext_adds = [formula for formula in te.coordinates[\"extended-1\"].formulas.values() if formula.name.startswith(\"add\") or formula.name.startswith(\"madd\")] + [\n",
    "    lib_formulas[\"add-sunec-v21-ed25519\"]]\n",
    "analyze_formulas(ext_adds)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fb60bd5f-47a6-4a8f-9864-66bc750bfe97",
   "metadata": {},
   "outputs": [],
   "source": [
    "ext_dbls = [formula for formula in te.coordinates[\"extended-1\"].formulas.values() if formula.name.startswith(\"dbl\")] + [\n",
    "    lib_formulas[\"dbl-sunec-v21-ed25519\"]]\n",
    "analyze_formulas(ext_dbls)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39f642c2-2350-4310-b365-3be6c3f6fcdb",
   "metadata": {},
   "source": [
    "## Expand\n",
    "We can also expand the set of formulas by applying various transformations (like swapping the order of commutative operations).\n",
    "\n",
    "Note that this computation is parallelized and you should set an appropriate number of workers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f3c2c134-a42d-4bb2-b729-0ec1ddcbcdef",
   "metadata": {},
   "outputs": [],
   "source": [
    "def expand_out(formulas, name):\n",
    "    from pyecsca.ec.formula.expand import expand_formula_set\n",
    "    import pickle\n",
    "\n",
    "    expanded = expand_formula_set(formulas)\n",
    "    with open(name, \"wb\") as f:\n",
    "        pickle.dump(expanded, f)\n",
    "    return name, len(expanded)\n",
    "\n",
    "# 22 is really a maximum usable number here, as there are only 22 \"jobs\".\n",
    "context = multiprocessing.get_context(\"spawn\")\n",
    "with ProcessPoolExecutor(max_workers=22, mp_context=context) as pool, enable_spawn(expand_out) as expand_func:\n",
    "    futures = []\n",
    "    args = []\n",
    "    for coord_name, coords in tqdm(sw.coordinates.items(), desc=\"Submitting\"):\n",
    "        adds = [formula for formula in coords.formulas.values() if formula.name.startswith(\"add\")]\n",
    "        lib_adds = [formula for formula in lib_formulas.values() if formula.coordinate_model == coords and formula.name.startswith(\"add\")]\n",
    "        dbls = [formula for formula in coords.formulas.values() if formula.name.startswith(\"dbl\")]\n",
    "        lib_dbls = [formula for formula in lib_formulas.values() if formula.coordinate_model == coords and formula.name.startswith(\"dbl\")]\n",
    "        futures.append(pool.submit(expand_func, adds + lib_adds, f\"sw_{coord_name}_adds.pickle\"))\n",
    "        args.append(f\"sw_{coord_name}_adds.pickle\")\n",
    "        futures.append(pool.submit(expand_func, dbls + lib_dbls, f\"sw_{coord_name}_dbls.pickle\"))\n",
    "        args.append(f\"sw_{coord_name}_dbls.pickle\")\n",
    "    for future in tqdm(as_completed(futures), total=len(futures), smoothing=0, desc=\"Computing\"):\n",
    "        j = futures.index(future)\n",
    "        arg = args[j]\n",
    "        error = future.exception()\n",
    "        if error:\n",
    "            print(arg, error)\n",
    "        else:\n",
    "            res = future.result()\n",
    "            print(*res)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5b1ccd8-1b2b-4e1a-85e1-5cab2b45300c",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}