{
  "cells": [
    {
      "cell_type": "markdown",
      "id": "cf69bb3f-94e6-4dba-92cd-ce08df117d67",
      "metadata": {
        "id": "cf69bb3f-94e6-4dba-92cd-ce08df117d67"
      },
      "source": [
        "## Solvency demo\n",
        "\n",
        "Here we create a demo of a solvency calculation in the manner of [summa-solvency](https://github.com/summa-dev/summa-solvency). The aim here is to demonstrate the use of the new kzgcommit method detailed [here](https://blog.ezkl.xyz/post/commits/). \n",
        "\n",
        "In this setup:\n",
        "- the commitments to users, respective balances, and total balance are known are publicly known to the prover and verifier. \n",
        "- We leave the outputs of the model as public as well (known to the  verifier and prover). \n",
        "\n",
        "The circuit calculates the total sum of the balances, and checks that it is less than the total balance which is precommited to."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "95613ee9",
      "metadata": {
        "id": "95613ee9"
      },
      "outputs": [],
      "source": [
        "# check if notebook is in colab\n",
        "try:\n",
        "    # install ezkl\n",
        "    import google.colab\n",
        "    import subprocess\n",
        "    import sys\n",
        "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"ezkl\"])\n",
        "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"onnx\"])\n",
        "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"pytest\"])\n",
        "\n",
        "# rely on local installation of ezkl if the notebook is not in colab\n",
        "except:\n",
        "    pass\n",
        "\n",
        "# uncomment to enable logging\n",
        "# import logging\n",
        "# FORMAT = '%(levelname)s %(name)s %(asctime)-15s %(filename)s:%(lineno)d %(message)s'\n",
        "# logging.basicConfig(format=FORMAT)\n",
        "# logging.getLogger().setLevel(logging.DEBUG)\n",
        "\n",
        "\n",
        "# here we create and (potentially train a model)\n",
        "\n",
        "# make sure you have the dependencies required here already installed\n",
        "from torch import nn\n",
        "import ezkl\n",
        "import os\n",
        "import json\n",
        "import torch\n",
        "\n",
        "class Circuit(nn.Module):\n",
        "    def __init__(self):\n",
        "        super(Circuit, self).__init__()\n",
        "\n",
        "    def forward(self, users, balances, total):\n",
        "        nil = torch.nn.Parameter(torch.tensor([0.0]))\n",
        "        # calculate the total balance across all users second term will be ignored by the optimizer but will force it to be included in a separate col for commitment\n",
        "        balances = torch.sum(balances, dim=1) + nil * users\n",
        "        # now check if the total balance is less than the total\n",
        "        return (balances[:,0] <= total)\n",
        "\n",
        "\n",
        "circuit = Circuit()\n",
        "\n",
        "# Train the model as you like here (skipped for brevity)\n",
        "\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b37637c4",
      "metadata": {
        "id": "b37637c4"
      },
      "outputs": [],
      "source": [
        "model_path = os.path.join('network.onnx')\n",
        "compiled_model_path = os.path.join('network.compiled')\n",
        "pk_path = os.path.join('test.pk')\n",
        "vk_path = os.path.join('test.vk')\n",
        "settings_path = os.path.join('settings.json')\n",
        "\n",
        "witness_path = os.path.join('witness.json')\n",
        "data_path = os.path.join('input.json')"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "1c21e56e",
      "metadata": {},
      "source": [
        "We create dummy data here for the sake of demonstration. In a real world scenario, the data would be provided by the users, and the commitments would be made by some trusted party.\n",
        "\n",
        "The users are generated as hashes of the integers 0 to 9. The balances are generated as integers between 0 and 10. \n",
        "\n",
        "The total balance is the sum of the balances."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "dfd6c7e7",
      "metadata": {},
      "outputs": [],
      "source": [
        "\n",
        "\n",
        "user_preimages = [0.0, 1.0, 2.0, 3.0, 4.0, 9.0]\n",
        "balances = torch.tensor([0, 2, 3, 4, 5, 10])\n",
        "balances = balances.reshape(1, 6)\n",
        "\n",
        "\n",
        "# Create an empty list to store the hashes of float -- which I guess we'll call the users here\n",
        "users = []\n",
        "\n",
        "# Loop through each element in the y tensor\n",
        "for e in user_preimages:\n",
        "    # Apply the custom function and append the result to the list\n",
        "    users.append(ezkl.poseidon_hash([ezkl.float_to_string(e, 0)])[0])\n",
        "\n",
        "users_t = torch.tensor(user_preimages)\n",
        "users_t = users_t.reshape(1, 6)\n",
        "\n",
        "total = torch.tensor([25])\n",
        "total = total.reshape(1, 1)\n",
        "\n",
        "# Flips thegraph into inference mode\n",
        "circuit.eval()\n",
        "\n",
        "    # Export the model\n",
        "torch.onnx.export(circuit,               # model being run\n",
        "                      (users_t,balances,total),                   # model input (or a tuple for multiple inputs)\n",
        "                      model_path,            # where to save the model (can be a file or file-like object)\n",
        "                      export_params=True,        # store the trained parameter weights inside the model file\n",
        "                      opset_version=17,          # the ONNX version to export the model to\n",
        "                      do_constant_folding=False,  # whether to execute constant folding for optimization\n",
        "                      input_names = ['input'],   # the model's input names\n",
        "                      output_names = ['output'], # the model's output names\n",
        "                      dynamic_axes={'input' : {0 : 'batch_size'},    # variable length axes\n",
        "                                    'output' : {0 : 'batch_size'}})\n",
        "\n",
        "\n",
        "    \n",
        "data_array_x = users\n",
        "data_array_y = ((balances).detach().numpy()).reshape([-1]).tolist()\n",
        "data_array_z = ((total).detach().numpy()).reshape([-1]).tolist()\n",
        "\n",
        "data = dict(input_data = [data_array_x, data_array_y, data_array_z])\n",
        "\n",
        "\n",
        "    # Serialize data into file:\n",
        "json.dump( data, open(data_path, 'w' ))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "22d3d8df",
      "metadata": {},
      "source": [
        "This is where the magic happens. We define our `PyRunArgs` objects which contains the visibility parameters for out model. \n",
        "- `input_visibility` defines the visibility of the model inputs\n",
        "- `param_visibility` defines the visibility of the model weights and constants and parameters \n",
        "- `output_visibility` defines the visibility of the model outputs\n",
        "\n",
        "There are currently 5 visibility settings:\n",
        "- `public`: known to both the verifier and prover (a subtle nuance is that this may not be the case for model parameters but until we have more rigorous theoretical results we don't want to make strong claims as to this). \n",
        "- `private`: known only to the prover\n",
        "- `hashed`: the hash pre-image is known to the prover, the prover and verifier know the hash. The prover proves that the they know the pre-image to the hash. \n",
        "- `encrypted`: the non-encrypted element and the secret key used for decryption are known to the prover. The prover and the verifier know the encrypted element, the public key used to encrypt, and the hash of the decryption hey. The prover proves that they know the pre-image of the hashed decryption key and that this key can in fact decrypt the encrypted message.\n",
        "- `kzgcommit`: unblinded advice column which generates a kzg commitment. This doesn't appear in the instances of the circuit and must instead be modified directly within the proof bytes.  \n",
        "\n",
        "Here we create the following setup:\n",
        "- `input_visibility`: \"kzgcommit\"\n",
        "- `param_visibility`: \"public\"\n",
        "- `output_visibility`: public\n",
        "\n",
        "We encourage you to play around with other setups :) \n",
        "\n",
        "Shoutouts: \n",
        "\n",
        "- [summa-solvency](https://github.com/summa-dev/summa-solvency) for their help with the poseidon hashing chip. \n",
        "- [timeofey](https://github.com/timoftime) for providing inspiration in our developement of the el-gamal encryption circuit in Halo2. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "d5e374a2",
      "metadata": {
        "id": "d5e374a2"
      },
      "outputs": [],
      "source": [
        "run_args = ezkl.PyRunArgs()\n",
        "# \"kzgcommit\" means that the output of the hashing is not visible to the verifier and is instead fed into the computational graph\n",
        "run_args.input_visibility = \"kzgcommit\"\n",
        "# the parameters are public\n",
        "run_args.param_visibility = \"fixed\"\n",
        "# the output is public (this is the inequality test)\n",
        "run_args.output_visibility = \"public\"\n",
        "run_args.variables = [(\"batch_size\", 1)]\n",
        "# never rebase the scale\n",
        "run_args.scale_rebase_multiplier = 1000\n",
        "# logrows\n",
        "run_args.logrows = 11\n",
        "run_args.lookup_range = (-1000,1000)\n",
        "run_args.input_scale = 0\n",
        "run_args.param_scale = 0\n",
        "\n",
        "\n",
        "# TODO: Dictionary outputs\n",
        "res = ezkl.gen_settings(model_path, settings_path, py_run_args=run_args)\n",
        "assert res == True\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "3aa4f090",
      "metadata": {
        "id": "3aa4f090"
      },
      "outputs": [],
      "source": [
        "res = ezkl.compile_circuit(model_path, compiled_model_path, settings_path)\n",
        "assert res == True"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8b74dcee",
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "8b74dcee",
        "outputId": "f7b9198c-2b3d-48bb-c67e-8478333cedb5"
      },
      "outputs": [],
      "source": [
        "# srs path\n",
        "res = ezkl.get_srs( settings_path)"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "f7c98c96",
      "metadata": {},
      "source": [
        "We'll generate two proofs, one with the correct total balance, and one with an incorrect total balance.\n",
        "\n",
        "## Correct total balance\n",
        "\n",
        "The data file above has a total balance of above the user total balance. We'll generate a proof with this total balance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "93e90079",
      "metadata": {},
      "outputs": [],
      "source": [
        "# setup keypair\n",
        "res = ezkl.setup(\n",
        "        compiled_model_path,\n",
        "        vk_path,\n",
        "        pk_path,\n",
        "        \n",
        "    )\n",
        "\n",
        "assert res == True\n",
        "assert os.path.isfile(vk_path)\n",
        "assert os.path.isfile(pk_path)\n",
        "assert os.path.isfile(settings_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "efa0ac91",
      "metadata": {},
      "outputs": [],
      "source": [
        "!export RUST_BACKTRACE=1\n",
        "\n",
        "witness_path = \"witness.json\"\n",
        "\n",
        "res = ezkl.gen_witness(data_path, compiled_model_path, witness_path, vk_path)\n",
        "assert os.path.isfile(witness_path)\n",
        "\n",
        "# we force the output to be 1 this corresponds to the solvency test being true -- and we set this to a fixed vis output\n",
        "# this means that the output is fixed and the verifier can see it but that if the input is not in the set the output will not be 0 and the verifier will reject\n",
        "witness = json.load(open(witness_path, \"r\"))\n",
        "witness[\"outputs\"][0] = [ezkl.float_to_string(1.0, 0)]\n",
        "json.dump(witness, open(witness_path, \"w\"))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "90973daf",
      "metadata": {},
      "outputs": [],
      "source": [
        "proof_path = os.path.join('proof.json')\n",
        "# proof path\n",
        "res = ezkl.prove(\n",
        "        witness_path,\n",
        "        compiled_model_path,\n",
        "        pk_path,\n",
        "        proof_path,\n",
        "        \n",
        "        \"single\",\n",
        "    )\n",
        "\n",
        "assert os.path.isfile(proof_path)\n",
        "\n",
        "print(res)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "ef79b4ee",
      "metadata": {},
      "source": [
        "- now we swap the commitments of the proof as a way to demonstrate that the proof is valid given some public inputs \n",
        "- this is just for testing purposes and would require fetching public commits from the blockchain or some other source\n",
        "- see https://blog.ezkl.xyz/post/commits/ for more details"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "8bb46735",
      "metadata": {},
      "outputs": [],
      "source": [
        "\n",
        "\n",
        "res = ezkl.swap_proof_commitments(proof_path, witness_path)\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "59f3b1e9",
      "metadata": {},
      "outputs": [],
      "source": [
        "# verify the proof\n",
        "res = ezkl.verify(\n",
        "        proof_path,\n",
        "        settings_path,\n",
        "        vk_path,\n",
        "        \n",
        "    )\n",
        "assert res == True"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "77dec3dd",
      "metadata": {},
      "source": [
        "### Faulty proof\n",
        "\n",
        "We'll generate a proof with a total balance of 10. This is below the user total balance."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "FQfGdcUNpvuK",
      "metadata": {
        "id": "FQfGdcUNpvuK"
      },
      "outputs": [],
      "source": [
        "# now generate a truthy input + witness file (x input not in the set)\n",
        "import random\n",
        "\n",
        "data_path_truthy = os.path.join('input.json')\n",
        "data = json.load(open(data_path, 'r' ))\n",
        "data['input_data'][2] = [10]\n",
        "\n",
        "data_path_faulty = os.path.join('input_faulty.json')\n",
        "# Serialize data into file:\n",
        "json.dump( data, open(data_path_faulty, 'w' ))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "b1c561a8",
      "metadata": {
        "id": "b1c561a8"
      },
      "outputs": [],
      "source": [
        "# now generate the witness file\n",
        "\n",
        "res = ezkl.gen_witness(data_path_faulty, compiled_model_path, witness_path, vk_path)\n",
        "assert os.path.isfile(witness_path)\n",
        "\n",
        "# we force the output to be 1 this corresponds to the solvency test being true -- and we set this to a fixed vis output\n",
        "# this means that the output is fixed and the verifier can see it but that if the input is not in the set the output will not be 0 and the verifier will reject\n",
        "witness = json.load(open(witness_path, \"r\"))\n",
        "witness[\"outputs\"][0] = [ezkl.float_to_string(1.0, 0)]\n",
        "json.dump(witness, open(witness_path, \"w\"))\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "c384cbc8",
      "metadata": {
        "id": "c384cbc8"
      },
      "outputs": [],
      "source": [
        "# GENERATE A PROOF\n",
        "\n",
        "\n",
        "proof_path = os.path.join('test.pf')\n",
        "\n",
        "res = ezkl.prove(\n",
        "        witness_path,\n",
        "        compiled_model_path,\n",
        "        pk_path,\n",
        "        proof_path,\n",
        "        \n",
        "        \"single\",\n",
        "    )\n",
        "\n",
        "print(res)\n",
        "assert os.path.isfile(proof_path)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4edaca46",
      "metadata": {},
      "outputs": [],
      "source": [
        "res = ezkl.swap_proof_commitments(proof_path, witness_path)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "id": "638d776f",
      "metadata": {},
      "source": [
        "Now we test that verification fails"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "id": "4nqEx7-qpciQ",
      "metadata": {
        "id": "4nqEx7-qpciQ"
      },
      "outputs": [],
      "source": [
        "import pytest\n",
        "\n",
        "def test_verification():\n",
        "    with pytest.raises(RuntimeError, match='Failed to run verify: The constraint system is not satisfied'):\n",
        "        ezkl.verify(\n",
        "            proof_path,\n",
        "            settings_path,\n",
        "            vk_path,\n",
        "            \n",
        "        )\n",
        "\n",
        "# Run the test function\n",
        "test_verification()"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "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.9.13"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}
