Plonk.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "39ae7f61-02f6-447c-93c4-bf0aa1d43499",
   "metadata": {},
   "outputs": [],
   "source": [
    "from compiler.program import Program\n",
    "from setup import Setup\n",
    "from prover import Prover\n",
    "from test.mini_poseidon import rc, mds, poseidon_hash\n",
    "from utils import *\n",
    "\n",
    "def prover_test():\n",
    "    print(\"Beginning prover test\")\n",
    "    # powers should be 2^n so that we can use roots of unity for FFT\n",
    "    # and should be bigger than len(coeffs) of polynomial to do KZG commitment\n",
    "    # the value here is: powers = 4 * group_order\n",
    "    # which is bigger than the order of quotient polynomial\n",
    "    group_order = 8\n",
    "    powers = group_order * 4\n",
    "    setup = Setup.generate_srs(powers)\n",
    "\n",
    "    program = Program([\"e public\", \"c <== a * b\", \"e <== c * d\"], group_order)\n",
    "    assignments = {\"a\": 3, \"b\": 4, \"c\": 12, \"d\": 5, \"e\": 60}\n",
    "    prover = Prover(setup, program)\n",
    "    proof = prover.prove(assignments)\n",
    "    print(\"Prover test success\")\n",
    "    return setup, proof, group_order"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fdd7202f-2e68-4d80-a3ea-cafcb1997024",
   "metadata": {},
   "source": [
    "Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "76a78f5b-fc3e-4945-98c1-ff76cd17a3e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import *\n",
    "import py_ecc.bn128 as b\n",
    "from curve import ec_lincomb, G1Point, G2Point\n",
    "from compiler.program import CommonPreprocessedInput\n",
    "from verifier import VerificationKey\n",
    "from dataclasses import dataclass\n",
    "from poly import Polynomial, Basis\n",
    "\n",
    "@dataclass\n",
    "class Setup(object):\n",
    "    # https://github.com/sec-bit/learning-zkp/blob/develop/plonk-intro-cn/plonk-polycom.md#kzg10-%E6%9E%84%E9%80%A0\n",
    "    #   ([1]₁, [x]₁, ..., [x^{d-1}]₁)\n",
    "    # = ( G,    xG,  ...,  x^{d-1}G ), where G is a generator of G_1\n",
    "    powers_of_x: list[G1Point]\n",
    "    # [x]₂ = xH, where H is a generator of G_2\n",
    "    X2: G2Point\n",
    "\n",
    "    @classmethod\n",
    "    def generate_srs(cls, powers: int):\n",
    "        print(\"Start to generate structured reference string\")\n",
    "        # tau is a random number whatever you choose\n",
    "        tau = 218313819403157342856071133\n",
    "\n",
    "        # Initialize powers_of_x with 0 values\n",
    "        powers_of_x = [0] * powers\n",
    "        # powers_of_x[0] =  b.G1 * tau**0 = b.G1\n",
    "        # powers_of_x[1] =  b.G1 * tau**1 = powers_of_x[0] * tau\n",
    "        # powers_of_x[2] =  b.G1 * tau**2 = powers_of_x[1] * tau\n",
    "        # ...\n",
    "        # powers_of_x[i] =  b.G1 * tau**i = powers_of_x[i - 1] * tau\n",
    "        powers_of_x[0] = b.G1\n",
    "\n",
    "        for i in range(powers):\n",
    "            if i > 0:\n",
    "                powers_of_x[i] = b.multiply(powers_of_x[i - 1], tau)\n",
    "\n",
    "        assert b.is_on_curve(powers_of_x[1], b.b)\n",
    "        print(\"Generated G1 side, X^1 point: {}\".format(powers_of_x[1]))\n",
    "\n",
    "        X2 = b.multiply(b.G2, tau)\n",
    "        assert b.is_on_curve(X2, b.b2)\n",
    "        print(\"Generated G2 side, X^1 point: {}\".format(X2))\n",
    "\n",
    "        assert b.pairing(b.G2, powers_of_x[1]) == b.pairing(X2, b.G1)\n",
    "        print(\"X^1 points checked consistent\")\n",
    "        print(\"Finished to generate structured reference string\")\n",
    "\n",
    "        return cls(powers_of_x, X2)\n",
    "\n",
    "    # Encodes the KZG commitment that evaluates to the given values in the group\n",
    "    def commit(self, values: Polynomial) -> G1Point:\n",
    "        if (values.basis == Basis.LAGRANGE):\n",
    "            # inverse FFT from Lagrange basis to monomial basis\n",
    "            coeffs = values.ifft().values\n",
    "        elif (values.basis == Basis.MONOMIAL):\n",
    "            coeffs = values.values\n",
    "        if len(coeffs) > len(self.powers_of_x):\n",
    "            raise Exception(\"Not enough powers in setup\")\n",
    "        return ec_lincomb([(s, x) for s, x in zip(self.powers_of_x, coeffs)])\n",
    "\n",
    "    # Generate the verification key for this program with the given setup\n",
    "    def verification_key(self, pk: CommonPreprocessedInput) -> VerificationKey:\n",
    "        return VerificationKey(\n",
    "            pk.group_order,\n",
    "            self.commit(pk.QM),\n",
    "            self.commit(pk.QL),\n",
    "            self.commit(pk.QR),\n",
    "            self.commit(pk.QO),\n",
    "            self.commit(pk.QC),\n",
    "            self.commit(pk.S1),\n",
    "            self.commit(pk.S2),\n",
    "            self.commit(pk.S3),\n",
    "            self.X2,\n",
    "            Scalar.root_of_unity(pk.group_order),\n",
    "        )\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbb13d5e-950e-45d0-ad82-d6025c6c1ab7",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Start to generate structured reference string\n",
      "Generated G1 side, X^1 point: (13294353531659665076299264371299131321133377949180224052095139292042656767801, 7244526365924412580786759495774941482824109386590049888405102905649868841718)\n",
      "Generated G2 side, X^1 point: ((19152636372783811233630472865897092822704646270638236258574538719932990329615, 5981775420279756813368727653284174372010905976217595528558124345993479956855), (15799937923252396087061091029963171915696870100807527253118347696384412331186, 9069635832515501441369801349489757510092424196796670011995129268707448587614))\n"
     ]
    }
   ],
   "source": [
    "    # powers should be 2^n so that we can use roots of unity for FFT\n",
    "    # and should be bigger than len(coeffs) of polynomial to do KZG commitment\n",
    "    # the value here is: powers = 4 * group_order\n",
    "    # which is bigger than the order of quotient polynomial\n",
    "    group_order = 8\n",
    "    powers = group_order * 4\n",
    "    setup = Setup.generate_srs(powers)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "d2beed06-e2fc-4ce7-ae43-61adebd134cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils import *\n",
    "from enum import Enum\n",
    "from dataclasses import dataclass\n",
    "\n",
    "\n",
    "class Column(Enum):\n",
    "    LEFT = 1\n",
    "    RIGHT = 2\n",
    "    OUTPUT = 3\n",
    "\n",
    "    def __lt__(self, other):\n",
    "        if self.__class__ is other.__class__:\n",
    "            return self.value < other.value\n",
    "        return NotImplemented\n",
    "\n",
    "    @staticmethod\n",
    "    def variants():\n",
    "        return [Column.LEFT, Column.RIGHT, Column.OUTPUT]\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Cell:\n",
    "    column: Column\n",
    "    row: int\n",
    "\n",
    "    def __key(self):\n",
    "        return (self.row, self.column.value)\n",
    "\n",
    "    def __hash__(self):\n",
    "        return hash(self.__key())\n",
    "\n",
    "    def __lt__(self, other):\n",
    "        if self.__class__ is other.__class__:\n",
    "            return self.__key() < other.__key()\n",
    "        return NotImplemented\n",
    "\n",
    "    def __repr__(self) -> str:\n",
    "        return \"(\" + str(self.row) + \", \" + str(self.column.value) + \")\"\n",
    "\n",
    "    def __str__(self) -> str:\n",
    "        return \"(\" + str(self.row) + \", \" + str(self.column.value) + \")\"\n",
    "\n",
    "    # Outputs the label (an inner-field element) representing a given\n",
    "    # (column, row) pair. Expects section = 1 for left, 2 right, 3 output\n",
    "    def label(self, group_order: int) -> Scalar:\n",
    "        assert self.row < group_order\n",
    "        return Scalar.roots_of_unity(group_order)[self.row] * self.column.value\n",
    "\n",
    "\n",
    "# Gets the key to use in the coeffs dictionary for the term for key1*key2,\n",
    "# where key1 and key2 can be constant(''), a variable, or product keys\n",
    "# Note that degrees higher than 2 are disallowed in the compiler, but we\n",
    "# still allow them in the parser in case we find a way to compile them later\n",
    "def get_product_key(key1, key2):\n",
    "    members = sorted((key1 or \"\").split(\"*\") + (key2 or \"\").split(\"*\"))\n",
    "    return \"*\".join([x for x in members if x])\n",
    "\n",
    "\n",
    "def is_valid_variable_name(name: str) -> bool:\n",
    "    return len(name) > 0 and name.isalnum() and name[0] not in \"0123456789\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "83edadde-8410-466f-9b63-50d8e0a3eb67",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from typing import Optional\n",
    "from dataclasses import dataclass\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class GateWires:\n",
    "    \"\"\"Variable names for Left, Right, and Output wires.\"\"\"\n",
    "\n",
    "    L: Optional[str]\n",
    "    R: Optional[str]\n",
    "    O: Optional[str]\n",
    "\n",
    "    def as_list(self) -> list[Optional[str]]:\n",
    "        return [self.L, self.R, self.O]\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Gate:\n",
    "    \"\"\"Gate polynomial\"\"\"\n",
    "\n",
    "    L: Scalar\n",
    "    R: Scalar\n",
    "    M: Scalar\n",
    "    O: Scalar\n",
    "    C: Scalar\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class AssemblyEqn:\n",
    "    \"\"\"Assembly equation mapping wires to coefficients.\"\"\"\n",
    "\n",
    "    wires: GateWires\n",
    "    coeffs: dict[Optional[str], int]\n",
    "\n",
    "    def L(self) -> Scalar:\n",
    "        return Scalar(-self.coeffs.get(self.wires.L, 0))\n",
    "\n",
    "    def R(self) -> Scalar:\n",
    "        if self.wires.R != self.wires.L:\n",
    "            return Scalar(-self.coeffs.get(self.wires.R, 0))\n",
    "        return Scalar(0)\n",
    "\n",
    "    def C(self) -> Scalar:\n",
    "        return Scalar(-self.coeffs.get(\"\", 0))\n",
    "\n",
    "    def O(self) -> Scalar:\n",
    "        return Scalar(self.coeffs.get(\"$output_coeff\", 1))\n",
    "\n",
    "    def M(self) -> Scalar:\n",
    "        if None not in self.wires.as_list():\n",
    "            return Scalar(\n",
    "                -self.coeffs.get(get_product_key(self.wires.L, self.wires.R), 0)\n",
    "            )\n",
    "        return Scalar(0)\n",
    "\n",
    "    def gate(self) -> Gate:\n",
    "        return Gate(self.L(), self.R(), self.M(), self.O(), self.C())\n",
    "\n",
    "\n",
    "# Converts a arithmetic expression containing numbers, variables and {+, -, *}\n",
    "# into a mapping of term to coefficient\n",
    "#\n",
    "# For example:\n",
    "# ['a', '+', 'b', '*', 'c', '*', '5'] becomes {'a': 1, 'b*c': 5}\n",
    "#\n",
    "# Note that this is a recursive algo, so the input can be a mix of tokens and\n",
    "# mapping expressions\n",
    "#\n",
    "def evaluate(exprs: list[str], first_is_negative=False) -> dict[Optional[str], int]:\n",
    "    # Splits by + and - first, then *, to follow order of operations\n",
    "    # The first_is_negative flag helps us correctly interpret expressions\n",
    "    # like 6000 - 700 - 80 + 9 (that's 5229)\n",
    "    if \"+\" in exprs:\n",
    "        L = evaluate(exprs[: exprs.index(\"+\")], first_is_negative)\n",
    "        R = evaluate(exprs[exprs.index(\"+\") + 1 :], False)\n",
    "        return {x: L.get(x, 0) + R.get(x, 0) for x in set(L.keys()).union(R.keys())}\n",
    "    elif \"-\" in exprs:\n",
    "        L = evaluate(exprs[: exprs.index(\"-\")], first_is_negative)\n",
    "        R = evaluate(exprs[exprs.index(\"-\") + 1 :], True)\n",
    "        return {x: L.get(x, 0) + R.get(x, 0) for x in set(L.keys()).union(R.keys())}\n",
    "    elif \"*\" in exprs:\n",
    "        L = evaluate(exprs[: exprs.index(\"*\")], first_is_negative)\n",
    "        R = evaluate(exprs[exprs.index(\"*\") + 1 :], first_is_negative)\n",
    "        o = {}\n",
    "        for k1 in L.keys():\n",
    "            for k2 in R.keys():\n",
    "                o[get_product_key(k1, k2)] = L[k1] * R[k2]\n",
    "        return o\n",
    "    elif len(exprs) > 1:\n",
    "        raise Exception(\"No ops, expected sub-expr to be a unit: {}\".format(exprs[1]))\n",
    "    elif exprs[0][0] == \"-\":\n",
    "        return evaluate([exprs[0][1:]], not first_is_negative)\n",
    "    elif exprs[0].isnumeric():\n",
    "        return {\"\": int(exprs[0]) * (-1 if first_is_negative else 1)}\n",
    "    elif is_valid_variable_name(exprs[0]):\n",
    "        return {exprs[0]: -1 if first_is_negative else 1}\n",
    "    else:\n",
    "        raise Exception(\"ok wtf is {}\".format(exprs[0]))\n",
    "\n",
    "\n",
    "# Converts an equation to a mapping of term to coefficient, and verifies that\n",
    "# the operations in the equation are valid.\n",
    "#\n",
    "# Also outputs a triple containing the L and R input variables and the output\n",
    "# variable\n",
    "#\n",
    "# Think of the list of (variable triples, coeffs) pairs as this language's\n",
    "# version of \"assembly\"\n",
    "#\n",
    "# Example valid equations, and output:\n",
    "# a === 9                      ([None, None, 'a'], {'': 9})\n",
    "# b <== a * c                  (['a', 'c', 'b'], {'a*c': 1})\n",
    "# d <== a * c - 45 * a + 987   (['a', 'c', 'd'], {'a*c': 1, 'a': -45, '': 987})\n",
    "#\n",
    "# Example invalid equations:\n",
    "# 7 === 7                      # Can't assign to non-variable\n",
    "# a <== b * * c                # Two times signs in a row\n",
    "# e <== a + b * c * d          # Multiplicative degree > 2\n",
    "#\n",
    "def eq_to_assembly(eq: str) -> AssemblyEqn:\n",
    "    tokens = eq.rstrip(\"\\n\").split(\" \")\n",
    "    if tokens[1] in (\"<==\", \"===\"):\n",
    "        # First token is the output variable\n",
    "        out = tokens[0]\n",
    "        # Convert the expression to coefficient map form\n",
    "        coeffs = evaluate(tokens[2:])\n",
    "        # Handle the \"-x === a * b\" case\n",
    "        if out[0] == \"-\":\n",
    "            out = out[1:]\n",
    "            coeffs[\"$output_coeff\"] = -1\n",
    "        # Check out variable name validity\n",
    "        if not is_valid_variable_name(out):\n",
    "            raise Exception(\"Invalid out variable name: {}\".format(out))\n",
    "        # Gather list of variables used in the expression\n",
    "        variables = []\n",
    "        for t in tokens[2:]:\n",
    "            var = t.lstrip(\"-\")\n",
    "            if is_valid_variable_name(var) and var not in variables:\n",
    "                variables.append(var)\n",
    "        # Construct the list of allowed coefficients\n",
    "        allowed_coeffs = variables + [\"\", \"$output_coeff\"]\n",
    "        if len(variables) == 0:\n",
    "            pass\n",
    "        elif len(variables) == 1:\n",
    "            variables.append(variables[0])\n",
    "            allowed_coeffs.append(get_product_key(*variables))\n",
    "        elif len(variables) == 2:\n",
    "            allowed_coeffs.append(get_product_key(*variables))\n",
    "        else:\n",
    "            raise Exception(\"Max 2 variables, found {}\".format(variables))\n",
    "        # Check that only allowed coefficients are in the coefficient map\n",
    "        for key in coeffs.keys():\n",
    "            if key not in allowed_coeffs:\n",
    "                raise Exception(\"Disallowed multiplication: {}\".format(key))\n",
    "        # Return output\n",
    "        wires = variables + [None] * (2 - len(variables)) + [out]\n",
    "        return AssemblyEqn(GateWires(wires[0], wires[1], wires[2]), coeffs)\n",
    "    elif tokens[1] == \"public\":\n",
    "        return AssemblyEqn(\n",
    "            GateWires(tokens[0], None, None),\n",
    "            {tokens[0]: -1, \"$output_coeff\": 0, \"$public\": True},\n",
    "        )\n",
    "    else:\n",
    "        raise Exception(\"Unsupported op: {}\".format(tokens[1]))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "e2723ce4-5bbc-4223-b0c7-0b19e1712406",
   "metadata": {},
   "outputs": [],
   "source": [
    "# A simple zk language, reverse-engineered to match https://zkrepl.dev/ output\n",
    "\n",
    "from typing import Optional, Set\n",
    "from poly import Polynomial, Basis\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class CommonPreprocessedInput:\n",
    "    \"\"\"Common preprocessed input\"\"\"\n",
    "\n",
    "    group_order: int\n",
    "    # q_M(X) multiplication selector polynomial\n",
    "    QM: Polynomial\n",
    "    # q_L(X) left selector polynomial\n",
    "    QL: Polynomial\n",
    "    # q_R(X) right selector polynomial\n",
    "    QR: Polynomial\n",
    "    # q_O(X) output selector polynomial\n",
    "    QO: Polynomial\n",
    "    # q_C(X) constants selector polynomial\n",
    "    QC: Polynomial\n",
    "    # S_σ1(X) first permutation polynomial S_σ1(X)\n",
    "    S1: Polynomial\n",
    "    # S_σ2(X) second permutation polynomial S_σ2(X)\n",
    "    S2: Polynomial\n",
    "    # S_σ3(X) third permutation polynomial S_σ3(X)\n",
    "    S3: Polynomial\n",
    "\n",
    "\n",
    "class Program:\n",
    "    constraints: list[AssemblyEqn]\n",
    "    group_order: int\n",
    "\n",
    "    def __init__(self, constraints: list[str], group_order: int):\n",
    "        if len(constraints) > group_order:\n",
    "            raise Exception(\"Group order too small\")\n",
    "        assembly = [eq_to_assembly(constraint) for constraint in constraints]\n",
    "        self.constraints = assembly\n",
    "        self.group_order = group_order\n",
    "\n",
    "    def common_preprocessed_input(self) -> CommonPreprocessedInput:\n",
    "        L, R, M, O, C = self.make_gate_polynomials()\n",
    "        S = self.make_s_polynomials()\n",
    "        return CommonPreprocessedInput(\n",
    "            self.group_order,\n",
    "            M,\n",
    "            L,\n",
    "            R,\n",
    "            O,\n",
    "            C,\n",
    "            S[Column.LEFT],\n",
    "            S[Column.RIGHT],\n",
    "            S[Column.OUTPUT],\n",
    "        )\n",
    "\n",
    "    @classmethod\n",
    "    def from_str(cls, constraints: str, group_order: int):\n",
    "        lines = [line.strip() for line in constraints.split(\"\\n\")]\n",
    "        return cls(lines, group_order)\n",
    "\n",
    "    def coeffs(self) -> list[dict[Optional[str], int]]:\n",
    "        return [constraint.coeffs for constraint in self.constraints]\n",
    "\n",
    "    def wires(self) -> list[GateWires]:\n",
    "        return [constraint.wires for constraint in self.constraints]\n",
    "\n",
    "    def make_s_polynomials(self) -> dict[Column, Polynomial]:\n",
    "        # For each variable, extract the list of (column, row) positions\n",
    "        # where that variable is used\n",
    "        variable_uses: dict[Optional[str], Set[Cell]] = {None: set()}\n",
    "        for row, constraint in enumerate(self.constraints):\n",
    "            for column, value in zip(Column.variants(), constraint.wires.as_list()):\n",
    "                if value not in variable_uses:\n",
    "                    variable_uses[value] = set()\n",
    "                variable_uses[value].add(Cell(column, row))\n",
    "\n",
    "        # Mark unused cells\n",
    "        for row in range(len(self.constraints), self.group_order):\n",
    "            for column in Column.variants():\n",
    "                variable_uses[None].add(Cell(column, row))\n",
    "\n",
    "        # For each list of positions, rotate by one.\n",
    "        #\n",
    "        # For example, if some variable is used in positions\n",
    "        # (LEFT, 4), (LEFT, 7) and (OUTPUT, 2), then we store:\n",
    "        #\n",
    "        # at S[LEFT][7] the field element representing (LEFT, 4)\n",
    "        # at S[OUTPUT][2] the field element representing (LEFT, 7)\n",
    "        # at S[LEFT][4] the field element representing (OUTPUT, 2)\n",
    "\n",
    "        S_values = {\n",
    "            Column.LEFT: [Scalar(0)] * self.group_order,\n",
    "            Column.RIGHT: [Scalar(0)] * self.group_order,\n",
    "            Column.OUTPUT: [Scalar(0)] * self.group_order,\n",
    "        }\n",
    "\n",
    "        for _, uses in variable_uses.items():\n",
    "            sorted_uses = sorted(uses)\n",
    "            for i, cell in enumerate(sorted_uses):\n",
    "                next_i = (i + 1) % len(sorted_uses)\n",
    "                next_column = sorted_uses[next_i].column\n",
    "                next_row = sorted_uses[next_i].row\n",
    "                S_values[next_column][next_row] = cell.label(self.group_order)\n",
    "\n",
    "        S = {}\n",
    "        S[Column.LEFT] = Polynomial(S_values[Column.LEFT], Basis.LAGRANGE)\n",
    "        S[Column.RIGHT] = Polynomial(S_values[Column.RIGHT], Basis.LAGRANGE)\n",
    "        S[Column.OUTPUT] = Polynomial(S_values[Column.OUTPUT], Basis.LAGRANGE)\n",
    "\n",
    "        return S\n",
    "\n",
    "    # Get the list of public variable assignments, in order\n",
    "    def get_public_assignments(self) -> list[Optional[str]]:\n",
    "        coeffs = self.coeffs()\n",
    "        o = []\n",
    "        no_more_allowed = False\n",
    "        for coeff in coeffs:\n",
    "            if coeff.get(\"$public\", False) is True:\n",
    "                if no_more_allowed:\n",
    "                    raise Exception(\"Public var declarations must be at the top\")\n",
    "                var_name = [x for x in list(coeff.keys()) if \"$\" not in str(x)][0]\n",
    "                if coeff != {\"$public\": True, \"$output_coeff\": 0, var_name: -1}:\n",
    "                    raise Exception(\"Malformatted coeffs: {}\", format(coeffs))\n",
    "                o.append(var_name)\n",
    "            else:\n",
    "                no_more_allowed = True\n",
    "        return o\n",
    "\n",
    "    # Generate the gate polynomials: L, R, M, O, C,\n",
    "    # each a list of length `group_order`\n",
    "    def make_gate_polynomials(\n",
    "        self,\n",
    "    ) -> tuple[Polynomial, Polynomial, Polynomial, Polynomial, Polynomial]:\n",
    "        L = [Scalar(0) for _ in range(self.group_order)]\n",
    "        R = [Scalar(0) for _ in range(self.group_order)]\n",
    "        M = [Scalar(0) for _ in range(self.group_order)]\n",
    "        O = [Scalar(0) for _ in range(self.group_order)]\n",
    "        C = [Scalar(0) for _ in range(self.group_order)]\n",
    "        for i, constraint in enumerate(self.constraints):\n",
    "            gate = constraint.gate()\n",
    "            L[i] = gate.L\n",
    "            R[i] = gate.R\n",
    "            M[i] = gate.M\n",
    "            O[i] = gate.O\n",
    "            C[i] = gate.C\n",
    "        return (\n",
    "            Polynomial(L, Basis.LAGRANGE),\n",
    "            Polynomial(R, Basis.LAGRANGE),\n",
    "            Polynomial(M, Basis.LAGRANGE),\n",
    "            Polynomial(O, Basis.LAGRANGE),\n",
    "            Polynomial(C, Basis.LAGRANGE),\n",
    "        )\n",
    "\n",
    "    # Attempts to \"run\" the program to fill in any intermediate variable\n",
    "    # assignments, starting from the given assignments. Eg. if\n",
    "    # `starting_assignments` contains {'a': 3, 'b': 5}, and the first line\n",
    "    # says `c <== a * b`, then it fills in `c: 15`.\n",
    "    def fill_variable_assignments(\n",
    "        self, starting_assignments: dict[Optional[str], int]\n",
    "    ) -> dict[Optional[str], int]:\n",
    "        out = {k: Scalar(v) for k, v in starting_assignments.items()}\n",
    "        out[None] = Scalar(0)\n",
    "        for constraint in self.constraints:\n",
    "            wires = constraint.wires\n",
    "            coeffs = constraint.coeffs\n",
    "            in_L = wires.L\n",
    "            in_R = wires.R\n",
    "            output = wires.O\n",
    "            out_coeff = coeffs.get(\"$output_coeff\", 1)\n",
    "            product_key = get_product_key(in_L, in_R)\n",
    "            if output is not None and out_coeff in (-1, 1):\n",
    "                new_value = (\n",
    "                    Scalar(\n",
    "                        coeffs.get(\"\", 0)\n",
    "                        + out[in_L] * coeffs.get(in_L, 0)\n",
    "                        + out[in_R] * coeffs.get(in_R, 0) * (1 if in_R != in_L else 0)\n",
    "                        + out[in_L] * out[in_R] * coeffs.get(product_key, 0)\n",
    "                    )\n",
    "                    * out_coeff\n",
    "                )  # should be / but equivalent for (1, -1)\n",
    "                if output in out:\n",
    "                    if out[output] != new_value:\n",
    "                        raise Exception(\n",
    "                            \"Failed assertion: {} = {}\".format(out[output], new_value)\n",
    "                        )\n",
    "                else:\n",
    "                    out[output] = new_value\n",
    "                    # print('filled in:', output, out[output])\n",
    "        return {k: v.n for k, v in out.items()}\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "419f5bfe-34a5-4683-8500-2b5db96bad8e",
   "metadata": {},
   "outputs": [],
   "source": [
    "program = Program([\"e public\", \"c <== a * b\", \"e <== c * d\"], group_order)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "0f3bcec4-a8b5-43ac-9fdb-55fd6dbed230",
   "metadata": {},
   "outputs": [],
   "source": [
    "assert program.group_order == 8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "f8759dd5-8b60-4645-844c-f6966be4f704",
   "metadata": {},
   "outputs": [],
   "source": [
    "    assignments = {\"a\": 3, \"b\": 4, \"c\": 12, \"d\": 5, \"e\": 60}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "1611f782-6e96-4081-b8d9-f48a0f843f38",
   "metadata": {},
   "outputs": [],
   "source": [
    "from compiler.program import Program, CommonPreprocessedInput\n",
    "from utils import *\n",
    "from setup import *\n",
    "from typing import Optional\n",
    "from dataclasses import dataclass\n",
    "from transcript import Transcript, Message1, Message2, Message3, Message4, Message5\n",
    "from poly import Polynomial, Basis\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Proof:\n",
    "    msg_1: Message1\n",
    "    msg_2: Message2\n",
    "    msg_3: Message3\n",
    "    msg_4: Message4\n",
    "    msg_5: Message5\n",
    "\n",
    "    def flatten(self):\n",
    "        proof = {}\n",
    "        proof[\"a_1\"] = self.msg_1.a_1\n",
    "        proof[\"b_1\"] = self.msg_1.b_1\n",
    "        proof[\"c_1\"] = self.msg_1.c_1\n",
    "        proof[\"z_1\"] = self.msg_2.z_1\n",
    "        proof[\"W_t\"] = self.msg_3.W_t\n",
    "        proof[\"a_eval\"] = self.msg_4.a_eval\n",
    "        proof[\"b_eval\"] = self.msg_4.b_eval\n",
    "        proof[\"c_eval\"] = self.msg_4.c_eval\n",
    "        proof[\"ql_eval\"] = self.msg_4.ql_eval\n",
    "        proof[\"qr_eval\"] = self.msg_4.qr_eval\n",
    "        proof[\"qm_eval\"] = self.msg_4.qm_eval\n",
    "        proof[\"qo_eval\"] = self.msg_4.qo_eval\n",
    "        proof[\"qc_eval\"] = self.msg_4.qc_eval\n",
    "        proof[\"s1_eval\"] = self.msg_4.s1_eval\n",
    "        proof[\"s2_eval\"] = self.msg_4.s2_eval\n",
    "        proof[\"s3_eval\"] = self.msg_4.s3_eval\n",
    "        proof[\"z_eval\"] = self.msg_4.z_eval\n",
    "        proof[\"zw_eval\"] = self.msg_4.zw_eval\n",
    "        proof[\"t_eval\"] = self.msg_4.t_eval\n",
    "        proof[\"W_a\"] = self.msg_5.W_a\n",
    "        proof[\"W_a_quot\"] = self.msg_5.W_a_quot\n",
    "        proof[\"W_b\"] = self.msg_5.W_b\n",
    "        proof[\"W_b_quot\"] = self.msg_5.W_b_quot\n",
    "        proof[\"W_c\"] = self.msg_5.W_c\n",
    "        proof[\"W_c_quot\"] = self.msg_5.W_c_quot\n",
    "        proof[\"W_ql\"] = self.msg_5.W_ql\n",
    "        proof[\"W_ql_quot\"] = self.msg_5.W_ql_quot\n",
    "        proof[\"W_qr\"] = self.msg_5.W_qr\n",
    "        proof[\"W_qr_quot\"] = self.msg_5.W_qr_quot\n",
    "        proof[\"W_qm\"] = self.msg_5.W_qm\n",
    "        proof[\"W_qm_quot\"] = self.msg_5.W_qm_quot\n",
    "        proof[\"W_qo\"] = self.msg_5.W_qo\n",
    "        proof[\"W_qo_quot\"] = self.msg_5.W_qo_quot\n",
    "        proof[\"W_qc\"] = self.msg_5.W_qc\n",
    "        proof[\"W_qc_quot\"] = self.msg_5.W_qc_quot\n",
    "        proof[\"W_s1\"] = self.msg_5.W_s1\n",
    "        proof[\"W_s1_quot\"] = self.msg_5.W_s1_quot\n",
    "        proof[\"W_s2\"] = self.msg_5.W_s2\n",
    "        proof[\"W_s2_quot\"] = self.msg_5.W_s2_quot\n",
    "        proof[\"W_s3\"] = self.msg_5.W_s3\n",
    "        proof[\"W_s3_quot\"] = self.msg_5.W_s3_quot\n",
    "        proof[\"W_z\"] = self.msg_5.W_z\n",
    "        proof[\"W_z_quot\"] = self.msg_5.W_z_quot\n",
    "        proof[\"W_zw\"] = self.msg_5.W_zw\n",
    "        proof[\"W_zw_quot\"] = self.msg_5.W_zw_quot\n",
    "        proof[\"W_t\"] = self.msg_5.W_t\n",
    "        proof[\"W_t_quot\"] = self.msg_5.W_t_quot\n",
    "        return proof\n",
    "\n",
    "\n",
    "@dataclass\n",
    "class Prover:\n",
    "    group_order: int\n",
    "    setup: Setup\n",
    "    program: Program\n",
    "    pk: CommonPreprocessedInput\n",
    "\n",
    "    def __init__(self, setup: Setup, program: Program):\n",
    "        self.group_order = program.group_order\n",
    "        self.setup = setup\n",
    "        self.program = program\n",
    "        self.pk = program.common_preprocessed_input()\n",
    "\n",
    "    def prove(self, witness: dict[Optional[str], int]) -> Proof:\n",
    "        # Initialise Fiat-Shamir transcript\n",
    "        transcript = Transcript(b\"plonk\")\n",
    "\n",
    "        # Collect fixed and public information\n",
    "        # FIXME: Hash pk and PI into transcript\n",
    "        public_vars = self.program.get_public_assignments()\n",
    "        # Public input polynomial\n",
    "        PI = Polynomial(\n",
    "            [Scalar(-witness[v]) for v in public_vars]\n",
    "            + [Scalar(0) for _ in range(self.group_order - len(public_vars))],\n",
    "            Basis.LAGRANGE,\n",
    "        )\n",
    "        self.PI = PI\n",
    "\n",
    "        # Round 1\n",
    "        msg_1 = self.round_1(witness)\n",
    "        self.beta, self.gamma = transcript.round_1(msg_1)\n",
    "\n",
    "        # Round 2\n",
    "        msg_2 = self.round_2()\n",
    "        self.alpha, self.fft_cofactor = transcript.round_2(msg_2)\n",
    "\n",
    "        # Round 3\n",
    "        msg_3 = self.round_3()\n",
    "        self.zeta = transcript.round_3(msg_3)\n",
    "\n",
    "        # Round 4\n",
    "        msg_4 = self.round_4()\n",
    "        self.v = transcript.round_4(msg_4)\n",
    "\n",
    "        # Round 5\n",
    "        msg_5 = self.round_5()\n",
    "\n",
    "        return Proof(msg_1, msg_2, msg_3, msg_4, msg_5)\n",
    "\n",
    "    def round_1(\n",
    "        self,\n",
    "        witness: dict[Optional[str], int],\n",
    "    ) -> Message1:\n",
    "        # https://github.com/sec-bit/learning-zkp/blob/develop/plonk-intro-cn/plonk-arithmetization.md\n",
    "        program = self.program\n",
    "        setup = self.setup\n",
    "        group_order = self.group_order\n",
    "\n",
    "        if None not in witness:\n",
    "            witness[None] = 0\n",
    "\n",
    "        # Compute wire assignments\n",
    "        A_values = [Scalar(0) for _ in range(group_order)]\n",
    "        B_values = [Scalar(0) for _ in range(group_order)]\n",
    "        C_values = [Scalar(0) for _ in range(group_order)]\n",
    "\n",
    "        for i, gate_wires in enumerate(program.wires()):\n",
    "            A_values[i] = Scalar(witness[gate_wires.L])\n",
    "            B_values[i] = Scalar(witness[gate_wires.R])\n",
    "            C_values[i] = Scalar(witness[gate_wires.O])\n",
    "\n",
    "        self.A = Polynomial(A_values, Basis.LAGRANGE)\n",
    "        self.B = Polynomial(B_values, Basis.LAGRANGE)\n",
    "        self.C = Polynomial(C_values, Basis.LAGRANGE)\n",
    "\n",
    "        a_1 = setup.commit(self.A)\n",
    "        b_1 = setup.commit(self.B)\n",
    "        c_1 = setup.commit(self.C)\n",
    "\n",
    "        # Sanity check that witness fulfils gate constraints\n",
    "        assert (\n",
    "            self.A * self.pk.QL\n",
    "            + self.B * self.pk.QR\n",
    "            + self.A * self.B * self.pk.QM\n",
    "            + self.C * self.pk.QO\n",
    "            + self.PI\n",
    "            + self.pk.QC\n",
    "            == Polynomial([Scalar(0)] * group_order, Basis.LAGRANGE)\n",
    "        )\n",
    "\n",
    "        return Message1(a_1, b_1, c_1)\n",
    "\n",
    "    def round_2(self) -> Message2:\n",
    "        # https://github.com/sec-bit/learning-zkp/blob/develop/plonk-intro-cn/plonk-permutation.md#%E5%AE%8C%E6%95%B4%E7%9A%84%E7%BD%AE%E6%8D%A2%E5%8D%8F%E8%AE%AE\n",
    "        group_order = self.group_order\n",
    "        setup = self.setup\n",
    "\n",
    "        Z_values = [Scalar(1)]\n",
    "        roots_of_unity = Scalar.roots_of_unity(group_order)\n",
    "        for i in range(group_order):\n",
    "            Z_values.append(\n",
    "                Z_values[-1]\n",
    "                * self.rlc(self.A.values[i], roots_of_unity[i])\n",
    "                * self.rlc(self.B.values[i], 2 * roots_of_unity[i])\n",
    "                * self.rlc(self.C.values[i], 3 * roots_of_unity[i])\n",
    "                / self.rlc(self.A.values[i], self.pk.S1.values[i])\n",
    "                / self.rlc(self.B.values[i], self.pk.S2.values[i])\n",
    "                / self.rlc(self.C.values[i], self.pk.S3.values[i])\n",
    "            )\n",
    "        # The last value is 1\n",
    "        assert Z_values.pop() == 1\n",
    "\n",
    "        # Sanity-check that Z was computed correctly\n",
    "        for i in range(group_order):\n",
    "            assert (\n",
    "                self.rlc(self.A.values[i], roots_of_unity[i])\n",
    "                * self.rlc(self.B.values[i], 2 * roots_of_unity[i])\n",
    "                * self.rlc(self.C.values[i], 3 * roots_of_unity[i])\n",
    "            ) * Z_values[i] - (\n",
    "                self.rlc(self.A.values[i], self.pk.S1.values[i])\n",
    "                * self.rlc(self.B.values[i], self.pk.S2.values[i])\n",
    "                * self.rlc(self.C.values[i], self.pk.S3.values[i])\n",
    "            ) * Z_values[\n",
    "                (i + 1) % group_order\n",
    "            ] == 0\n",
    "\n",
    "        Z = Polynomial(Z_values, Basis.LAGRANGE)\n",
    "        z_1 = setup.commit(Z)\n",
    "        print(\"Permutation accumulator polynomial successfully generated\")\n",
    "\n",
    "        self.Z = Z\n",
    "        return Message2(z_1)\n",
    "\n",
    "    def round_3(self) -> Message3:\n",
    "        # https://github.com/sec-bit/learning-zkp/blob/develop/plonk-intro-cn/plonk-constraints.md\n",
    "        group_order = self.group_order\n",
    "        setup = self.setup\n",
    "\n",
    "        # Compute the quotient polynomial\n",
    "\n",
    "        alpha = self.alpha\n",
    "\n",
    "        roots_of_unity = Scalar.roots_of_unity(group_order)\n",
    "\n",
    "        A_coeff, B_coeff, C_coeff, S1_coeff, S2_coeff, S3_coeff, Z_coeff, QL_coeff, QR_coeff, QM_coeff, QO_coeff, QC_coeff, PI_coeff = (\n",
    "            x.ifft()\n",
    "            for x in (\n",
    "                self.A,\n",
    "                self.B,\n",
    "                self.C,\n",
    "                self.pk.S1,\n",
    "                self.pk.S2,\n",
    "                self.pk.S3,\n",
    "                self.Z,\n",
    "                self.pk.QL,\n",
    "                self.pk.QR,\n",
    "                self.pk.QM,\n",
    "                self.pk.QO,\n",
    "                self.pk.QC,\n",
    "                self.PI,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        L0_coeff = (\n",
    "            Polynomial([Scalar(1)] + [Scalar(0)] * (group_order - 1), Basis.LAGRANGE)\n",
    "        ).ifft()\n",
    "\n",
    "        # x^8 - 1 coeffs are [-1, 0, 0, 0, 0, 0, 0, 0, 1]\n",
    "        # which needs 9 points(n + 1) to determine the polynomial\n",
    "        ZH_array = [Scalar(-1)] + [Scalar(0)] * (group_order - 1) + [Scalar(1)]\n",
    "        ZH_coeff = Polynomial(ZH_array, Basis.MONOMIAL)\n",
    "\n",
    "        gate_constraints_coeff = (\n",
    "            A_coeff * QL_coeff\n",
    "            + B_coeff * QR_coeff\n",
    "            + A_coeff * B_coeff * QM_coeff\n",
    "            + C_coeff * QO_coeff\n",
    "            + PI_coeff\n",
    "            + QC_coeff\n",
    "        )\n",
    "\n",
    "        normal_roots = Polynomial(\n",
    "            roots_of_unity, Basis.LAGRANGE\n",
    "        )\n",
    "\n",
    "        roots_coeff = normal_roots.ifft()\n",
    "        # z * w\n",
    "        ZW = self.Z.shift(1)\n",
    "        ZW_coeff = ZW.ifft()\n",
    "\n",
    "        for i in range(group_order):\n",
    "            assert (\n",
    "                self.rlc(self.A.values[i], roots_of_unity[i])\n",
    "                * self.rlc(self.B.values[i], 2 * roots_of_unity[i])\n",
    "                * self.rlc(self.C.values[i], 3 * roots_of_unity[i])\n",
    "            ) * self.Z.values[i] - (\n",
    "                self.rlc(self.A.values[i], self.pk.S1.values[i])\n",
    "                * self.rlc(self.B.values[i], self.pk.S2.values[i])\n",
    "                * self.rlc(self.C.values[i], self.pk.S3.values[i])\n",
    "            ) * ZW.values[\n",
    "                i % group_order\n",
    "            ] == 0\n",
    "\n",
    "        permutation_grand_product_coeff = (\n",
    "            (\n",
    "                self.rlc(A_coeff, roots_coeff)\n",
    "                * self.rlc(B_coeff, roots_coeff * Scalar(2))\n",
    "                * self.rlc(C_coeff, roots_coeff * Scalar(3))\n",
    "            )\n",
    "            * Z_coeff\n",
    "            - (\n",
    "                self.rlc(A_coeff, S1_coeff)\n",
    "                * self.rlc(B_coeff, S2_coeff)\n",
    "                * self.rlc(C_coeff, S3_coeff)\n",
    "            )\n",
    "            * ZW_coeff\n",
    "        )\n",
    "\n",
    "        permutation_first_row_coeff = (Z_coeff - Scalar(1)) * L0_coeff\n",
    "\n",
    "        all_constraints = (\n",
    "            gate_constraints_coeff\n",
    "            + permutation_grand_product_coeff * alpha\n",
    "            + permutation_first_row_coeff * alpha**2\n",
    "        )\n",
    "\n",
    "        # quotient polynomial\n",
    "        T_coeff = all_constraints / ZH_coeff\n",
    "\n",
    "        print(\"Generated the quotient polynomial\")\n",
    "\n",
    "        W_t = setup.commit(T_coeff)\n",
    "\n",
    "        self.A_coeff = A_coeff\n",
    "        self.B_coeff = B_coeff\n",
    "        self.C_coeff = C_coeff\n",
    "        self.S1_coeff = S1_coeff\n",
    "        self.S2_coeff = S2_coeff\n",
    "        self.S3_coeff = S3_coeff\n",
    "        self.Z_coeff = Z_coeff\n",
    "        self.ZW_coeff = ZW_coeff\n",
    "        self.QL_coeff = QL_coeff\n",
    "        self.QR_coeff = QR_coeff\n",
    "        self.QM_coeff = QM_coeff\n",
    "        self.QO_coeff = QO_coeff\n",
    "        self.QC_coeff = QC_coeff\n",
    "        self.PI_coeff = PI_coeff\n",
    "        self.T_coeff = T_coeff\n",
    "\n",
    "        return Message3(W_t)\n",
    "\n",
    "    def round_4(self) -> Message4:\n",
    "        # https://github.com/sec-bit/learning-zkp/blob/develop/plonk-intro-cn/plonk-constraints.md\n",
    "        group_order = self.group_order\n",
    "        zeta = self.zeta\n",
    "\n",
    "        a_eval = self.A_coeff.coeff_eval(zeta)\n",
    "        b_eval = self.B_coeff.coeff_eval(zeta)\n",
    "        c_eval = self.C_coeff.coeff_eval(zeta)\n",
    "        s1_eval = self.S1_coeff.coeff_eval(zeta)\n",
    "        s2_eval = self.S2_coeff.coeff_eval(zeta)\n",
    "        s3_eval = self.S3_coeff.coeff_eval(zeta)\n",
    "        root_of_unity = Scalar.root_of_unity(group_order)\n",
    "        z_eval = self.Z_coeff.coeff_eval(zeta)\n",
    "        zw_eval = self.Z_coeff.coeff_eval(zeta * root_of_unity)\n",
    "        ql_eval = self.QL_coeff.coeff_eval(zeta)\n",
    "        qr_eval = self.QR_coeff.coeff_eval(zeta)\n",
    "        qm_eval = self.QM_coeff.coeff_eval(zeta)\n",
    "        qo_eval = self.QO_coeff.coeff_eval(zeta)\n",
    "        qc_eval = self.QC_coeff.coeff_eval(zeta)\n",
    "        t_eval = self.T_coeff.coeff_eval(zeta)\n",
    "\n",
    "        self.a_eval = a_eval\n",
    "        self.b_eval = b_eval\n",
    "        self.c_eval = c_eval\n",
    "        self.ql_eval = ql_eval\n",
    "        self.qr_eval = qr_eval\n",
    "        self.qm_eval = qm_eval\n",
    "        self.qo_eval = qo_eval\n",
    "        self.qc_eval = qc_eval\n",
    "        self.s1_eval = s1_eval\n",
    "        self.s2_eval = s2_eval\n",
    "        self.s3_eval = s3_eval\n",
    "        self.z_eval = z_eval\n",
    "        self.zw_eval = zw_eval\n",
    "        self.t_eval = t_eval\n",
    "\n",
    "        return Message4(\n",
    "            a_eval,\n",
    "            b_eval,\n",
    "            c_eval,\n",
    "            ql_eval,\n",
    "            qr_eval,\n",
    "            qm_eval,\n",
    "            qo_eval,\n",
    "            qc_eval,\n",
    "            s1_eval,\n",
    "            s2_eval,\n",
    "            s3_eval,\n",
    "            z_eval,\n",
    "            zw_eval,\n",
    "            t_eval\n",
    "        )\n",
    "\n",
    "    def round_5(self) -> Message5:\n",
    "        W_a, W_a_quot = self.generate_commitment(self.A_coeff, self.a_eval)\n",
    "        W_b, W_b_quot = self.generate_commitment(self.B_coeff, self.b_eval)\n",
    "        W_c, W_c_quot = self.generate_commitment(self.C_coeff, self.c_eval)\n",
    "        W_ql, W_ql_quot = self.generate_commitment(self.QL_coeff, self.ql_eval)\n",
    "        W_qr, W_qr_quot = self.generate_commitment(self.QR_coeff, self.qr_eval)\n",
    "        W_qm, W_qm_quot = self.generate_commitment(self.QM_coeff, self.qm_eval)\n",
    "        W_qo, W_qo_quot = self.generate_commitment(self.QO_coeff, self.qo_eval)\n",
    "        W_qc, W_qc_quot = self.generate_commitment(self.QC_coeff, self.qc_eval)\n",
    "        W_s1, W_s1_quot = self.generate_commitment(self.S1_coeff, self.s1_eval)\n",
    "        W_s2, W_s2_quot = self.generate_commitment(self.S2_coeff, self.s2_eval)\n",
    "        W_s3, W_s3_quot = self.generate_commitment(self.S3_coeff, self.s3_eval)\n",
    "        W_z, W_z_quot = self.generate_commitment(self.Z_coeff, self.z_eval)\n",
    "        W_zw, W_zw_quot = self.generate_commitment(self.ZW_coeff, self.zw_eval)\n",
    "        W_t, W_t_quot = self.generate_commitment(self.T_coeff, self.t_eval)\n",
    "\n",
    "        print(\"Generated final quotient witness polynomials\")\n",
    "        return Message5(\n",
    "            W_a, W_a_quot,\n",
    "            W_b, W_b_quot,\n",
    "            W_c, W_c_quot,\n",
    "            W_ql, W_ql_quot,\n",
    "            W_qr, W_qr_quot,\n",
    "            W_qm, W_qm_quot,\n",
    "            W_qo, W_qo_quot,\n",
    "            W_qc, W_qc_quot,\n",
    "            W_s1, W_s1_quot,\n",
    "            W_s2, W_s2_quot,\n",
    "            W_s3, W_s3_quot,\n",
    "            W_z, W_z_quot,\n",
    "            W_zw, W_zw_quot,\n",
    "            W_t, W_t_quot,\n",
    "        )\n",
    "\n",
    "    def rlc(self, term_1, term_2):\n",
    "        return term_1 + term_2 * self.beta + self.gamma\n",
    "\n",
    "    def generate_commitment(self, coeff: Polynomial, eval: Scalar):\n",
    "        setup = self.setup\n",
    "        zeta = self.zeta\n",
    "        # Polynomial for (X - zeta)\n",
    "        ZH_zeta_coeff = Polynomial([-zeta, Scalar(1)], Basis.MONOMIAL)\n",
    "        quot_coeff = (coeff - eval) / ZH_zeta_coeff\n",
    "        # witness for polynomial itself\n",
    "        w = setup.commit(coeff)\n",
    "        # witness for quotient polynomial\n",
    "        w_quot = setup.commit(quot_coeff)\n",
    "        return w, w_quot\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "781f63c9-f04b-42a2-a418-b462f47a4f96",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "__main__.Proof"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Proof"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "c459d635-f246-423e-bfc0-c901f83537cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "prover = Prover(setup, program)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "af18bdc6-2330-4ca2-8157-aa2da71e2065",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Permutation accumulator polynomial successfully generated\n",
      "Generated the quotient polynomial\n",
      "Generated final quotient witness polynomials\n"
     ]
    }
   ],
   "source": [
    "proof = prover.prove(assignments)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f072c329-2072-4845-9757-de8a5c585265",
   "metadata": {},
   "source": [
    "Verify"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "cba02134-5068-4e62-8d93-b402e84ee4c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "    program = Program([\"e public\", \"c <== a * b\", \"e <== c * d\"], group_order)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "0bd8e66a-4587-4628-88c1-410d4ed62c37",
   "metadata": {},
   "outputs": [],
   "source": [
    "    vk = setup.verification_key(program.common_preprocessed_input())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "1d39b418-998c-4b81-b14c-dc19fc749eff",
   "metadata": {},
   "outputs": [],
   "source": [
    "    public = [60]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "57d64172-94c1-4bf8-86f7-f753c5a015a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Done KZG10 commitment check for a_eval polynomial\n",
      "Done KZG10 commitment check for b_eval polynomial\n",
      "Done KZG10 commitment check for c_eval polynomial\n",
      "Done KZG10 commitment check for z_eval polynomial\n",
      "Done KZG10 commitment check for zw_eval polynomial\n",
      "Done KZG10 commitment check for t_eval polynomial\n",
      "Done KZG10 commitment check for ql_eval polynomial\n",
      "Done KZG10 commitment check for qr_eval polynomial\n",
      "Done KZG10 commitment check for qm_eval polynomial\n",
      "Done KZG10 commitment check for qo_eval polynomial\n",
      "Done KZG10 commitment check for qc_eval polynomial\n",
      "Done KZG10 commitment check for s1_eval polynomial\n",
      "Done KZG10 commitment check for s2_eval polynomial\n",
      "Done KZG10 commitment check for s3_eval polynomial\n",
      "Done equation check for all constraints\n"
     ]
    }
   ],
   "source": [
    "    assert vk.verify_proof(group_order, proof, public)\n"
   ]
  }
 ],
 "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}