[
  {
    "path": ".gitignore",
    "content": "venv\n*.pyc\n# Remove jupyter notebook stuff\n*.ipynb\nsettings.json"
  },
  {
    "path": "README.md",
    "content": "# PlonKathon\n**PlonKathon** is part of the program for [MIT IAP 2023] [Modern Zero Knowledge Cryptography](https://zkiap.com/). Over the course of this weekend, we will get into the weeds of the PlonK protocol through a series of exercises and extensions. This repository contains a simple python implementation of PlonK adapted from [py_plonk](https://github.com/ethereum/research/tree/master/py_plonk), and targeted to be close to compatible with the implementation at https://zkrepl.dev.\n\n### Exercises\n\nEach step of the exercise is accompanied by tests in `test.py` to check your progress.\n\n#### Step 1: Implement setup.py\n\nImplement `Setup.commit` and `Setup.verification_key`.\n\n#### Step 2: Implement prover.py\n\n1. Implement Round 1 of the PlonK prover\n2. Implement Round 2 of the PlonK prover\n3. Implement Round 3 of the PlonK prover\n4. Implement Round 4 of the PlonK prover\n5. Implement Round 5 of the PlonK prover\n\n#### Step 3: Implement verifier.py\n\nImplement `VerificationKey.verify_proof_unoptimized` and `VerificationKey.verify_proof`. See the comments for the differences.\n\n#### Step 4: Pass all the tests!\n\nPass a number of miscellaneous tests that test your implementation end-to-end.\n\n### Extensions\n1. Add support for custom gates.\n[TurboPlonK](https://docs.zkproof.org/pages/standards/accepted-workshop3/proposal-turbo_plonk.pdf) introduced support for custom constraints, beyond the addition and multiplication gates supported here. Try to generalise this implementation to allow circuit writers to define custom constraints.\n2. Add zero-knowledge.\nThe parts of PlonK that are responsible for ensuring strong privacy are left out of this implementation. See if you can identify them in the [original paper](https://eprint.iacr.org/2019/953.pdf) and add them here.\n3. Add support for lookups.\nA lookup argument allows us to prove that a certain element can be found in a public lookup table. [PlonKup](https://eprint.iacr.org/2022/086.pdf) introduces lookup arguments to PlonK. Try to understand the construction in the paper and implement it here.\n4. Implement Merlin transcript.\nCurrently, this implementation uses the [merlin transcript package](https://github.com/nalinbhardwaj/curdleproofs.pie/tree/master/merlin). Learn about the [Merlin transcript construction](https://merlin.cool) and the [STROBE framework](https://www.cryptologie.net/article/416/the-strobe-protocol-framework/) which Merlin is based upon, and then implement the transcript class `MerlinTranscript` yourself!\n\n## Getting started\n\nTo get started, you'll need to have a Python version >= 3.8 and [`poetry`](https://python-poetry.org) installed: `curl -sSL https://install.python-poetry.org | python3 -`.\n\nThen, run `poetry install` in the root of the repository. This will install all the dependencies in a virtualenv.\n\nThen, to see the proof system in action, run `poetry run python test.py` from the root of the repository. This will take you through the workflow of setup, proof generation, and verification for several example programs.\n\nThe `main` branch contains code stubbed out with comments to guide you through the tests. The `hardcore` branch removes the comments for the more adventurous amongst you. The `reference` branch contains a completed implementation.\n\nFor linting and types, the repo also provides `poetry run black .` and `poetry run mypy .` \n\n### Compiler\n#### Program\nWe specify our program logic in a high-level language involving constraints and variable assignments. Here is a program that lets you prove that you know two small numbers that multiply to a given number (in our example we'll use 91) without revealing what those numbers are:\n\n```\nn public\npb0 === pb0 * pb0\npb1 === pb1 * pb1\npb2 === pb2 * pb2\npb3 === pb3 * pb3\nqb0 === qb0 * qb0\nqb1 === qb1 * qb1\nqb2 === qb2 * qb2\nqb3 === qb3 * qb3\npb01 <== pb0 + 2 * pb1\npb012 <== pb01 + 4 * pb2\np <== pb012 + 8 * pb3\nqb01 <== qb0 + 2 * qb1\nqb012 <== qb01 + 4 * qb2\nq <== qb012 + 8 * qb3\nn <== p * q\n```\n\nExamples of valid program constraints:\n- `a === 9`\n- `b <== a * c`\n- `d <== a * c - 45 * a + 987`\n\nExamples of invalid program constraints:\n- `7 === 7` (can't assign to non-variable)\n- `a <== b * * c` (two multiplications in a row)\n- `e <== a + b * c * d` (multiplicative degree > 2)\n\nGiven a `Program`, we can derive the `CommonPreprocessedInput`, which are the polynomials representing the fixed constraints of the program. The prover later uses these polynomials to construct the quotient polynomial, and to compute their evaluations at a given challenge point.\n\n```python\n@dataclass\nclass CommonPreprocessedInput:\n    \"\"\"Common preprocessed input\"\"\"\n\n    group_order: int\n    # q_M(X) multiplication selector polynomial\n    QM: list[Scalar]\n    # q_L(X) left selector polynomial\n    QL: list[Scalar]\n    # q_R(X) right selector polynomial\n    QR: list[Scalar]\n    # q_O(X) output selector polynomial\n    QO: list[Scalar]\n    # q_C(X) constants selector polynomial\n    QC: list[Scalar]\n    # S_σ1(X) first permutation polynomial S_σ1(X)\n    S1: list[Scalar]\n    # S_σ2(X) second permutation polynomial S_σ2(X)\n    S2: list[Scalar]\n    # S_σ3(X) third permutation polynomial S_σ3(X)\n    S3: list[Scalar]\n```\n\n#### Assembly\nOur \"assembly\" language consists of `AssemblyEqn`s:\n\n```python\nclass AssemblyEqn:\n    \"\"\"Assembly equation mapping wires to coefficients.\"\"\"\n    wires: GateWires\n    coeffs: dict[Optional[str], int]\n```\n\nwhere:\n```python\n@dataclass\nclass GateWires:\n    \"\"\"Variable names for Left, Right, and Output wires.\"\"\"\n    L: Optional[str]\n    R: Optional[str]\n    O: Optional[str]\n```\n\nExamples of valid program constraints, and corresponding assembly:\n| program constraint         | assembly                                         |\n| -------------------------- | ------------------------------------------------ |\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\n### Setup\nLet $\\mathbb{G}_1$ and $\\mathbb{G}_2$ be two elliptic curves with a pairing $e : \\mathbb{G}_1 \\times \\mathbb{G}_2 \\rightarrow \\mathbb{G}_T$. Let $p$ be the order of $\\mathbb{G}_1$ and $\\mathbb{G}_2$, and $G$ and $H$ be generators of $\\mathbb{G}_1$ and $\\mathbb{G}_2$. We will use the shorthand notation\n\n$$[x]_1 = xG \\in \\mathbb{G}_1 \\text{ and } [x]_2 = xH \\in \\mathbb{G}_2$$\n\nfor any $x \\in \\mathbb{F}_p$.\n\nThe trusted setup is a preprocessing step that produces a structured reference string:\n$$\\mathsf{srs} = ([1]_1, [x]_1, \\cdots, [x^{d-1}]_1, [x]_2),$$\nwhere:\n- $x \\in \\mathbb{F}$ is a randomly chosen, **secret** evaluation point; and\n- $d$ is the size of the trusted setup, corresponding to the maximum degree polynomial that it can support.\n\n```python\n@dataclass\nclass Setup(object):\n    #   ([1]₁, [x]₁, ..., [x^{d-1}]₁)\n    # = ( G,    xG,  ...,  x^{d-1}G ), where G is a generator of G_2\n    powers_of_x: list[G1Point]\n    # [x]₂ = xH, where H is a generator of G_2\n    X2: G2Point\n```\n\nIn this repository, we are using the pairing-friendly [BN254 curve](https://hackmd.io/@jpw/bn254), where:\n- `p = 21888242871839275222246405745257275088696311157297823662689037894645226208583`\n- $\\mathbb{G}_1$ is the curve $y^2 = x^3 + 3$ over $\\mathbb{F}_p$;\n- $\\mathbb{G}_2$ is the twisted curve $y^2 = x^3 + 3/(9+u)$ over $\\mathbb{F}_{p^2}$; and\n- $\\mathbb{G}_T = {\\mu}_r \\subset \\mathbb{F}_{p^{12}}^{\\times}$.\n\nWe are using an existing setup for $d = 2^{11}$, from this [ceremony](https://github.com/iden3/snarkjs/blob/master/README.md). You can find out more about trusted setup ceremonies [here](https://github.com/weijiekoh/perpetualpowersoftau).\n\n### Prover\nThe prover creates a proof of knowledge of some satisfying witness to a program.\n\n```python\n@dataclass\nclass Prover:\n    group_order: int\n    setup: Setup\n    program: Program\n    pk: CommonPreprocessedInput\n```\n\nThe prover progresses in five rounds, and produces a message at the end of each. After each round, the message is hashed into the `Transcript`.\n\nThe `Proof` consists of all the round messages (`Message1`, `Message2`, `Message3`, `Message4`, `Message5`).\n\n#### Round 1\n```python\ndef round_1(\n    self,\n    witness: dict[Optional[str], int],\n) -> Message1\n\n@dataclass\nclass Message1:\n    # - [a(x)]₁ (commitment to left wire polynomial)\n    a_1: G1Point\n    # - [b(x)]₁ (commitment to right wire polynomial)\n    b_1: G1Point\n    # - [c(x)]₁ (commitment to output wire polynomial)\n    c_1: G1Point\n```\n\n#### Round 2\n```python\ndef round_2(self) -> Message2\n\n@dataclass\nclass Message2:\n    # [z(x)]₁ (commitment to permutation polynomial)\n    z_1: G1Point\n```\n\n#### Round 3\n```python\ndef round_3(self) -> Message3\n\n@dataclass\nclass Message3:\n    # [t_lo(x)]₁ (commitment to t_lo(X), the low chunk of the quotient polynomial t(X))\n    t_lo_1: G1Point\n    # [t_mid(x)]₁ (commitment to t_mid(X), the middle chunk of the quotient polynomial t(X))\n    t_mid_1: G1Point\n    # [t_hi(x)]₁ (commitment to t_hi(X), the high chunk of the quotient polynomial t(X))\n    t_hi_1: G1Point\n```\n\n#### Round 4\n```python\ndef round_4(self) -> Message4\n\n@dataclass\nclass Message4:\n    # Evaluation of a(X) at evaluation challenge ζ\n    a_eval: Scalar\n    # Evaluation of b(X) at evaluation challenge ζ\n    b_eval: Scalar\n    # Evaluation of c(X) at evaluation challenge ζ\n    c_eval: Scalar\n    # Evaluation of the first permutation polynomial S_σ1(X) at evaluation challenge ζ\n    s1_eval: Scalar\n    # Evaluation of the second permutation polynomial S_σ2(X) at evaluation challenge ζ\n    s2_eval: Scalar\n    # Evaluation of the shifted permutation polynomial z(X) at the shifted evaluation challenge ζω\n    z_shifted_eval: Scalar\n```\n\n#### Round 5\n```python\ndef round_5(self) -> Message5\n\n@dataclass\nclass Message5:\n    # [W_ζ(X)]₁ (commitment to the opening proof polynomial)\n    W_z_1: G1Point\n    # [W_ζω(X)]₁ (commitment to the opening proof polynomial)\n    W_zw_1: G1Point\n```\n\n### Verifier\nGiven a `Setup` and a `Program`, we can generate a verification key for the program:\n\n```python\ndef verification_key(self, pk: CommonPreprocessedInput) -> VerificationKey\n```\n\nThe `VerificationKey` contains:\n\n| verification key element | remark                                                           |\n| ------------------------ | ---------------------------------------------------------------- |\n| $[q_M(x)]_1$             | commitment to multiplication selector polynomial                 |\n| $[q_L(x)]_1$             | commitment to left selector polynomial                           |\n| $[q_R(x)]_1$             | commitment to right selector polynomial                          |\n| $[q_O(x)]_1$             | commitment to output selector polynomial                         |\n| $[q_C(x)]_1$             | commitment to constants selector polynomial                      |\n| $[S_{\\sigma1}(x)]_1$     | commitment to the first permutation polynomial $S_{\\sigma1}(X)$  |\n| $[S_{\\sigma2}(x)]_1$     | commitment to the second permutation polynomial $S_{\\sigma2}(X)$ |\n| $[S_{\\sigma3}(x)]_1$     | commitment to the third permutation polynomial $S_{\\sigma3}(X)$  |\n| $[x]_2 = xH$             | (from the $\\mathsf{srs}$)                                        |\n| $\\omega$                 | an $n$-th root of unity, where $n$ is the program's group order. |\n"
  },
  {
    "path": "TESTING_verifier_DO_NOT_OPEN.py",
    "content": "import py_ecc.bn128 as b\nfrom utils import *\nfrom dataclasses import dataclass\nfrom curve import *\nfrom transcript import Transcript\nfrom poly import Polynomial, Basis\n\n\n@dataclass\nclass TestingVerificationKey:\n    \"\"\"Testing Verification key: DO NOT READ THIS CODE, only for testing prover implementations\"\"\"\n\n    group_order: int\n    # [q_M(x)]₁ (commitment to multiplication selector polynomial)\n    Qm: G1Point\n    # [q_L(x)]₁ (commitment to left selector polynomial)\n    Ql: G1Point\n    # [q_R(x)]₁ (commitment to right selector polynomial)\n    Qr: G1Point\n    # [q_O(x)]₁ (commitment to output selector polynomial)\n    Qo: G1Point\n    # [q_C(x)]₁ (commitment to constants selector polynomial)\n    Qc: G1Point\n    # [S_σ1(x)]₁ (commitment to the first permutation polynomial S_σ1(X))\n    S1: G1Point\n    # [S_σ2(x)]₁ (commitment to the second permutation polynomial S_σ2(X))\n    S2: G1Point\n    # [S_σ3(x)]₁ (commitment to the third permutation polynomial S_σ3(X))\n    S3: G1Point\n    # [x]₂ = xH, where H is a generator of G_2\n    X_2: G2Point\n    # nth root of unity, where n is the program's group order.\n    w: Scalar\n\n    # More optimized version that tries hard to minimize pairings and\n    # elliptic curve multiplications, but at the cost of being harder\n    # to understand and mixing together a lot of the computations to\n    # efficiently batch them\n    def verify_proof(self, group_order: int, pf, public=[]) -> bool:\n        # 4. Compute challenges\n        beta, gamma, alpha, zeta, v, u = self.compute_challenges(pf)\n        proof = pf.flatten()\n\n        # 5. Compute zero polynomial evaluation Z_H(ζ) = ζ^n - 1\n        root_of_unity = Scalar.root_of_unity(group_order)\n        ZH_ev = zeta**group_order - 1\n\n        # 6. Compute Lagrange polynomial evaluation L_0(ζ)\n        L0_ev = ZH_ev / (group_order * (zeta - 1))\n\n        # 7. Compute public input polynomial evaluation PI(ζ).\n        PI = Polynomial(\n            [Scalar(-x) for x in public]\n            + [Scalar(0) for _ in range(group_order - len(public))],\n            Basis.LAGRANGE,\n        )\n        PI_ev = PI.barycentric_eval(zeta)\n\n        # Compute the constant term of R. This is not literally the degree-0\n        # term of the R polynomial; rather, it's the portion of R that can\n        # be computed directly, without resorting to elliptic cutve commitments\n        r0 = (\n            PI_ev\n            - L0_ev * alpha**2\n            - (\n                alpha\n                * (proof[\"a_eval\"] + beta * proof[\"s1_eval\"] + gamma)\n                * (proof[\"b_eval\"] + beta * proof[\"s2_eval\"] + gamma)\n                * (proof[\"c_eval\"] + gamma)\n                * proof[\"z_shifted_eval\"]\n            )\n        )\n\n        # D = (R - r0) + u * Z\n        D_pt = ec_lincomb(\n            [\n                (self.Qm, proof[\"a_eval\"] * proof[\"b_eval\"]),\n                (self.Ql, proof[\"a_eval\"]),\n                (self.Qr, proof[\"b_eval\"]),\n                (self.Qo, proof[\"c_eval\"]),\n                (self.Qc, 1),\n                (\n                    proof[\"z_1\"],\n                    (\n                        (proof[\"a_eval\"] + beta * zeta + gamma)\n                        * (proof[\"b_eval\"] + beta * 2 * zeta + gamma)\n                        * (proof[\"c_eval\"] + beta * 3 * zeta + gamma)\n                        * alpha\n                        + L0_ev * alpha**2\n                        + u\n                    ),\n                ),\n                (\n                    self.S3,\n                    (\n                        -(proof[\"a_eval\"] + beta * proof[\"s1_eval\"] + gamma)\n                        * (proof[\"b_eval\"] + beta * proof[\"s2_eval\"] + gamma)\n                        * alpha\n                        * beta\n                        * proof[\"z_shifted_eval\"]\n                    ),\n                ),\n                (proof[\"t_lo_1\"], -ZH_ev),\n                (proof[\"t_mid_1\"], -ZH_ev * zeta**group_order),\n                (proof[\"t_hi_1\"], -ZH_ev * zeta ** (group_order * 2)),\n            ]\n        )\n\n        F_pt = ec_lincomb(\n            [\n                (D_pt, 1),\n                (proof[\"a_1\"], v),\n                (proof[\"b_1\"], v**2),\n                (proof[\"c_1\"], v**3),\n                (self.S1, v**4),\n                (self.S2, v**5),\n            ]\n        )\n\n        E_pt = ec_mul(\n            b.G1,\n            (\n                -r0\n                + v * proof[\"a_eval\"]\n                + v**2 * proof[\"b_eval\"]\n                + v**3 * proof[\"c_eval\"]\n                + v**4 * proof[\"s1_eval\"]\n                + v**5 * proof[\"s2_eval\"]\n                + u * proof[\"z_shifted_eval\"]\n            ),\n        )\n\n        # What's going on here is a clever re-arrangement of terms to check\n        # the same equations that are being checked in the basic version,\n        # but in a way that minimizes the number of EC muls and even\n        # compressed the two pairings into one. The 2 pairings -> 1 pairing\n        # trick is basically to replace checking\n        #\n        # Y1 = A * (X - a) and Y2 = B * (X - b)\n        #\n        # with\n        #\n        # Y1 + A * a = A * X\n        # Y2 + B * b = B * X\n        #\n        # so at this point we can take a random linear combination of the two\n        # checks, and verify it with only one pairing.\n        assert b.pairing(\n            self.X_2, ec_lincomb([(proof[\"W_z_1\"], 1), (proof[\"W_zw_1\"], u)])\n        ) == b.pairing(\n            b.G2,\n            ec_lincomb(\n                [\n                    (proof[\"W_z_1\"], zeta),\n                    (proof[\"W_zw_1\"], u * zeta * root_of_unity),\n                    (F_pt, 1),\n                    (E_pt, -1),\n                ]\n            ),\n        )\n\n        print(\"done combined check\")\n        return True\n\n    # Basic, easier-to-understand version of what's going on\n    def verify_proof_unoptimized(self, group_order: int, pf, public=[]) -> bool:\n        # 4. Compute challenges\n        beta, gamma, alpha, zeta, v, _ = self.compute_challenges(pf)\n        proof = pf.flatten()\n\n        # 5. Compute zero polynomial evaluation Z_H(ζ) = ζ^n - 1\n        root_of_unity = Scalar.root_of_unity(group_order)\n        ZH_ev = zeta**group_order - 1\n\n        # 6. Compute Lagrange polynomial evaluation L_0(ζ)\n        L0_ev = ZH_ev / (group_order * (zeta - 1))\n\n        # 7. Compute public input polynomial evaluation PI(ζ).\n        PI = Polynomial(\n            [Scalar(-x) for x in public]\n            + [Scalar(0) for _ in range(group_order - len(public))],\n            Basis.LAGRANGE,\n        )\n        PI_ev = PI.barycentric_eval(zeta)\n\n        # Recover the commitment to the linearization polynomial R,\n        # exactly the same as what was created by the prover\n        R_pt = ec_lincomb(\n            [\n                (self.Qm, proof[\"a_eval\"] * proof[\"b_eval\"]),\n                (self.Ql, proof[\"a_eval\"]),\n                (self.Qr, proof[\"b_eval\"]),\n                (self.Qo, proof[\"c_eval\"]),\n                (b.G1, PI_ev),\n                (self.Qc, 1),\n                (\n                    proof[\"z_1\"],\n                    (\n                        (proof[\"a_eval\"] + beta * zeta + gamma)\n                        * (proof[\"b_eval\"] + beta * 2 * zeta + gamma)\n                        * (proof[\"c_eval\"] + beta * 3 * zeta + gamma)\n                        * alpha\n                    ),\n                ),\n                (\n                    self.S3,\n                    (\n                        -(proof[\"a_eval\"] + beta * proof[\"s1_eval\"] + gamma)\n                        * (proof[\"b_eval\"] + beta * proof[\"s2_eval\"] + gamma)\n                        * beta\n                        * alpha\n                        * proof[\"z_shifted_eval\"]\n                    ),\n                ),\n                (\n                    b.G1,\n                    (\n                        -(proof[\"a_eval\"] + beta * proof[\"s1_eval\"] + gamma)\n                        * (proof[\"b_eval\"] + beta * proof[\"s2_eval\"] + gamma)\n                        * (proof[\"c_eval\"] + gamma)\n                        * alpha\n                        * proof[\"z_shifted_eval\"]\n                    ),\n                ),\n                (proof[\"z_1\"], L0_ev * alpha**2),\n                (b.G1, -L0_ev * alpha**2),\n                (proof[\"t_lo_1\"], -ZH_ev),\n                (proof[\"t_mid_1\"], -ZH_ev * zeta**group_order),\n                (proof[\"t_hi_1\"], -ZH_ev * zeta ** (group_order * 2)),\n            ]\n        )\n\n        print(\"verifier R_pt\", R_pt)\n\n        # Verify that R(z) = 0 and the prover-provided evaluations\n        # A(z), B(z), C(z), S1(z), S2(z) are all correct\n        assert b.pairing(\n            b.G2,\n            ec_lincomb(\n                [\n                    (R_pt, 1),\n                    (proof[\"a_1\"], v),\n                    (b.G1, -v * proof[\"a_eval\"]),\n                    (proof[\"b_1\"], v**2),\n                    (b.G1, -(v**2) * proof[\"b_eval\"]),\n                    (proof[\"c_1\"], v**3),\n                    (b.G1, -(v**3) * proof[\"c_eval\"]),\n                    (self.S1, v**4),\n                    (b.G1, -(v**4) * proof[\"s1_eval\"]),\n                    (self.S2, v**5),\n                    (b.G1, -(v**5) * proof[\"s2_eval\"]),\n                ]\n            ),\n        ) == b.pairing(b.add(self.X_2, ec_mul(b.G2, -zeta)), proof[\"W_z_1\"])\n        print(\"done check 1\")\n\n        # Verify that the provided value of Z(zeta*w) is correct\n        assert b.pairing(\n            b.G2, ec_lincomb([(proof[\"z_1\"], 1), (b.G1, -proof[\"z_shifted_eval\"])])\n        ) == b.pairing(\n            b.add(self.X_2, ec_mul(b.G2, -zeta * root_of_unity)), proof[\"W_zw_1\"]\n        )\n        print(\"done check 2\")\n        return True\n\n    # Compute challenges (should be same as those computed by prover)\n    def compute_challenges(\n        self, proof\n    ) -> tuple[Scalar, Scalar, Scalar, Scalar, Scalar, Scalar]:\n        transcript = Transcript(b\"plonk\")\n        beta, gamma = transcript.round_1(proof.msg_1)\n        alpha, _fft_cofactor = transcript.round_2(proof.msg_2)\n        zeta = transcript.round_3(proof.msg_3)\n        v = transcript.round_4(proof.msg_4)\n        u = transcript.round_5(proof.msg_5)\n\n        return beta, gamma, alpha, zeta, v, u\n"
  },
  {
    "path": "__init__.py",
    "content": ""
  },
  {
    "path": "compiler/__init__.py",
    "content": ""
  },
  {
    "path": "compiler/assembly.py",
    "content": "from utils import *\nfrom .utils import *\nfrom typing import Optional\nfrom dataclasses import dataclass\n\n\n@dataclass\nclass 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\nclass Gate:\n    \"\"\"Gate polynomial\"\"\"\n\n    L: Scalar\n    R: Scalar\n    M: Scalar\n    O: Scalar\n    C: Scalar\n\n\n@dataclass\nclass 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#\ndef 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#\ndef 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"
  },
  {
    "path": "compiler/program.py",
    "content": "# A simple zk language, reverse-engineered to match https://zkrepl.dev/ output\n\nfrom utils import *\nfrom .assembly import *\nfrom .utils import *\nfrom typing import Optional, Set\nfrom poly import Polynomial, Basis\n\n\n@dataclass\nclass 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\nclass 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"
  },
  {
    "path": "compiler/utils.py",
    "content": "from utils import *\nfrom enum import Enum\nfrom dataclasses import dataclass\n\n\nclass 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\nclass 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\ndef 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\ndef is_valid_variable_name(name: str) -> bool:\n    return len(name) > 0 and name.isalnum() and name[0] not in \"0123456789\"\n"
  },
  {
    "path": "curve.py",
    "content": "from py_ecc.fields.field_elements import FQ as Field\nimport py_ecc.bn128 as b\nfrom typing import NewType\n\nprimitive_root = 5\nG1Point = NewType(\"G1Point\", tuple[b.FQ, b.FQ])\nG2Point = NewType(\"G2Point\", tuple[b.FQ2, b.FQ2])\n\n\nclass Scalar(Field):\n    field_modulus = b.curve_order\n\n    # Gets the first root of unity of a given group order\n    @classmethod\n    def root_of_unity(cls, group_order: int):\n        return Scalar(5) ** ((cls.field_modulus - 1) // group_order)\n\n    # Gets the full list of roots of unity of a given group order\n    @classmethod\n    def roots_of_unity(cls, group_order: int):\n        o = [Scalar(1), cls.root_of_unity(group_order)]\n        while len(o) < group_order:\n            o.append(o[-1] * o[1])\n        return o\n\n\nBase = NewType(\"Base\", b.FQ)\n\n\ndef ec_mul(pt, coeff):\n    if hasattr(coeff, \"n\"):\n        coeff = coeff.n\n    return b.multiply(pt, coeff % b.curve_order)\n\n\n# Elliptic curve linear combination. A truly optimized implementation\n# would replace this with a fast lin-comb algo, see https://ethresear.ch/t/7238\ndef ec_lincomb(pairs):\n    return lincomb(\n        [pt for (pt, _) in pairs],\n        [int(n) % b.curve_order for (_, n) in pairs],\n        b.add,\n        b.Z1,\n    )\n    # Equivalent to:\n    # o = b.Z1\n    # for pt, coeff in pairs:\n    #     o = b.add(o, ec_mul(pt, coeff))\n    # return o\n\n\n################################################################\n# multicombs\n################################################################\n\nimport random, sys, math\n\n\ndef multisubset(numbers, subsets, adder=lambda x, y: x + y, zero=0):\n    # Split up the numbers into partitions\n    partition_size = 1 + int(math.log(len(subsets) + 1))\n    # Align number count to partition size (for simplicity)\n    numbers = numbers[::]\n    while len(numbers) % partition_size != 0:\n        numbers.append(zero)\n    # Compute power set for each partition (eg. a, b, c -> {0, a, b, a+b, c, a+c, b+c, a+b+c})\n    power_sets = []\n    for i in range(0, len(numbers), partition_size):\n        new_power_set = [zero]\n        for dimension, value in enumerate(numbers[i : i + partition_size]):\n            new_power_set += [adder(n, value) for n in new_power_set]\n        power_sets.append(new_power_set)\n    # Compute subset sums, using elements from power set for each range of values\n    # ie. with a single power set lookup you can get the sum of _all_ elements in\n    # the range partition_size*k...partition_size*(k+1) that are in that subset\n    subset_sums = []\n    for subset in subsets:\n        o = zero\n        for i in range(len(power_sets)):\n            index_in_power_set = 0\n            for j in range(partition_size):\n                if i * partition_size + j in subset:\n                    index_in_power_set += 2**j\n            o = adder(o, power_sets[i][index_in_power_set])\n        subset_sums.append(o)\n    return subset_sums\n\n\n# Reduces a linear combination `numbers[0] * factors[0] + numbers[1] * factors[1] + ...`\n# into a multi-subset problem, and computes the result efficiently\ndef lincomb(numbers, factors, adder=lambda x, y: x + y, zero=0):\n    # Maximum bit length of a number; how many subsets we need to make\n    maxbitlen = max(len(bin(f)) - 2 for f in factors)\n    # Compute the subsets: the ith subset contains the numbers whose corresponding factor\n    # has a 1 at the ith bit\n    subsets = [\n        {i for i in range(len(numbers)) if factors[i] & (1 << j)}\n        for j in range(maxbitlen + 1)\n    ]\n    subset_sums = multisubset(numbers, subsets, adder=adder, zero=zero)\n    # For example, suppose a value V has factor 6 (011 in increasing-order binary). Subset 0\n    # will not have V, subset 1 will, and subset 2 will. So if we multiply the output of adding\n    # subset 0 with twice the output of adding subset 1, with four times the output of adding\n    # subset 2, then V will be represented 0 + 2 + 4 = 6 times. This reasoning applies for every\n    # value. So `subset_0_sum + 2 * subset_1_sum + 4 * subset_2_sum` gives us the result we want.\n    # Here, we compute this as `((subset_2_sum * 2) + subset_1_sum) * 2 + subset_0_sum` for\n    # efficiency: an extra `maxbitlen * 2` group operations.\n    o = zero\n    for i in range(len(subsets) - 1, -1, -1):\n        o = adder(adder(o, o), subset_sums[i])\n    return o\n\n\n# Tests go here\ndef make_mock_adder():\n    counter = [0]\n\n    def adder(x, y):\n        if x and y:\n            counter[0] += 1\n        return x + y\n\n    return adder, counter\n\n\ndef test_multisubset(numcount, setcount):\n    numbers = [random.randrange(10**20) for _ in range(numcount)]\n    subsets = [\n        {i for i in range(numcount) if random.randrange(2)} for i in range(setcount)\n    ]\n    adder, counter = make_mock_adder()\n    o = multisubset(numbers, subsets, adder=adder)\n    for output, subset in zip(o, subsets):\n        assert output == sum([numbers[x] for x in subset])\n\n\ndef test_lincomb(numcount, bitlength=256):\n    numbers = [random.randrange(10**20) for _ in range(numcount)]\n    factors = [random.randrange(2**bitlength) for _ in range(numcount)]\n    adder, counter = make_mock_adder()\n    o = lincomb(numbers, factors, adder=adder)\n    assert o == sum([n * f for n, f in zip(numbers, factors)])\n    total_ones = sum(bin(f).count(\"1\") for f in factors)\n    print(\"Naive operation count: %d\" % (bitlength * numcount + total_ones))\n    print(\"Optimized operation count: %d\" % (bitlength * 2 + counter[0]))\n    print(\n        \"Optimization factor: %.2f\"\n        % ((bitlength * numcount + total_ones) / (bitlength * 2 + counter[0]))\n    )\n\n\nif __name__ == \"__main__\":\n    test_lincomb(int(sys.argv[1]) if len(sys.argv) >= 2 else 80)\n"
  },
  {
    "path": "poly.py",
    "content": "from curve import Scalar\nfrom enum import Enum\n\n\nclass Basis(Enum):\n    LAGRANGE = 1\n    MONOMIAL = 2\n\n\nclass Polynomial:\n    values: list[Scalar]\n    basis: Basis\n\n    def __init__(self, values: list[Scalar], basis: Basis):\n        assert all(isinstance(x, Scalar) for x in values)\n        assert isinstance(basis, Basis)\n        self.values = values\n        self.basis = basis\n\n    def __eq__(self, other):\n        return (self.basis == other.basis) and (self.values == other.values)\n\n    def __add__(self, other):\n        if isinstance(other, Polynomial):\n            assert len(self.values) == len(other.values)\n            assert self.basis == other.basis\n\n            return Polynomial(\n                [x + y for x, y in zip(self.values, other.values)],\n                self.basis,\n            )\n        else:\n            assert isinstance(other, Scalar)\n            if self.basis == Basis.LAGRANGE:\n                return Polynomial(\n                    [x + other for x in self.values],\n                    self.basis,\n                )\n            else:\n                return Polynomial(\n                    [self.values[0] + other] + self.values[1:],\n                    self.basis\n                )\n\n    def __sub__(self, other):\n        if isinstance(other, Polynomial):\n            assert len(self.values) == len(other.values)\n            assert self.basis == other.basis\n\n            return Polynomial(\n                [x - y for x, y in zip(self.values, other.values)],\n                self.basis,\n            )\n        else:\n            assert isinstance(other, Scalar)\n            if self.basis == Basis.LAGRANGE:\n                return Polynomial(\n                    [x - other for x in self.values],\n                    self.basis,\n                )\n            else:\n                return Polynomial(\n                    [self.values[0] - other] + self.values[1:],\n                    self.basis\n                )\n\n\n    def __mul__(self, other):\n        if isinstance(other, Polynomial):\n            assert self.basis == Basis.LAGRANGE\n            assert self.basis == other.basis\n            assert len(self.values) == len(other.values)\n\n            return Polynomial(\n                [x * y for x, y in zip(self.values, other.values)],\n                self.basis,\n            )\n        else:\n            assert isinstance(other, Scalar)\n            return Polynomial(\n                [x * other for x in self.values],\n                self.basis,\n            )\n\n    def __truediv__(self, other):\n        if isinstance(other, Polynomial):\n            assert self.basis == Basis.LAGRANGE\n            assert self.basis == other.basis\n            assert len(self.values) == len(other.values)\n\n            return Polynomial(\n                [x / y for x, y in zip(self.values, other.values)],\n                self.basis,\n            )\n        else:\n            assert isinstance(other, Scalar)\n            return Polynomial(\n                [x / other for x in self.values],\n                self.basis,\n            )\n\n    def shift(self, shift: int):\n        assert self.basis == Basis.LAGRANGE\n        assert shift < len(self.values)\n\n        return Polynomial(\n            self.values[shift:] + self.values[:shift],\n            self.basis,\n        )\n\n    # Convenience method to do FFTs specifically over the subgroup over which\n    # all of the proofs are operating\n    def fft(self, inv=False):\n        # Fast Fourier transform, used to convert between polynomial coefficients\n        # and a list of evaluations at the roots of unity\n        # See https://vitalik.ca/general/2019/05/12/fft.html\n        def _fft(vals, modulus, roots_of_unity):\n            if len(vals) == 1:\n                return vals\n            L = _fft(vals[::2], modulus, roots_of_unity[::2])\n            R = _fft(vals[1::2], modulus, roots_of_unity[::2])\n            o = [0] * len(vals)\n            for i, (x, y) in enumerate(zip(L, R)):\n                y_times_root = y * roots_of_unity[i]\n                o[i] = (x + y_times_root) % modulus\n                o[i + len(L)] = (x - y_times_root) % modulus\n            return o\n\n        roots = [x.n for x in Scalar.roots_of_unity(len(self.values))]\n        o, nvals = Scalar.field_modulus, [x.n for x in self.values]\n        if inv:\n            assert self.basis == Basis.LAGRANGE\n            # Inverse FFT\n            invlen = Scalar(1) / len(self.values)\n            reversed_roots = [roots[0]] + roots[1:][::-1]\n            return Polynomial(\n                [Scalar(x) * invlen for x in _fft(nvals, o, reversed_roots)],\n                Basis.MONOMIAL,\n            )\n        else:\n            assert self.basis == Basis.MONOMIAL\n            # Regular FFT\n            return Polynomial(\n                [Scalar(x) for x in _fft(nvals, o, roots)], Basis.LAGRANGE\n            )\n\n    def ifft(self):\n        return self.fft(True)\n\n    # Converts a list of evaluations at [1, w, w**2... w**(n-1)] to\n    # a list of evaluations at\n    # [offset, offset * q, offset * q**2 ... offset * q**(4n-1)] where q = w**(1/4)\n    # This lets us work with higher-degree polynomials, and the offset lets us\n    # avoid the 0/0 problem when computing a division (as long as the offset is\n    # chosen randomly)\n    def to_coset_extended_lagrange(self, offset):\n        assert self.basis == Basis.LAGRANGE\n        group_order = len(self.values)\n        x_powers = self.ifft().values\n        x_powers = [(offset**i * x) for i, x in enumerate(x_powers)] + [Scalar(0)] * (\n            group_order * 3\n        )\n        return Polynomial(x_powers, Basis.MONOMIAL).fft()\n\n    # Convert from offset form into coefficients\n    # Note that we can't make a full inverse function of to_coset_extended_lagrange\n    # because the output of this might be a deg >= n polynomial, which cannot\n    # be expressed via evaluations at n roots of unity\n    def coset_extended_lagrange_to_coeffs(self, offset):\n        assert self.basis == Basis.LAGRANGE\n\n        shifted_coeffs = self.ifft().values\n        inv_offset = 1 / offset\n        return Polynomial(\n            [v * inv_offset**i for (i, v) in enumerate(shifted_coeffs)],\n            Basis.MONOMIAL,\n        )\n\n    # Given a polynomial expressed as a list of evaluations at roots of unity,\n    # evaluate it at x directly, without using an FFT to covert to coeffs first\n    def barycentric_eval(self, x: Scalar):\n        assert self.basis == Basis.LAGRANGE\n\n        order = len(self.values)\n        roots_of_unity = Scalar.roots_of_unity(order)\n        return (\n            (Scalar(x) ** order - 1)\n            / order\n            * sum(\n                [\n                    value * root / (x - root)\n                    for value, root in zip(self.values, roots_of_unity)\n                ]\n            )\n        )\n"
  },
  {
    "path": "prover.py",
    "content": "from compiler.program import Program, CommonPreprocessedInput\nfrom utils import *\nfrom setup import *\nfrom typing import Optional\nfrom dataclasses import dataclass\nfrom transcript import Transcript, Message1, Message2, Message3, Message4, Message5\nfrom poly import Polynomial, Basis\n\n\n@dataclass\nclass 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[\"t_lo_1\"] = self.msg_3.t_lo_1\n        proof[\"t_mid_1\"] = self.msg_3.t_mid_1\n        proof[\"t_hi_1\"] = self.msg_3.t_hi_1\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[\"s1_eval\"] = self.msg_4.s1_eval\n        proof[\"s2_eval\"] = self.msg_4.s2_eval\n        proof[\"z_shifted_eval\"] = self.msg_4.z_shifted_eval\n        proof[\"W_z_1\"] = self.msg_5.W_z_1\n        proof[\"W_zw_1\"] = self.msg_5.W_zw_1\n        return proof\n\n\n@dataclass\nclass 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        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        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 for A, B, C, corresponding:\n        # - A_values: witness[program.wires()[i].L]\n        # - B_values: witness[program.wires()[i].R]\n        # - C_values: witness[program.wires()[i].O]\n\n        # Construct A, B, C Lagrange interpolation polynomials for\n        # A_values, B_values, C_values\n\n        # Compute a_1, b_1, c_1 commitments to A, B, C polynomials\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 a_1, b_1, c_1\n        return Message1(a_1, b_1, c_1)\n\n    def round_2(self) -> Message2:\n        group_order = self.group_order\n        setup = self.setup\n\n        # Using A, B, C, values, and pk.S1, pk.S2, pk.S3, compute\n        # Z_values for permutation grand product polynomial Z\n        #\n        # Note the convenience function:\n        #       self.rlc(val1, val2) = val_1 + self.beta * val_2 + gamma\n\n        # Check that the last term Z_n = 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        # Construct Z, Lagrange interpolation polynomial for Z_values\n        # Cpmpute z_1 commitment to Z polynomial\n\n        # Return z_1\n        return Message2(z_1)\n\n    def round_3(self) -> Message3:\n        group_order = self.group_order\n        setup = self.setup\n\n        # Compute the quotient polynomial\n\n        # List of roots of unity at 4x fineness, i.e. the powers of µ\n        # where µ^(4n) = 1\n\n        # Using self.fft_expand, move A, B, C into coset extended Lagrange basis\n\n        # Expand public inputs polynomial PI into coset extended Lagrange\n\n        # Expand selector polynomials pk.QL, pk.QR, pk.QM, pk.QO, pk.QC\n        # into the coset extended Lagrange basis\n\n        # Expand permutation grand product polynomial Z into coset extended\n        # Lagrange basis\n\n        # Expand shifted Z(ω) into coset extended Lagrange basis\n\n        # Expand permutation polynomials pk.S1, pk.S2, pk.S3 into coset\n        # extended Lagrange basis\n\n        # Compute Z_H = X^N - 1, also in evaluation form in the coset\n\n        # Compute L0, the Lagrange basis polynomial that evaluates to 1 at x = 1 = ω^0\n        # and 0 at other roots of unity\n\n        # Expand L0 into the coset extended Lagrange basis\n        L0_big = self.fft_expand(\n            Polynomial([Scalar(1)] + [Scalar(0)] * (group_order - 1), Basis.LAGRANGE)\n        )\n\n        # Compute the quotient polynomial (called T(x) in the paper)\n        # It is only possible to construct this polynomial if the following\n        # equations are true at all roots of unity {1, w ... w^(n-1)}:\n        # 1. All gates are correct:\n        #    A * QL + B * QR + A * B * QM + C * QO + PI + QC = 0\n        #\n        # 2. The permutation accumulator is valid:\n        #    Z(wx) = Z(x) * (rlc of A, X, 1) * (rlc of B, 2X, 1) *\n        #                   (rlc of C, 3X, 1) / (rlc of A, S1, 1) /\n        #                   (rlc of B, S2, 1) / (rlc of C, S3, 1)\n        #    rlc = random linear combination: term_1 + beta * term2 + gamma * term3\n        #\n        # 3. The permutation accumulator equals 1 at the start point\n        #    (Z - 1) * L0 = 0\n        #    L0 = Lagrange polynomial, equal at all roots of unity except 1\n\n        # Sanity check: QUOT has degree < 3n\n        assert (\n            self.expanded_evals_to_coeffs(QUOT_big).values[-group_order:]\n            == [0] * group_order\n        )\n        print(\"Generated the quotient polynomial\")\n\n        # Split up T into T1, T2 and T3 (needed because T has degree 3n - 4, so is\n        # too big for the trusted setup)\n\n        # Sanity check that we've computed T1, T2, T3 correctly\n        assert (\n            T1.barycentric_eval(fft_cofactor)\n            + T2.barycentric_eval(fft_cofactor) * fft_cofactor**group_order\n            + T3.barycentric_eval(fft_cofactor) * fft_cofactor ** (group_order * 2)\n        ) == QUOT_big.values[0]\n\n        print(\"Generated T1, T2, T3 polynomials\")\n\n        # Compute commitments t_lo_1, t_mid_1, t_hi_1 to T1, T2, T3 polynomials\n\n        # Return t_lo_1, t_mid_1, t_hi_1\n        return Message3(t_lo_1, t_mid_1, t_hi_1)\n\n    def round_4(self) -> Message4:\n        # Compute evaluations to be used in constructing the linearization polynomial.\n\n        # Compute a_eval = A(zeta)\n        # Compute b_eval = B(zeta)\n        # Compute c_eval = C(zeta)\n        # Compute s1_eval = pk.S1(zeta)\n        # Compute s2_eval = pk.S2(zeta)\n        # Compute z_shifted_eval = Z(zeta * ω)\n\n        # Return a_eval, b_eval, c_eval, s1_eval, s2_eval, z_shifted_eval\n        return Message4(a_eval, b_eval, c_eval, s1_eval, s2_eval, z_shifted_eval)\n\n    def round_5(self) -> Message5:\n        # Evaluate the Lagrange basis polynomial L0 at zeta\n        # Evaluate the vanishing polynomial Z_H(X) = X^n - 1 at zeta\n\n        # Move T1, T2, T3 into the coset extended Lagrange basis\n        # Move pk.QL, pk.QR, pk.QM, pk.QO, pk.QC into the coset extended Lagrange basis\n        # Move Z into the coset extended Lagrange basis\n        # Move pk.S3 into the coset extended Lagrange basis\n\n        # Compute the \"linearization polynomial\" R. This is a clever way to avoid\n        # needing to provide evaluations of _all_ the polynomials that we are\n        # checking an equation betweeen: instead, we can \"skip\" the first\n        # multiplicand in each term. The idea is that we construct a\n        # polynomial which is constructed to equal 0 at Z only if the equations\n        # that we are checking are correct, and which the verifier can reconstruct\n        # the KZG commitment to, and we provide proofs to verify that it actually\n        # equals 0 at Z\n        #\n        # In order for the verifier to be able to reconstruct the commitment to R,\n        # it has to be \"linear\" in the proof items, hence why we can only use each\n        # proof item once; any further multiplicands in each term need to be\n        # replaced with their evaluations at Z, which do still need to be provided\n\n        # Commit to R\n\n        # Sanity-check R\n        assert R.barycentric_eval(zeta) == 0\n\n        print(\"Generated linearization polynomial R\")\n\n        # Generate proof that W(z) = 0 and that the provided evaluations of\n        # A, B, C, S1, S2 are correct\n\n        # Move A, B, C into the coset extended Lagrange basis\n        # Move pk.S1, pk.S2 into the coset extended Lagrange basis\n\n        # In the COSET EXTENDED LAGRANGE BASIS,\n        # Construct W_Z = (\n        #     R\n        #   + v * (A - a_eval)\n        #   + v**2 * (B - b_eval)\n        #   + v**3 * (C - c_eval)\n        #   + v**4 * (S1 - s1_eval)\n        #   + v**5 * (S2 - s2_eval)\n        # ) / (X - zeta)\n\n        # Check that degree of W_z is not greater than n\n        assert W_z_coeffs[group_order:] == [0] * (group_order * 3)\n\n        # Compute W_z_1 commitment to W_z\n\n        # Generate proof that the provided evaluation of Z(z*w) is correct. This\n        # awkwardly different term is needed because the permutation accumulator\n        # polynomial Z is the one place where we have to check between adjacent\n        # coordinates, and not just within one coordinate.\n        # In other words: Compute W_zw = (Z - z_shifted_eval) / (X - zeta * ω)\n\n        # Check that degree of W_z is not greater than n\n        assert W_zw_coeffs[group_order:] == [0] * (group_order * 3)\n\n        # Compute W_z_1 commitment to W_z\n\n        print(\"Generated final quotient witness polynomials\")\n\n        # Return W_z_1, W_zw_1\n        return Message5(W_z_1, W_zw_1)\n\n    def fft_expand(self, x: Polynomial):\n        return x.to_coset_extended_lagrange(self.fft_cofactor)\n\n    def expanded_evals_to_coeffs(self, x: Polynomial):\n        return x.coset_extended_lagrange_to_coeffs(self.fft_cofactor)\n\n    def rlc(self, term_1, term_2):\n        return term_1 + term_2 * self.beta + self.gamma\n"
  },
  {
    "path": "pyproject.toml",
    "content": "[tool.poetry]\nname = \"plonkathon\"\nversion = \"0.1.0\"\ndescription = \"A simple Python implementation of PLONK adapted from py_plonk\"\nauthors = [\"0xPARC / Vitalik Buterin\"]\nlicense = \"MIT\"\nreadme = \"README.md\"\n\n[tool.poetry.dependencies]\npython = \"^3.9\"\npy-ecc = \"^6.0.0\"\nmerlin = {git = \"https://github.com/nalinbhardwaj/curdleproofs.pie\", rev = \"master\", subdirectory = \"merlin\"}\n\n[tool.poetry.group.dev.dependencies]\nmypy = \"^0.991\"\nblack = \"^22.12.0\"\n\n[build-system]\nrequires = [\"poetry-core\"]\nbuild-backend = \"poetry.core.masonry.api\"\n\n[tool.mypy]\nexplicit_package_bases = true"
  },
  {
    "path": "setup.py",
    "content": "from utils import *\nimport py_ecc.bn128 as b\nfrom curve import ec_lincomb, G1Point, G2Point\nfrom compiler.program import CommonPreprocessedInput\nfrom verifier import VerificationKey\nfrom dataclasses import dataclass\nfrom poly import Polynomial, Basis\n\n# Recover the trusted setup from a file in the format used in\n# https://github.com/iden3/snarkjs#7-prepare-phase-2\nSETUP_FILE_G1_STARTPOS = 80\nSETUP_FILE_POWERS_POS = 60\n\n\n@dataclass\nclass Setup(object):\n    #   ([1]₁, [x]₁, ..., [x^{d-1}]₁)\n    # = ( G,    xG,  ...,  x^{d-1}G ), where G is a generator of G_2\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 from_file(cls, filename):\n        contents = open(filename, \"rb\").read()\n        # Byte 60 gives you the base-2 log of how many powers there are\n        powers = 2 ** contents[SETUP_FILE_POWERS_POS]\n        # Extract G1 points, which start at byte 80\n        values = [\n            int.from_bytes(contents[i : i + 32], \"little\")\n            for i in range(\n                SETUP_FILE_G1_STARTPOS, SETUP_FILE_G1_STARTPOS + 32 * powers * 2, 32\n            )\n        ]\n        assert max(values) < b.field_modulus\n        # The points are encoded in a weird encoding, where all x and y points\n        # are multiplied by a factor (for montgomery optimization?). We can\n        # extract the factor because we know the first point is the generator.\n        factor = b.FQ(values[0]) / b.G1[0]\n        values = [b.FQ(x) / factor for x in values]\n        powers_of_x = [(values[i * 2], values[i * 2 + 1]) for i in range(powers)]\n        print(\"Extracted G1 side, X^1 point: {}\".format(powers_of_x[1]))\n        # Search for start of G2 points. We again know that the first point is\n        # the generator.\n        pos = SETUP_FILE_G1_STARTPOS + 32 * powers * 2\n        target = (factor * b.G2[0].coeffs[0]).n\n        while pos < len(contents):\n            v = int.from_bytes(contents[pos : pos + 32], \"little\")\n            if v == target:\n                break\n            pos += 1\n        print(\"Detected start of G2 side at byte {}\".format(pos))\n        X2_encoding = contents[pos + 32 * 4 : pos + 32 * 8]\n        X2_values = [\n            b.FQ(int.from_bytes(X2_encoding[i : i + 32], \"little\")) / factor\n            for i in range(0, 128, 32)\n        ]\n        X2 = (b.FQ2(X2_values[:2]), b.FQ2(X2_values[2:]))\n        assert b.is_on_curve(X2, b.b2)\n        print(\"Extracted G2 side, X^1 point: {}\".format(X2))\n        # assert b.pairing(b.G2, powers_of_x[1]) == b.pairing(X2, b.G1)\n        # print(\"X^1 points checked consistent\")\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        assert values.basis == Basis.LAGRANGE\n\n        # Run inverse FFT to convert values from Lagrange basis to monomial basis\n        # Optional: Check values size does not exceed maximum power setup can handle\n        # Compute linear combination of setup with values\n        return NotImplemented\n\n    # Generate the verification key for this program with the given setup\n    def verification_key(self, pk: CommonPreprocessedInput) -> VerificationKey:\n        # Create the appropriate VerificationKey object\n        return NotImplemented\n"
  },
  {
    "path": "test/__init__.py",
    "content": ""
  },
  {
    "path": "test/main.plonk.vkey-58.json",
    "content": "{\n  \"protocol\": \"plonk\",\n  \"curve\": \"bn128\",\n  \"nPublic\": 0,\n  \"power\": 3,\n  \"k1\": \"2\",\n  \"k2\": \"3\",\n  \"Qm\": [\n    \"10294367845524522889674980414658158979115219665406612861401259333422895729896\",\n    \"17339696279167455564514853058684962930296864414660175742312401951183098671156\",\n    \"1\"\n  ],\n  \"Ql\": [\n    \"14297155691368363150439281660551929853142513799648244067851273621337387750022\",\n    \"12012534117624137096359211205114297110997558611632571627258110765766724342420\",\n    \"1\"\n  ],\n  \"Qr\": [\n    \"14297155691368363150439281660551929853142513799648244067851273621337387750022\",\n    \"9875708754215138125887194540142977977698752545665252035430927128878501866163\",\n    \"1\"\n  ],\n  \"Qo\": [\n    \"9979011916860674833876589507703522468002507052909770370053733666594800845094\",\n    \"14729630273693641999471892596240127436204275281179182844090045631345819602206\",\n    \"1\"\n  ],\n  \"Qc\": [\n    \"0\",\n    \"1\",\n    \"0\"\n  ],\n  \"S1\": [\n    \"12298616283483778451735558489316417273025867766441336478133033639422100496973\",\n    \"15676047185797786755474337421099990612245072514136746262179062664360148735456\",\n    \"1\"\n  ],\n  \"S2\": [\n    \"3285281999993628756611785365316971613389533139400689700975461452204136786218\",\n    \"1617843731485072605069101897885581356842856822348044818321641451407219765267\",\n    \"1\"\n  ],\n  \"S3\": [\n    \"15275313728188330177132414910553475149781975424851628063518044954012175053863\",\n    \"6189857575419431040150205313691546343485409281989411607909724166107587561616\",\n    \"1\"\n  ],\n  \"X_2\": [\n    [\n      \"21831381940315734285607113342023901060522397560371972897001948545212302161822\",\n      \"17231025384763736816414546592865244497437017442647097510447326538965263639101\"\n    ],\n    [\n      \"2388026358213174446665280700919698872609886601280537296205114254867301080648\",\n      \"11507326595632554467052522095592665270651932854513688777769618397986436103170\"\n    ],\n    [\n      \"1\",\n      \"0\"\n    ]\n  ],\n  \"w\": \"19540430494807482326159819597004422086093766032135589407132600596362845576832\"\n}\n"
  },
  {
    "path": "test/main.plonk.vkey-59.json",
    "content": "{\n  \"protocol\": \"plonk\",\n  \"curve\": \"bn128\",\n  \"nPublic\": 1,\n  \"power\": 3,\n  \"k1\": \"2\",\n  \"k2\": \"3\",\n  \"Qm\": [\n    \"10294367845524522889674980414658158979115219665406612861401259333422895729896\",\n    \"17339696279167455564514853058684962930296864414660175742312401951183098671156\",\n    \"1\"\n  ],\n  \"Ql\": [\n    \"14297155691368363150439281660551929853142513799648244067851273621337387750022\",\n    \"9875708754215138125887194540142977977698752545665252035430927128878501866163\",\n    \"1\"\n  ],\n  \"Qr\": [\n    \"0\",\n    \"1\",\n    \"0\"\n  ],\n  \"Qo\": [\n    \"10294367845524522889674980414658158979115219665406612861401259333422895729896\",\n    \"4548546592671819657731552686572312158399446742637647920376635943462127537427\",\n    \"1\"\n  ],\n  \"Qc\": [\n    \"0\",\n    \"1\",\n    \"0\"\n  ],\n  \"S1\": [\n    \"2694761611667402433549058650401049833973608710551146850129171008254242491412\",\n    \"4407815622841625592989621790140274705225113068749062773093818734897989348824\",\n    \"1\"\n  ],\n  \"S2\": [\n    \"174950894878901504258554221888959060942005622520060188523927601854050691737\",\n    \"4218570225917094256073281485929194442981718242484973947497628954679969816940\",\n    \"1\"\n  ],\n  \"S3\": [\n    \"5287191920074181963852791792640835974097875195213946104826736553520071559657\",\n    \"20309358409622118500558363825899879958276827259299017541094321370900713472551\",\n    \"1\"\n  ],\n  \"X_2\": [\n    [\n      \"21831381940315734285607113342023901060522397560371972897001948545212302161822\",\n      \"17231025384763736816414546592865244497437017442647097510447326538965263639101\"\n    ],\n    [\n      \"2388026358213174446665280700919698872609886601280537296205114254867301080648\",\n      \"11507326595632554467052522095592665270651932854513688777769618397986436103170\"\n    ],\n    [\n      \"1\",\n      \"0\"\n    ]\n  ],\n  \"w\": \"19540430494807482326159819597004422086093766032135589407132600596362845576832\"\n}\n"
  },
  {
    "path": "test/main.plonk.vkey.json",
    "content": "{\n  \"protocol\": \"plonk\",\n  \"curve\": \"bn128\",\n  \"nPublic\": 0,\n  \"power\": 3,\n  \"k1\": \"2\",\n  \"k2\": \"3\",\n  \"Qm\": [\n    \"14297155691368363150439281660551929853142513799648244067851273621337387750022\",\n    \"12012534117624137096359211205114297110997558611632571627258110765766724342420\",\n    \"1\"\n  ],\n  \"Ql\": [\n    \"0\",\n    \"1\",\n    \"0\"\n  ],\n  \"Qr\": [\n    \"0\",\n    \"1\",\n    \"0\"\n  ],\n  \"Qo\": [\n    \"14297155691368363150439281660551929853142513799648244067851273621337387750022\",\n    \"9875708754215138125887194540142977977698752545665252035430927128878501866163\",\n    \"1\"\n  ],\n  \"Qc\": [\n    \"0\",\n    \"1\",\n    \"0\"\n  ],\n  \"S1\": [\n    \"3514537020795837778176804346752472700561147058607627192700022438897457917853\",\n    \"8628914924381244881833041972511347057770541232212012588202494401021191294886\",\n    \"1\"\n  ],\n  \"S2\": [\n    \"10270239105545965856893095866462682446854660969221881327142885898762754802290\",\n    \"18254124398103062012489482857980956345119701560599595388652594443232359794243\",\n    \"1\"\n  ],\n  \"S3\": [\n    \"7581161054812340292915025932276548889828390062297736519460377063075571422119\",\n    \"11644391033980170207157104232723275903489480225585820213777359398785819753031\",\n    \"1\"\n  ],\n  \"X_2\": [\n    [\n      \"21831381940315734285607113342023901060522397560371972897001948545212302161822\",\n      \"17231025384763736816414546592865244497437017442647097510447326538965263639101\"\n    ],\n    [\n      \"2388026358213174446665280700919698872609886601280537296205114254867301080648\",\n      \"11507326595632554467052522095592665270651932854513688777769618397986436103170\"\n    ],\n    [\n      \"1\",\n      \"0\"\n    ]\n  ],\n  \"w\": \"19540430494807482326159819597004422086093766032135589407132600596362845576832\"\n}"
  },
  {
    "path": "test/mini_poseidon.py",
    "content": "from py_ecc.fields.field_elements import FQ as Field\nfrom py_ecc import bn128 as b\nimport json\nfrom curve import Scalar\n\n# Mimics the Poseidon hash for params:\n#\n# p                    = b.curve_order\n# security level       = 128\n# alpha                = 5\n# input size           = 2\n# t (inner state size) = 3\n# full round count     = 8 (4 on each side)\n# partial round count  = 56\n#\n# Tested compatible with the implementation at\n# https://github.com/ingonyama-zk/poseidon-hash\n\nrc = [\n    [Scalar(a), Scalar(b), Scalar(c)]\n    for (a, b, c) in json.load(open(\"test/poseidon_rc.json\"))\n]\n\nmds = [Scalar(1) / i for i in range(3, 8)]\n\n\ndef poseidon_hash(in1, in2):\n    L, M, R = Scalar(in1), Scalar(in2), Scalar(0)\n    for i in range(64):\n        L = (L + rc[i][0]) ** 5\n        M += rc[i][1]\n        R += rc[i][2]\n        if i < 4 or i >= 60:\n            M = M**5\n            R = R**5\n\n        (L, M, R) = (\n            (L * mds[0] + M * mds[1] + R * mds[2]),\n            (L * mds[1] + M * mds[2] + R * mds[3]),\n            (L * mds[2] + M * mds[3] + R * mds[4]),\n        )\n    return M\n"
  },
  {
    "path": "test/poseidon_rc.json",
    "content": "[\n    [\n        20036611579827150559091469005844175073625940102952070649817884191797764107075,\n        12833584042949159986565784794014151972247796182705941809242049488642050965764,\n        20460265239335923814507658708649753625122635556480788121847315929099732630160\n    ],\n    [\n        2327433556176440050072122789158937095349015830966296688052004527413893088211,\n        18928950063089170035907725405866336756839578021518946851617532132424258251714,\n        1498343560697108822498273562774300261830404487630763796790963549875888044170\n    ],\n    [\n        10894060797649430325830102265601652830696874259935838091685396334824506870974,\n        10724282231429846920698162462239380161919181405938370346044877808218935812722,\n        12900822782410361357346122719301635014682733276151991957222022994303272655598\n    ],\n    [\n        12281187892168564254318627327513968787860396596558174241050192282039059721015,\n        20650821582781615100585145227212437582142977940055542470636589623435638972974,\n        10629768372315475914205126397268924392851627409938344639024631693382468091238\n    ],\n    [\n        9863412385841952362091849999302344670551536862597895561278739824943955436428,\n        12647165719356049717276082841324782064533712829755618494899697790036659206896,\n        7086910375195160419007414095698529055230556352178368086440281946317648386880\n    ],\n    [\n        3534616497502400285225524656623573157846352519939301302938694874082052072531,\n        7614193770971378304733014379518957809467088529876055633743098033063991937988,\n        5524517084220026554140171636731805466222414790303448093874175843707842149296\n    ],\n    [\n        6465594825680448812389988135904471523354179427789234651961801931723687100590,\n        13069961729933398081974430219710353490873542026760226429212509555519019152029,\n        20442858075434477469830277030056901344749316967720547608602986649544991223254\n    ],\n    [\n        2545061208850080799716748664877827143438576541109233212439269028123126995283,\n        10590804083025333327978466704873260869373451484886744716135293454855421776182,\n        2971801848580094041349173303826341900320100639936132696391170461147651791599\n    ],\n    [\n        6457488363353205666653234338157121013068090150457850441006724736538329117150,\n        6828448758803946230526683011338306528814963654399318058641145666219257935090,\n        21220571338383411884578252416062534137998565615530587243864120751033430171639\n    ],\n    [\n        3026330605417380631686907051427895875297995068925155493166857059368062462182,\n        731369258216230615295136666016875570919026708063642185179381726677850241678,\n        3976056293038938116541801753309319302894026070844887224129384450572018141696\n    ],\n    [\n        18191771098193274629372635247562817907522214770330118615218435763543549998849,\n        5725653743840821273222706240430501486301231493172435748281584336969930288424,\n        11315632284569557454743403424858271195359234409995993988812038704370836780935\n    ],\n    [\n        1554548474507260015710470936845569353450016357484523407182389210520378329681,\n        4043703307392904605083268263455220115794158610223459176828550627728253474746,\n        5147341731800747507000152527876279909425699775431448328748957228886538246347\n    ],\n    [\n        316000240737142429159443531715814797123772606639483557601339985787766020188,\n        9201639581870600108429525788569689133841787929111628476773761218482387877013,\n        76452313668236664875560655448227033496493983051536430171258861062052191226\n    ],\n    [\n        5503629712269636063700501314848921574392241749467380315893984871270736181455,\n        6403315002483104841708481320831358853734637389275574672824410469941470947501,\n        18767770632814906108167838116228345063195454403262264488888241387554264238660\n    ],\n    [\n        373847975877840160326133445669241099117560443926620032996032499044947786206,\n        21583421457099893419202355352294471270616995660291443431182584120442735921995,\n        9983253793391637270425117523681631923584433357117997691348261314215409417977\n    ],\n    [\n        15636580057941931488663342701921482616732902735747145545560675417041959528168,\n        13987037342707768451704041319080865947736730484739494822133564732867708182129,\n        17012807810939334924306150791435050869173026502739324471672493058506335677602\n    ],\n    [\n        20445006141370558272839096434547759512337597485395700957462476332341263470795,\n        8919512276260973838722894975323643726180891487427800112399328013257265098884,\n        5121935724947519798577565968119435009768355186616305612599005246528173860072\n    ],\n    [\n        14390107556389236497954661819339690487046314759239155128415464040636056995180,\n        15584006578697113728990467783042934569634784959449117771168091586065677667289,\n        1894113083403807703900670244801865115743187370284525340414788037116278755190\n    ],\n    [\n        10705487296676443016268478693800719585977460625398657826409495078078625785779,\n        19873128801331021326057263165428495080386730321288221112754217581989746046187,\n        9535868612309525085410638759100096355215966759273780340615254181304597180876\n    ],\n    [\n        19343524237095964073059512688602276354221693507832098135091590955813103860375,\n        6822002759027581661397720296774923367218318024525870595858692787302610345838,\n        20926315677022871764300366782923548421637739455364254388830334528197327379743\n    ],\n    [\n        21136323309320513006223521228388165385265667945268848018308200096388016784087,\n        6225919113652258029432904641421409709120407088870082365870799895213524311811,\n        21311480157558702908741067111471547471593313246226635294078364811203226453373\n    ],\n    [\n        14198190262525448445981718442399104668575545443532459171696901751739593839857,\n        12459931960399741184291910088669153631848894479486118281597822681015713751918,\n        6907332911214986094624912982222847460679912779076715919016253686198271600395\n    ],\n    [\n        3393759605179403680839415959192212038674751416897927769901727596993250674985,\n        16080342864802322672358399934807873723002376453489101187414456500323890285012,\n        16917851076945834679745654051495359311013500015444379399867812928847812812510\n    ],\n    [\n        5864181025763406655115777223889072827293022750325008731169039770623910915277,\n        10954281474369541712624748628862576225398375445664590376470032986861704444851,\n        13047511983320021446527738191176357616721874803157974120155396052149484216634\n    ],\n    [\n        17384764166393886603954067639141617082485023441458666471633998399316404519396,\n        2362259915869885612166233838782699563177506073802792748904427749547307239091,\n        3449611877995076637446149781868130604254776727586032907573864421212958534148\n    ],\n    [\n        8105036071757751706553034047983212897131570345551297243248892142702129409678,\n        4057622885794382454162574102227879765926272192471614700281108386186208378108,\n        19253655284986776930766796883935217167812408888646382124396693895764793065618\n    ],\n    [\n        14562077372515720927783340475811850803822078325552042696568002577429594161808,\n        7014243471382661219949882659461009076113783570411660059539677464905457050066,\n        11119991552991434794635736236084481059792337090668881972430146540524522150444\n    ],\n    [\n        6920339933326447431851243390609339708453004519540083864540711586753457519075,\n        11162222425158351635121226742909449004096853940732577587985131267177476511460,\n        21584563699769356971941778663228493736429564166612353926541774478985941223655\n    ],\n    [\n        15214147914792054366849870322421302165632582490745675437262348323314349777499,\n        13722597825431950116506669418346236306300603980563381228999662274560534928630,\n        5918152979416198971639649518325331620625448907020294484718605076371986899062\n    ],\n    [\n        15678541263797653248465557518970463974123241861658113965486586969940493538159,\n        8175777680824278600606955282259166614282815213949021367965714288908365703981,\n        8574569171387093031118013554489261429423593158447011529314100526972055673141\n    ],\n    [\n        10918716126595714712314376320709488869257965519405486565051522412268603317681,\n        2532325191174382117711281885521238232961564999737136355944431967345212530064,\n        9114912432384157184380243669694781216466912274227843462720070724726171771343\n    ],\n    [\n        9691400567437687859118635218970787863193933880668112713672540475341700900162,\n        10028377139165856509931030519434327238309496767889858927872034975860560686709,\n        2849683057915704668087183484859329979263625107757384166719222732911939917981\n    ],\n    [\n        20815778942199761651592702319981868700183832939349741057726702071066236466370,\n        3621409685183637105257042956474539264816126311779453543298291481462284916375,\n        10642444461722023103186165843304806163340395145625014310520389928965799104093\n    ],\n    [\n        17907899976867921198783978602917240207228046151371192195898939004844562290779,\n        9146061947045243980414113080125198805700065329387764273046439006263867106802,\n        13405550010907066703864845579707432131752948528064342334643991332494785370038\n    ],\n    [\n        5273152122127100210394298964572549609798968014900036928547501207482676600182,\n        20346921636268799423535709508901024522528095837141171053270676908508776437195,\n        12308379321943972271977727620557057340570329211508147240067063174565210570182\n    ],\n    [\n        10770632218611704372085484530618691838222212586567712779168513852598498268856,\n        3259010024842168233969061619991654422834134029900348332440765888676313413487,\n        13595627212987049470638603008526754259254399837052142232458256128156882979850\n    ],\n    [\n        13686702530817919707738525718876497735372528847997348246778693823968904671413,\n        4460395761756237799418597895559539723093635900761406744282005033814770269920,\n        919904521071258469189758023145643742474862093406711345641316491878850805232\n    ],\n    [\n        19552067628327891935919890073139134992135919431541039306364452733276397042650,\n        10232151001993153761133312595886476849690259266379133316759758019413141271347,\n        19100291637068471015012417085176588278037788688620609514545142049341033683807\n    ],\n    [\n        17297878401483595514567069610271673264004048391410881314538210656783184883353,\n        9840592683860937355653821679751062912936830160133502129601080738765222581867,\n        12297898491098218555414910863121396110527474713633732916696718763044349837748\n    ],\n    [\n        15412894648987998885409062475029640813182264551444190226337524042414883759931,\n        4319236626081961726054666064727097684638760777375371898340440545864486507582,\n        18191973605703047070391590784739729000864209799853952244185039146413344974340\n    ],\n    [\n        5752404706735255136529409306453204906911390186311487545512488044087647004363,\n        19230596513152256569445337244613095254390639822535622716862845651176210035567,\n        2252265175955857562316757498871300534719440896551852902250776607815209442724\n    ],\n    [\n        9073931270846747895810603986335991449117293042147443815945435448584060282081,\n        6483856960596610771808987619300220938594487623302140190538043615094727502464,\n        637083280333510402345876123020430035110637297697948006826849236896141939084\n    ],\n    [\n        7961522592388438366485301812068782850165520204310798157690585659140643523427,\n        17678862179510708909425310131356019028135505241560682522031533174595818508320,\n        574878632188556442314330639459430853202335502659322827394752921941768574011\n    ],\n    [\n        3568116434701175901536016957970225427791238673348886522778060096654834078838,\n        209057191470494933977718281886471474348594564432534881082944258569955784687,\n        3800165250578186427973769620301873054994730397398528230413324355039263408752\n    ],\n    [\n        3639100050381197573204744243844108375909703224935458547540586483692218704877,\n        9218466811348574054855689012570491528574600673776655812307423675082651391951,\n        8596059170152859617707385585788415986680436507191992176622088596772810472999\n    ],\n    [\n        322732566805048175662260816011517006326081754353072089063398671083697385738,\n        7049453550072299437444402478358913479930492212486482712275051349319498748499,\n        20246786653614752389192328734689452300173121506317343804072039327223658303388\n    ],\n    [\n        20278117822137273667620939538959902468837011105813146111671754380981924989927,\n        1654160784815709253179744420396162626998410370662765393562093472726921196108,\n        5926459500982186461721807733417927873652323980102172461303859048787673285969\n    ],\n    [\n        9814654683585465702825961842499223555384560422004063024601314839443288937571,\n        20667359178767651545823868376438771969042007215942165947678357086227810746146,\n        15377207745218347185681329705726178336802704421362747654022727216790137478882\n    ],\n    [\n        19356050391225890781843193058227673417005845306283197872873417841669492437242,\n        12166372468984248061354090302429129252134172582728018687944601026858163494693,\n        18717100384016177736416179070153621959142213446659284814208081064231797419015\n    ],\n    [\n        5204341573325740834696285877748784802324386535160173798458534920759432537787,\n        19782883278004953315498469484333242486546087363994023561941134846660210215049,\n        16034955990523880906964021599001738382034941228857279292438387818659914390371\n    ],\n    [\n        14818160808196029852784679856539666583166658183438355275636086056715415083081,\n        17716280573405858003273309028354201201039005385690741490326620374206055544700,\n        20325174322201219249837009573437620189627639570996454573810011699858417613068\n    ],\n    [\n        3392148968210856354607867739690107473804056082222277341573277534219216150766,\n        19428278646034050357542301711807143225530095360718614464231862005201274256649,\n        12276431613202649950362446507866899415646220636973437014689940717149756340651\n    ],\n    [\n        2587427326658891351835083368358925497178797694851506580681724992222840444228,\n        15153709835422806446107529654085221072935438340920609707727892556855417866058,\n        15451366934210083469013986585953772077331429640485131755078569432020850609582\n    ],\n    [\n        7529206459496272002258451885407469611325302405318004333109879967079729143354,\n        11074797962501795518507846987686168533587240327062952171521584310243255490527,\n        12764552282499399390083087144902783986037307873849260653087645304074666033424\n    ],\n    [\n        13042650929397207114336562675355065717548627481583423774346222039669532727946,\n        16043108230643325415237962657577223570300108689567622235410238284480628436709,\n        9512396575137793662413202544406785640755021654701449580404337449461432257619\n    ],\n    [\n        1647285959695493539499324566741747088943357828231148957465068152734531543503,\n        21709648504024848772276178029483179772486495316295743876700927299213559938506,\n        21724258692495205310896265709064144808671841367275300438694185214351785954314\n    ],\n    [\n        1867255786401891769870675572668606218710278873748335003682799683714442945229,\n        7987721022707500619696671044839066336452195788164827456425525148091870118251,\n        6539200631055259414716243770171674046221430025597614000179562787351853958064\n    ],\n    [\n        8813689471242022713194444818607001333164230629904238460550879892354711745295,\n        20714180249817511375952366875224782387429058136430649457777093547955191502176,\n        18539711572616584819185239447633561813340291234100075574593487873259164787756\n    ],\n    [\n        17667981635370223776813093657005199235669153259421463967392158997697754245871,\n        5376160060481176473583403267888691728170726259835912852329136073184696997147,\n        4827362235909576508007307863435628112155590770759625133213539477210553500723\n    ],\n    [\n        13370319446619231361317124489467185868983690187607611069576690369092652259991,\n        7028188533891957864971355746639290885773084850501201875941442998993682972992,\n        13540296832979359410552476169105468590730420434093705252588595073215852455795\n    ],\n    [\n        12990891770605679281297891471045699934999991085025534577757867831342762325899,\n        12462076986357124836994040219766123609086672012765820033289064632297189251090,\n        1828489075147172049835809292018292715706474000231898985692281252669405838335\n    ],\n    [\n        3765746554097871584971351502123536140065794926761110131907071999370632640132,\n        3439916369559340392520868084385525387338766489009161266575789323989654664195,\n        20809812814789097985689246615784687916156945595313708545672843850284131223766\n    ],\n    [\n        13027707712686586554032640557539411722798759615913193118482673776883949392685,\n        6182352544550527301330938836100871624958555541496136265556139550324548491339,\n        12013693095010211530598130838725639174031637655633969485927984978914903872850\n    ],\n    [\n        16544062834867868837759941955456105371508093239200809853834322062372359189123,\n        14286779287597934647998556545702143097785912165468635922756805434029327391090,\n        3870485400879314278951355285041632089577731427251451910320055748195753746665\n    ]\n]"
  },
  {
    "path": "test.py",
    "content": "import pickle\nfrom TESTING_verifier_DO_NOT_OPEN import TestingVerificationKey\nfrom compiler.program import Program\nfrom curve import G1Point\nfrom poly import Basis, Polynomial\nfrom setup import Setup\nfrom prover import Prover\nfrom verifier import VerificationKey\nimport json\nfrom test.mini_poseidon import rc, mds, poseidon_hash\nfrom utils import *\n\n\ndef setup_test():\n    print(\"===setup_test===\")\n\n    setup = Setup.from_file(\"test/powersOfTau28_hez_final_11.ptau\")\n    dummy_values = Polynomial(\n        list(map(Scalar, [1, 2, 3, 4, 5, 6, 7, 8])), Basis.LAGRANGE\n    )\n    program = Program([\"c <== a * b\"], 8)\n    commitment = setup.commit(dummy_values)\n    assert commitment == G1Point(\n        (\n            16120260411117808045030798560855586501988622612038310041007562782458075125622,\n            3125847109934958347271782137825877642397632921923926105820408033549219695465,\n        )\n    )\n    vk = setup.verification_key(program.common_preprocessed_input())\n    assert (\n        vk.w\n        == 19540430494807482326159819597004422086093766032135589407132600596362845576832\n    )\n    print(\"Successfully created dummy commitment and verification key\")\n\n\ndef basic_test():\n    print(\"===basic_test===\")\n\n    # Extract 2^28 powers of tau\n    setup = Setup.from_file(\"test/powersOfTau28_hez_final_11.ptau\")\n    print(\"Extracted setup\")\n    program = Program([\"c <== a * b\"], 8)\n    vk = setup.verification_key(program.common_preprocessed_input())\n    print(\"Generated verification key\")\n    their_output = json.load(open(\"test/main.plonk.vkey.json\"))\n    for key in (\"Qm\", \"Ql\", \"Qr\", \"Qo\", \"Qc\", \"S1\", \"S2\", \"S3\", \"X_2\"):\n        if interpret_json_point(their_output[key]) != getattr(vk, key):\n            raise Exception(\n                \"Mismatch {}: ours {} theirs {}\".format(\n                    key, getattr(vk, key), their_output[key]\n                )\n            )\n    assert getattr(vk, \"w\") == int(their_output[\"w\"])\n    print(\"Basic test success\")\n    return setup\n\n\n# Equivalent to this zkrepl code:\n#\n# template Example () {\n#    signal input a;\n#    signal input b;\n#    signal c;\n#    c <== a * b + a;\n# }\ndef ab_plus_a_test(setup):\n    print(\"===ab_plus_a_test===\")\n\n    program = Program([\"ab === a - c\", \"-ab === a * b\"], 8)\n    vk = setup.verification_key(program.common_preprocessed_input())\n    print(\"Generated verification key\")\n    their_output = json.load(open(\"test/main.plonk.vkey-58.json\"))\n    for key in (\"Qm\", \"Ql\", \"Qr\", \"Qo\", \"Qc\", \"S1\", \"S2\", \"S3\", \"X_2\"):\n        if interpret_json_point(their_output[key]) != getattr(vk, key):\n            raise Exception(\n                \"Mismatch {}: ours {} theirs {}\".format(\n                    key, getattr(vk, key), their_output[key]\n                )\n            )\n    assert getattr(vk, \"w\") == int(their_output[\"w\"])\n    print(\"ab+a test success\")\n\n\ndef one_public_input_test(setup):\n    print(\"===one_public_input_test===\")\n\n    program = Program([\"c public\", \"c === a * b\"], 8)\n    vk = setup.verification_key(program.common_preprocessed_input())\n    print(\"Generated verification key\")\n    their_output = json.load(open(\"test/main.plonk.vkey-59.json\"))\n    for key in (\"Qm\", \"Ql\", \"Qr\", \"Qo\", \"Qc\", \"S1\", \"S2\", \"S3\", \"X_2\"):\n        if interpret_json_point(their_output[key]) != getattr(vk, key):\n            raise Exception(\n                \"Mismatch {}: ours {} theirs {}\".format(\n                    key, getattr(vk, key), their_output[key]\n                )\n            )\n    assert getattr(vk, \"w\") == int(their_output[\"w\"])\n    print(\"One public input test success\")\n\n\ndef prover_test_dummy_verifier(setup):\n    print(\"===prover_test_dummy_verifier===\")\n\n    print(\"Beginning prover test with test verifier\")\n    program = Program([\"e public\", \"c <== a * b\", \"e <== c * d\"], 8)\n    assignments = {\"a\": 3, \"b\": 4, \"c\": 12, \"d\": 5, \"e\": 60}\n    prover = Prover(setup, program)\n    proof = prover.prove(assignments)\n\n    print(\"Beginning test verification\")\n    program = Program([\"e public\", \"c <== a * b\", \"e <== c * d\"], 8)\n    public = [60]\n    vk = setup.verification_key(program.common_preprocessed_input())\n\n    vk_test = TestingVerificationKey(\n        group_order=vk.group_order,\n        Qm=vk.Qm,\n        Ql=vk.Ql,\n        Qr=vk.Qr,\n        Qo=vk.Qo,\n        Qc=vk.Qc,\n        S1=vk.S1,\n        S2=vk.S2,\n        S3=vk.S3,\n        X_2=vk.X_2,\n        w=vk.w,\n    )\n\n    assert vk_test.verify_proof_unoptimized(8, proof, public)\n    assert vk_test.verify_proof(8, proof, public)\n    print(\"Prover test with dummy verifier success\")\n\n\ndef prover_test(setup):\n    print(\"===prover_test===\")\n\n    print(\"Beginning prover test\")\n    program = Program([\"e public\", \"c <== a * b\", \"e <== c * d\"], 8)\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 proof\n\n\ndef verifier_test_unoptimized(setup, proof):\n    print(\"===verifier_test_unoptimized===\")\n\n    print(\"Beginning verifier test\")\n    program = Program([\"e public\", \"c <== a * b\", \"e <== c * d\"], 8)\n    public = [60]\n    vk = setup.verification_key(program.common_preprocessed_input())\n    assert vk.verify_proof_unoptimized(8, proof, public)\n    print(\"Verifier test success\")\n\n\ndef verifier_test_full(setup, proof):\n    print(\"===verifier_test_full===\")\n\n    print(\"Beginning verifier test\")\n    program = Program([\"e public\", \"c <== a * b\", \"e <== c * d\"], 8)\n    public = [60]\n    vk = setup.verification_key(program.common_preprocessed_input())\n    assert vk.verify_proof_unoptimized(8, proof, public)\n    assert vk.verify_proof(8, proof, public)\n    print(\"Verifier test success\")\n\n\ndef factorization_test(setup):\n    print(\"===factorization_test===\")\n\n    print(\"Beginning test: prove you know small integers that multiply to 91\")\n    program = Program.from_str(\n        \"\"\"n public\n        pb0 === pb0 * pb0\n        pb1 === pb1 * pb1\n        pb2 === pb2 * pb2\n        pb3 === pb3 * pb3\n        qb0 === qb0 * qb0\n        qb1 === qb1 * qb1\n        qb2 === qb2 * qb2\n        qb3 === qb3 * qb3\n        pb01 <== pb0 + 2 * pb1\n        pb012 <== pb01 + 4 * pb2\n        p <== pb012 + 8 * pb3\n        qb01 <== qb0 + 2 * qb1\n        qb012 <== qb01 + 4 * qb2\n        q <== qb012 + 8 * qb3\n        n <== p * q\"\"\",\n        16,\n    )\n    public = [91]\n    vk = setup.verification_key(program.common_preprocessed_input())\n    print(\"Generated verification key\")\n    assignments = program.fill_variable_assignments(\n        {\n            \"pb3\": 1,\n            \"pb2\": 1,\n            \"pb1\": 0,\n            \"pb0\": 1,\n            \"qb3\": 0,\n            \"qb2\": 1,\n            \"qb1\": 1,\n            \"qb0\": 1,\n        }\n    )\n    prover = Prover(setup, program)\n    proof = prover.prove(assignments)\n    print(\"Generated proof\")\n    assert vk.verify_proof(16, proof, public)\n    print(\"Factorization test success!\")\n\n\ndef output_proof_lang() -> str:\n    o = []\n    o.append(\"L0 public\")\n    o.append(\"M0 public\")\n    o.append(\"M64 public\")\n    o.append(\"R0 <== 0\")\n    for i in range(64):\n        for j, pos in enumerate((\"L\", \"M\", \"R\")):\n            f = {\"x\": i, \"r\": rc[i][j], \"p\": pos}\n            if i < 4 or i >= 60 or pos == \"L\":\n                o.append(\"{p}adj{x} <== {p}{x} + {r}\".format(**f))\n                o.append(\"{p}sq{x} <== {p}adj{x} * {p}adj{x}\".format(**f))\n                o.append(\"{p}qd{x} <== {p}sq{x} * {p}sq{x}\".format(**f))\n                o.append(\"{p}qn{x} <== {p}qd{x} * {p}adj{x}\".format(**f))\n            else:\n                o.append(\"{p}qn{x} <== {p}{x} + {r}\".format(**f))\n        for j, pos in enumerate((\"L\", \"M\", \"R\")):\n            f = {\"x\": i, \"p\": pos, \"m\": mds[j]}\n            o.append(\"{p}suma{x} <== Lqn{x} * {m}\".format(**f))\n            f = {\"x\": i, \"p\": pos, \"m\": mds[j + 1]}\n            o.append(\"{p}sumb{x} <== {p}suma{x} + Mqn{x} * {m}\".format(**f))\n            f = {\"x\": i, \"xp1\": i + 1, \"p\": pos, \"m\": mds[j + 2]}\n            o.append(\"{p}{xp1} <== {p}sumb{x} + Rqn{x} * {m}\".format(**f))\n    return \"\\n\".join(o)\n\n\ndef poseidon_test(setup):\n    print(\"===poseidon_test===\")\n\n    # PLONK-prove the correctness of a Poseidon execution. Note that this is\n    # a very suboptimal way to do it: an optimized implementation would use\n    # a custom PLONK gate to do a round in a single gate\n    expected_value = poseidon_hash(1, 2)\n    # Generate code for proof\n    program = Program.from_str(output_proof_lang(), 1024)\n    print(\"Generated code for Poseidon test\")\n    assignments = program.fill_variable_assignments({\"L0\": 1, \"M0\": 2})\n    vk = setup.verification_key(program.common_preprocessed_input())\n    print(\"Generated verification key\")\n    prover = Prover(setup, program)\n    proof = prover.prove(assignments)\n    print(\"Generated proof\")\n    assert vk.verify_proof(1024, proof, [1, 2, expected_value])\n    print(\"Verified proof!\")\n\n\nif __name__ == \"__main__\":\n    # Step 1: Pass setup test\n    setup_test()\n\n    setup = basic_test()\n\n    # Step 2: Pass prover test using verifier we provide (DO NOT READ TEST VERIFIER CODE)\n    prover_test_dummy_verifier(setup)\n\n    # Step 3: Pass verifier test using your own verifier\n    with open(\"test/proof.pickle\", \"rb\") as f:\n        proof = pickle.load(f)\n    verifier_test_unoptimized(setup, proof)\n    verifier_test_full(setup, proof)\n\n    # Step 4: Pass end-to-end tests for prover and verifier\n    ab_plus_a_test(setup)\n    one_public_input_test(setup)\n    proof = prover_test(setup)\n    verifier_test_full(setup, proof)\n    factorization_test(setup)\n    poseidon_test(setup)\n"
  },
  {
    "path": "transcript.py",
    "content": "from utils import Scalar\nfrom curve import G1Point\nfrom merlin import MerlinTranscript\nfrom py_ecc.secp256k1.secp256k1 import bytes_to_int\nfrom dataclasses import dataclass\n\n\n@dataclass\nclass Message1:\n    # [a(x)]₁ (commitment to left wire polynomial)\n    a_1: G1Point\n    # [b(x)]₁ (commitment to right wire polynomial)\n    b_1: G1Point\n    # [c(x)]₁ (commitment to output wire polynomial)\n    c_1: G1Point\n\n\n@dataclass\nclass Message2:\n    # [z(x)]₁ (commitment to permutation polynomial)\n    z_1: G1Point\n\n\n@dataclass\nclass Message3:\n    # [t_lo(x)]₁ (commitment to t_lo(X), the low chunk of the quotient polynomial t(X))\n    t_lo_1: G1Point\n    # [t_mid(x)]₁ (commitment to t_mid(X), the middle chunk of the quotient polynomial t(X))\n    t_mid_1: G1Point\n    # [t_hi(x)]₁ (commitment to t_hi(X), the high chunk of the quotient polynomial t(X))\n    t_hi_1: G1Point\n\n\n@dataclass\nclass Message4:\n    # Evaluation of a(X) at evaluation challenge ζ\n    a_eval: Scalar\n    # Evaluation of b(X) at evaluation challenge ζ\n    b_eval: Scalar\n    # Evaluation of c(X) at evaluation challenge ζ\n    c_eval: Scalar\n    # Evaluation of the first permutation polynomial S_σ1(X) at evaluation challenge ζ\n    s1_eval: Scalar\n    # Evaluation of the second permutation polynomial S_σ2(X) at evaluation challenge ζ\n    s2_eval: Scalar\n    # Evaluation of the shifted permutation polynomial z(X) at the shifted evaluation challenge ζω\n    z_shifted_eval: Scalar\n\n\n@dataclass\nclass Message5:\n    # [W_ζ(X)]₁ (commitment to the opening proof polynomial)\n    W_z_1: G1Point\n    # [W_ζω(X)]₁ (commitment to the opening proof polynomial)\n    W_zw_1: G1Point\n\n\nclass Transcript(MerlinTranscript):\n    def append(self, label: bytes, item: bytes) -> None:\n        self.append_message(label, item)\n\n    def append_scalar(self, label: bytes, item: Scalar):\n        self.append_message(label, item.n.to_bytes(32, \"big\"))\n\n    def append_point(self, label: bytes, item: G1Point):\n        self.append_message(label, item[0].n.to_bytes(32, \"big\"))\n        self.append_message(label, item[1].n.to_bytes(32, \"big\"))\n\n    def get_and_append_challenge(self, label: bytes) -> Scalar:\n        while True:\n            challenge_bytes = self.challenge_bytes(label, 255)\n            f = Scalar(bytes_to_int(challenge_bytes))\n            if f != Scalar.zero():  # Enforce challenge != 0\n                self.append(label, challenge_bytes)\n                return f\n\n    def round_1(self, message: Message1) -> tuple[Scalar, Scalar]:\n        self.append_point(b\"a_1\", message.a_1)\n        self.append_point(b\"b_1\", message.b_1)\n        self.append_point(b\"c_1\", message.c_1)\n\n        # The first two Fiat-Shamir challenges\n        beta = self.get_and_append_challenge(b\"beta\")\n        gamma = self.get_and_append_challenge(b\"gamma\")\n\n        return beta, gamma\n\n    def round_2(self, message: Message2) -> tuple[Scalar, Scalar]:\n        self.append_point(b\"z_1\", message.z_1)\n\n        alpha = self.get_and_append_challenge(b\"alpha\")\n        # This value could be anything, it just needs to be unpredictable. Lets us\n        # have evaluation forms at cosets to avoid zero evaluations, so we can\n        # divide polys without the 0/0 issue\n        fft_cofactor = self.get_and_append_challenge(b\"fft_cofactor\")\n\n        return alpha, fft_cofactor\n\n    def round_3(self, message: Message3) -> Scalar:\n        self.append_point(b\"t_lo_1\", message.t_lo_1)\n        self.append_point(b\"t_mid_1\", message.t_mid_1)\n        self.append_point(b\"t_hi_1\", message.t_hi_1)\n\n        zeta = self.get_and_append_challenge(b\"zeta\")\n        return zeta\n\n    def round_4(self, message: Message4) -> Scalar:\n        self.append_scalar(b\"a_eval\", message.a_eval)\n        self.append_scalar(b\"b_eval\", message.b_eval)\n        self.append_scalar(b\"c_eval\", message.c_eval)\n        self.append_scalar(b\"s1_eval\", message.s1_eval)\n        self.append_scalar(b\"s2_eval\", message.s2_eval)\n        self.append_scalar(b\"z_shifted_eval\", message.z_shifted_eval)\n\n        v = self.get_and_append_challenge(b\"v\")\n        return v\n\n    def round_5(self, message: Message5) -> Scalar:\n        self.append_point(b\"W_z_1\", message.W_z_1)\n        self.append_point(b\"W_zw_1\", message.W_zw_1)\n\n        u = self.get_and_append_challenge(b\"u\")\n        return u\n"
  },
  {
    "path": "utils.py",
    "content": "import py_ecc.bn128 as b\nfrom curve import Scalar\n\nf = b.FQ\nf2 = b.FQ2\n\nprimitive_root = 5\n\n# Extracts a point from JSON in zkrepl's format\ndef interpret_json_point(p):\n    if len(p) == 3 and isinstance(p[0], str) and p[2] == \"1\":\n        return (f(int(p[0])), f(int(p[1])))\n    elif len(p) == 3 and p == [\"0\", \"1\", \"0\"]:\n        return b.Z1\n    elif len(p) == 3 and isinstance(p[0], list) and p[2] == [\"1\", \"0\"]:\n        return (\n            f2([int(p[0][0]), int(p[0][1])]),\n            f2([int(p[1][0]), int(p[1][1])]),\n        )\n    elif len(p) == 3 and p == [[\"0\", \"0\"], [\"1\", \"0\"], [\"0\", \"0\"]]:\n        return b.Z2\n    raise Exception(\"cannot interpret that point: {}\".format(p))\n"
  },
  {
    "path": "verifier.py",
    "content": "import py_ecc.bn128 as b\nfrom utils import *\nfrom dataclasses import dataclass\nfrom curve import *\nfrom transcript import Transcript\nfrom poly import Polynomial, Basis\n\n\n@dataclass\nclass VerificationKey:\n    \"\"\"Verification key\"\"\"\n\n    # we set this to some power of 2 (so that we can FFT over it), that is at least the number of constraints we have (so we can Lagrange interpolate them)\n    group_order: int\n    # [q_M(x)]₁ (commitment to multiplication selector polynomial)\n    Qm: G1Point\n    # [q_L(x)]₁ (commitment to left selector polynomial)\n    Ql: G1Point\n    # [q_R(x)]₁ (commitment to right selector polynomial)\n    Qr: G1Point\n    # [q_O(x)]₁ (commitment to output selector polynomial)\n    Qo: G1Point\n    # [q_C(x)]₁ (commitment to constants selector polynomial)\n    Qc: G1Point\n    # [S_σ1(x)]₁ (commitment to the first permutation polynomial S_σ1(X))\n    S1: G1Point\n    # [S_σ2(x)]₁ (commitment to the second permutation polynomial S_σ2(X))\n    S2: G1Point\n    # [S_σ3(x)]₁ (commitment to the third permutation polynomial S_σ3(X))\n    S3: G1Point\n    # [x]₂ = xH, where H is a generator of G_2\n    X_2: G2Point\n    # nth root of unity (i.e. ω^1), where n is the program's group order.\n    w: Scalar\n\n    # More optimized version that tries hard to minimize pairings and\n    # elliptic curve multiplications, but at the cost of being harder\n    # to understand and mixing together a lot of the computations to\n    # efficiently batch them\n    def verify_proof(self, group_order: int, pf, public=[]) -> bool:\n        # 4. Compute challenges\n\n        # 5. Compute zero polynomial evaluation Z_H(ζ) = ζ^n - 1\n\n        # 6. Compute Lagrange polynomial evaluation L_0(ζ)\n\n        # 7. Compute public input polynomial evaluation PI(ζ).\n\n        # Compute the constant term of R. This is not literally the degree-0\n        # term of the R polynomial; rather, it's the portion of R that can\n        # be computed directly, without resorting to elliptic cutve commitments\n\n        # Compute D = (R - r0) + u * Z, and E and F\n\n        # Run one pairing check to verify the last two checks.\n        # What's going on here is a clever re-arrangement of terms to check\n        # the same equations that are being checked in the basic version,\n        # but in a way that minimizes the number of EC muls and even\n        # compressed the two pairings into one. The 2 pairings -> 1 pairing\n        # trick is basically to replace checking\n        #\n        # Y1 = A * (X - a) and Y2 = B * (X - b)\n        #\n        # with\n        #\n        # Y1 + A * a = A * X\n        # Y2 + B * b = B * X\n        #\n        # so at this point we can take a random linear combination of the two\n        # checks, and verify it with only one pairing.\n\n        return False\n\n    # Basic, easier-to-understand version of what's going on\n    def verify_proof_unoptimized(self, group_order: int, pf, public=[]) -> bool:\n        # 4. Compute challenges\n\n        # 5. Compute zero polynomial evaluation Z_H(ζ) = ζ^n - 1\n\n        # 6. Compute Lagrange polynomial evaluation L_0(ζ)\n\n        # 7. Compute public input polynomial evaluation PI(ζ).\n\n        # Recover the commitment to the linearization polynomial R,\n        # exactly the same as what was created by the prover\n\n        # Verify that R(z) = 0 and the prover-provided evaluations\n        # A(z), B(z), C(z), S1(z), S2(z) are all correct\n\n        # Verify that the provided value of Z(zeta*w) is correct\n\n        return False\n\n    # Compute challenges (should be same as those computed by prover)\n    def compute_challenges(\n        self, proof\n    ) -> tuple[Scalar, Scalar, Scalar, Scalar, Scalar, Scalar]:\n        transcript = Transcript(b\"plonk\")\n        beta, gamma = transcript.round_1(proof.msg_1)\n        alpha, _fft_cofactor = transcript.round_2(proof.msg_2)\n        zeta = transcript.round_3(proof.msg_3)\n        v = transcript.round_4(proof.msg_4)\n        u = transcript.round_5(proof.msg_5)\n\n        return beta, gamma, alpha, zeta, v, u\n"
  }
]