plonk-to-ccs-plus.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "7c5beac9-7549-4a24-adcc-84c6bbf01324",
   "metadata": {},
   "source": [
    "This notebook is based on https://github.com/dmpierre/air-ccs/blob/main/sage/air-to-ccs.ipynb. This notebook implements\n",
    "- Plonk-to-CCS conversion algorithm, found in the [CCS paper](https://eprint.iacr.org/2023/552), Section 2.2.\n",
    "- A PoC of incorporating Halo2 lookup into CCS+.\n",
    "[The explanation of math behind the algorithm is here](https://hackmd.io/@pnyda/SkG0qaUuA)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96f62303-cf3e-401e-a4d5-675ba0be28ed",
   "metadata": {},
   "source": [
    "# Derive a CCS instance from a custom gate and constant columns with no respect to copy constraints and lookups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "bae1fefc",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Classes to represent a custom gate\n",
    "\n",
    "p = 41\n",
    "F = GF(p)\n",
    "R.<x> = PolynomialRing(F)\n",
    "# The modulus of the field doesn't really matter.\n",
    "# Here I use one from https://youtu.be/A3edAQDPnDY?si=K5E4EP3tTwyrKY5P&t=2544\n",
    "\n",
    "from dataclasses import dataclass\n",
    "from typing import Literal\n",
    "\n",
    "@dataclass(eq=True, frozen=True, order=True)\n",
    "class Column:\n",
    "    index: int\n",
    "    kind: Literal[\"advice\", \"constant\"]\n",
    "\n",
    "@dataclass(eq=True, frozen=True, order=True)\n",
    "class RelativeCellPointer:\n",
    "    column: Column\n",
    "    row_offset: int\n",
    "\n",
    "@dataclass(eq=True, frozen=True, order=True)\n",
    "class Monomial:\n",
    "    coefficient: F\n",
    "    variables: list[RelativeCellPointer]\n",
    "\n",
    "@dataclass(eq=True, frozen=True, order=True)\n",
    "class CustomGate:\n",
    "    monomials: list[Monomial]\n",
    "\n",
    "    def cell_pointers(self) -> list[RelativeCellPointer]:\n",
    "        # We sometimes refer to the same set of cells from multiple monomials.\n",
    "        # In that case a naive implementation would generate redundant M_j.\n",
    "        # To avoid it we take distinct set of RelativeCellPointers.\n",
    "        # and assign ID j for M_j by sorting it\n",
    "        return sorted({variable for monomial in self.monomials for variable in monomial.variables})\n",
    "\n",
    "    def t(self) -> list[int]:\n",
    "        return len(self.cell_pointers())\n",
    "\n",
    "    def c(self) -> list[int]:\n",
    "        return [monomial.coefficient for monomial in self.monomials]\n",
    "\n",
    "    def s(self) -> list[set[int]]:\n",
    "        return [{j := self.cell_pointers().index(variable) for variable in monomial.variables} for monomial in self.monomials]\n",
    "        \n",
    "    def m(self, constant_columns: list[list[F]]):\n",
    "        table_height = len(constant_columns[0])\n",
    "        matrices = []\n",
    "    \n",
    "        for cell_pointer in self.cell_pointers():\n",
    "            # Here I add 1 because 1 is always in Z\n",
    "            m_j = matrix(F, table_height, self.num_advice_columns() * table_height + 1)\n",
    "            for y in range(table_height):\n",
    "                if cell_pointer.column.kind == \"advice\":\n",
    "                    x = 1 + self.l() + cell_pointer.column.index * table_height + (y + cell_pointer.row_offset) % table_height\n",
    "                    m_j[y, x] = 1\n",
    "                elif cell_pointer.column.kind == \"constant\":\n",
    "                    # We assume here constant columns are followed by advice columns. This is a PoC code and it's not essential.\n",
    "                    constant_column_index = cell_pointer.column.index - self.num_advice_columns()\n",
    "                    # 0 here means that we multiply the rhs by the first element of Z which is always 1\n",
    "                    m_j[y, 0] = constant_columns[constant_column_index][(y + cell_pointer.row_offset) % table_height]\n",
    "                else:\n",
    "                    pass\n",
    "\n",
    "            matrices.append(m_j)\n",
    "    \n",
    "        return matrices\n",
    "\n",
    "    def l(self) -> int:\n",
    "        # Number of public i/o\n",
    "        # To make this PoC code simple we assume l=0\n",
    "        return 0\n",
    "\n",
    "    def num_advice_columns(self) -> int:\n",
    "        return len({cell_pointer.column.index for cell_pointer in self.cell_pointers() if cell_pointer.column.kind == \"advice\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "757ae84c-d379-4954-a9e3-740d9e28a8a8",
   "metadata": {},
   "source": [
    "# Check if a CCS instance is satisfied"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1a7096b1-0537-4631-87a6-06a277bb5fe0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# https://github.com/dmpierre/air-ccs/blob/main/sage/air-to-ccs.ipynb\n",
    "\n",
    "def ccs_is_satisfied(F, z, matrices, multisets, constants):\n",
    "    satisfied_instance_witness = vector(F, [0 for i in range(matrices[0].dimensions()[0])])\n",
    "    z_final = vector(F, [0 for i in range(matrices[0].dimensions()[0])])\n",
    "    for i, c in enumerate(constants):\n",
    "        multiset = multisets[i]\n",
    "        z_i = vector(F, [1 for i in range(matrices[0].dimensions()[0])])\n",
    "        for j in multiset:\n",
    "            z_i = z_i.pairwise_product(matrices[j] * z)\n",
    "        z_final += c * z_i\n",
    "    return z_final == satisfied_instance_witness"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34b952e8-cb4d-45fe-b661-97cab5d53cc2",
   "metadata": {},
   "source": [
    "# Original Plonk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "d010fad2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "c_plonk [1, 1, 1, 1, 1]\n",
      "s_plonk [{0, 1, 3}, {0, 4}, {1, 5}, {2, 6}, {7}]\n",
      "m_plonk\n",
      "\n",
      "[0 1 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 1 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 1 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 1 0 0 0 0 0 0 0 0]\n",
      "\n",
      "[0 0 0 0 0 1 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 1 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 1 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 1 0 0 0 0]\n",
      "\n",
      "[0 0 0 0 0 0 0 0 0 1 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 1 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 1 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 1]\n",
      "\n",
      "[1 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[1 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[2 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[2 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "\n",
      "[40  0  0  0  0  0  0  0  0  0  0  0  0]\n",
      "[40  0  0  0  0  0  0  0  0  0  0  0  0]\n",
      "[40  0  0  0  0  0  0  0  0  0  0  0  0]\n",
      "[ 0  0  0  0  0  0  0  0  0  0  0  0  0]\n",
      "\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "z_plonk (1, 0, 1, 1, 1, 0, 1, 2, 2, 0, 1, 6, 3)\n"
     ]
    }
   ],
   "source": [
    "# g = q_m * a * b + q_l * a + q_r * b + q_o * c + q_c\n",
    "\n",
    "## advice columns\n",
    "a = Column(0, \"advice\")\n",
    "b = Column(1, \"advice\")\n",
    "c = Column(2, \"advice\")\n",
    "\n",
    "## constant columns\n",
    "q_m = Column(3, \"constant\")\n",
    "q_l = Column(4, \"constant\")\n",
    "q_r = Column(5, \"constant\")\n",
    "q_o = Column(6, \"constant\")\n",
    "q_c = Column(7, \"constant\")\n",
    "\n",
    "a_cur = RelativeCellPointer(a, 0)\n",
    "b_cur = RelativeCellPointer(b, 0)\n",
    "c_cur = RelativeCellPointer(c, 0)\n",
    "q_m_cur = RelativeCellPointer(q_m, 0)\n",
    "q_l_cur = RelativeCellPointer(q_l, 0)\n",
    "q_r_cur = RelativeCellPointer(q_r, 0)\n",
    "q_o_cur = RelativeCellPointer(q_o, 0)\n",
    "q_c_cur = RelativeCellPointer(q_c, 0)\n",
    "\n",
    "g_plonk = CustomGate([\n",
    "    Monomial(F(1), [q_m_cur, a_cur, b_cur]),\n",
    "    Monomial(F(1), [q_l_cur, a_cur]),\n",
    "    Monomial(F(1), [q_r_cur, b_cur]),\n",
    "    Monomial(F(1), [q_o_cur, c_cur]),\n",
    "    Monomial(F(1), [q_c_cur]),\n",
    "])\n",
    "\n",
    "print(\"c_plonk\", g_plonk.c())\n",
    "print(\"s_plonk\", g_plonk.s())\n",
    "\n",
    "# Example values taken from https://hackmd.io/nQtquuk9QCGiB9EhUPgXvg#CCS-and-Plonkish-intro\n",
    "q_m_values = [F(1), F(1), F(0), F(0)]\n",
    "q_l_values = [F(0), F(0), F(2), F(0)]\n",
    "q_r_values = [F(0), F(0), F(2), F(0)]\n",
    "q_o_values = [F(-1), F(-1), F(-1), F(0)]\n",
    "q_c_values = [F(0), F(0), F(0), F(0)]\n",
    "constant_columns = [q_m_values, q_l_values, q_r_values, q_o_values, q_c_values]\n",
    "m_plonk = g_plonk.m(constant_columns)\n",
    "print(\"m_plonk\", *m_plonk, sep=\"\\n\\n\")\n",
    "\n",
    "a_values = [F(0), F(1), F(1), F(1)]\n",
    "b_values = [F(0), F(1), F(2), F(2)]\n",
    "c_values = [F(0), F(1), F(6), F(3)]\n",
    "\n",
    "z_plonk = vector(F, [F(1)] + a_values + b_values + c_values)\n",
    "print(\"z_plonk\", z_plonk)\n",
    "assert ccs_is_satisfied(F, z_plonk, m_plonk, g_plonk.s(), g_plonk.c())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a78c16f-f1b9-4bd5-82b6-7f8611f8c1e8",
   "metadata": {},
   "source": [
    "# Fibonacci"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "df5c7601",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "c_fib [1, 1, -1]\n",
      "s_fib [{0, 3}, {1, 3}, {2, 3}]\n",
      "m_fib\n",
      "\n",
      "[0 1 0 0 0]\n",
      "[0 0 1 0 0]\n",
      "[0 0 0 1 0]\n",
      "[0 0 0 0 1]\n",
      "\n",
      "[0 0 1 0 0]\n",
      "[0 0 0 1 0]\n",
      "[0 0 0 0 1]\n",
      "[0 1 0 0 0]\n",
      "\n",
      "[0 0 0 1 0]\n",
      "[0 0 0 0 1]\n",
      "[0 1 0 0 0]\n",
      "[0 0 1 0 0]\n",
      "\n",
      "[1 0 0 0 0]\n",
      "[1 0 0 0 0]\n",
      "[0 0 0 0 0]\n",
      "[0 0 0 0 0]\n"
     ]
    }
   ],
   "source": [
    "# enable(X) * (fib(X) + fib(Xω) - fib(Xω^2)) = 0\n",
    "fib = Column(0, \"advice\")\n",
    "enable = Column(1, \"constant\")\n",
    "\n",
    "fib_cur = RelativeCellPointer(fib, 0)\n",
    "fib_next = RelativeCellPointer(fib, 1)\n",
    "fib_next_next = RelativeCellPointer(fib, 2)\n",
    "enable_cur = RelativeCellPointer(enable, 0)\n",
    "\n",
    "g_fib = CustomGate([\n",
    "    Monomial(1, [enable_cur, fib_cur]),\n",
    "    Monomial(1, [enable_cur, fib_next]),\n",
    "    Monomial(-1, [enable_cur, fib_next_next]),\n",
    "])\n",
    "print(\"c_fib\", g_fib.c())\n",
    "print(\"s_fib\", g_fib.s())\n",
    "\n",
    "enable_values = [F(1), F(1), F(0), F(0)]\n",
    "m_fib = g_fib.m([enable_values])\n",
    "print(\"m_fib\", *m_fib, sep=\"\\n\\n\")\n",
    "\n",
    "fib_values = [F(1), F(1), F(2), F(3)]\n",
    "z_fib = vector(F, [F(1)] + fib_values)\n",
    "assert ccs_is_satisfied(F, z_fib, m_fib, g_fib.s(), g_fib.c())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23e7e231-dc09-44c9-b251-14a73e7200e3",
   "metadata": {},
   "source": [
    "# Apply copy constraints to a CCS instance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "4a05b087",
   "metadata": {},
   "outputs": [],
   "source": [
    "from dataclasses import dataclass\n",
    "\n",
    "@dataclass\n",
    "class AbsoluteCellPointer:\n",
    "    column: Column\n",
    "    row_index: int\n",
    "\n",
    "def z_index(cell: AbsoluteCellPointer, table_height: int) -> int: \n",
    "    return 1 + cell.column.index * table_height + cell.row_index\n",
    "\n",
    "# Step 1:\n",
    "# If M_j refers to multiple elements of Z that have to be the same,\n",
    "# update the reference in such a way that it all refers to the same element.\n",
    "\n",
    "# We also generate deduplication_map whose key is the original z_index and the value is z_index into which the original was deduplicated\n",
    "# This is needed to add lookup constraints afterwards\n",
    "\n",
    "def apply_copy_constraints(copy_constraints: list[list[AbsoluteCellPointer]], matrices: list[matrix], constant_columns: list[list[F]]) -> (list[matrix], list[int]):\n",
    "    table_height = len(constant_columns[0])\n",
    "    num_advice_columns = (matrices[0].ncols() - 1) // table_height\n",
    "\n",
    "    deduplication_map = [i for i in range(matrices[0].ncols())]\n",
    "    matrices = [copy(m_j) for m_j in matrices]\n",
    "    for m_j in matrices:\n",
    "        for row_index in range(m_j.nrows()):\n",
    "            for equal_cells in copy_constraints:\n",
    "                for cell in equal_cells[1:]:\n",
    "                    if m_j[row_index, z_index(cell, table_height)] != 0:\n",
    "                        # Deduplicate equal_cells[1:] into equal_cells[0]\n",
    "                        if equal_cells[0].column.kind == \"advice\":\n",
    "                            m_j[row_index, z_index(equal_cells[0], table_height)] = m_j[row_index, z_index(cell, table_height)]\n",
    "                            m_j[row_index, z_index(cell, table_height)] = 0\n",
    "                            deduplication_map[z_index(cell, table_height)] = z_index(equal_cells[0], table_height)\n",
    "                        elif equal_cells[0].column.kind == \"constant\":\n",
    "                            # We assume here constant columns are followed by advice columns. This is a PoC code and it's not essential.\n",
    "                            constant_column_index = equal_cells[0].column.index - num_advice_columns\n",
    "                            # 0 here means that we multiply the rhs by the first element of Z which is always 1\n",
    "                            m_j[row_index, 0] = constant_columns[constant_column_index][equal_cells[0].row_index]\n",
    "                            m_j[row_index, z_index(cell, table_height)] = 0\n",
    "                        else:\n",
    "                            pass\n",
    "\n",
    "    return matrices, deduplication_map\n",
    "\n",
    "\n",
    "# Step 2:\n",
    "# Remove elements of Z that is referenced by no M_j\n",
    "def reduce_witness_size(old_matrices: list[matrix], old_z: vector, dedup_map: list[int]) -> (list[matrix], vector):\n",
    "    used_z_indices = {x for m_j in old_matrices for y in range(m_j.nrows()) for x in range(m_j.ncols()) if m_j[y, x] != 0}\n",
    "    unused_z_indices = set(range(len(old_z))) - used_z_indices\n",
    "    print(\"unused_z_indices\", unused_z_indices)\n",
    "    print(\"used_z_indices\", used_z_indices)\n",
    "\n",
    "    # Update Z\n",
    "    new_z = list(old_z)\n",
    "    for unused_z_index in reversed(sorted(list(unused_z_indices))):\n",
    "        new_z.pop(unused_z_index)\n",
    "\n",
    "    new_matrices = []\n",
    "    for (j, old_m_j) in enumerate(old_matrices):\n",
    "        new_m_j = matrix(F, old_m_j.nrows(), len(used_z_indices))\n",
    "        for y in range(old_m_j.nrows()):\n",
    "            for x in range(old_m_j.ncols()):\n",
    "                offset = len([unused_z_index for unused_z_index in unused_z_indices if unused_z_index < x])\n",
    "                if old_m_j[y, x] != 0:\n",
    "                    # Update M_j\n",
    "                    new_m_j[y, x - offset] = old_m_j[y, x]\n",
    "\n",
    "                    # Update deduplication map\n",
    "                    for original_z_index in range(len(dedup_map)):\n",
    "                        if x == dedup_map[original_z_index]:\n",
    "                            dedup_map[original_z_index] -= offset\n",
    "        \n",
    "        new_matrices.append(new_m_j)\n",
    "\n",
    "    return (new_matrices, vector(F, new_z))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce4a2da1-30eb-4c84-bd78-e6c2b75c960a",
   "metadata": {},
   "source": [
    "# Apply copy constraints to m_plonk"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "ec9e000b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "unused_z_indices {5, 6, 7, 8, 9, 10, 12}\n",
      "used_z_indices {0, 1, 2, 3, 4, 11}\n",
      "original\n",
      "\n",
      "[0 1 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 1 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 1 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 1 0 0 0 0 0 0 0 0]\n",
      "\n",
      "[0 0 0 0 0 1 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 1 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 1 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 1 0 0 0 0]\n",
      "\n",
      "[0 0 0 0 0 0 0 0 0 1 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 1 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 1 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 1]\n",
      "\n",
      "[1 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[1 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[2 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[2 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "\n",
      "[40  0  0  0  0  0  0  0  0  0  0  0  0]\n",
      "[40  0  0  0  0  0  0  0  0  0  0  0  0]\n",
      "[40  0  0  0  0  0  0  0  0  0  0  0  0]\n",
      "[ 0  0  0  0  0  0  0  0  0  0  0  0  0]\n",
      "\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "deduplicated\n",
      "\n",
      "[0 1 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 1 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 1 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 1 0 0 0 0 0 0 0 0]\n",
      "\n",
      "[0 1 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 1 0 0 0 0 0 0 0 0 0 0]\n",
      "[2 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 1 0 0 0 0 0 0 0 0]\n",
      "\n",
      "[0 1 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 1 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 1 0]\n",
      "[0 0 0 0 1 0 0 0 0 0 0 0 0]\n",
      "\n",
      "[1 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[1 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[2 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[2 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "\n",
      "[40  0  0  0  0  0  0  0  0  0  0  0  0]\n",
      "[40  0  0  0  0  0  0  0  0  0  0  0  0]\n",
      "[40  0  0  0  0  0  0  0  0  0  0  0  0]\n",
      "[ 0  0  0  0  0  0  0  0  0  0  0  0  0]\n",
      "\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0 0 0 0 0 0 0 0]\n",
      "cleaned_up_m\n",
      "\n",
      "[0 1 0 0 0 0]\n",
      "[0 0 1 0 0 0]\n",
      "[0 0 0 1 0 0]\n",
      "[0 0 0 0 1 0]\n",
      "\n",
      "[0 1 0 0 0 0]\n",
      "[0 0 1 0 0 0]\n",
      "[2 0 0 0 0 0]\n",
      "[0 0 0 0 1 0]\n",
      "\n",
      "[0 1 0 0 0 0]\n",
      "[0 0 1 0 0 0]\n",
      "[0 0 0 0 0 1]\n",
      "[0 0 0 0 1 0]\n",
      "\n",
      "[1 0 0 0 0 0]\n",
      "[1 0 0 0 0 0]\n",
      "[0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0]\n",
      "\n",
      "[0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0]\n",
      "[2 0 0 0 0 0]\n",
      "[0 0 0 0 0 0]\n",
      "\n",
      "[0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0]\n",
      "[2 0 0 0 0 0]\n",
      "[0 0 0 0 0 0]\n",
      "\n",
      "[40  0  0  0  0  0]\n",
      "[40  0  0  0  0  0]\n",
      "[40  0  0  0  0  0]\n",
      "[ 0  0  0  0  0  0]\n",
      "\n",
      "[0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0]\n",
      "[0 0 0 0 0 0]\n",
      "original_z (1, 0, 1, 1, 1, 0, 1, 2, 2, 0, 1, 6, 3)\n",
      "cleaned_up_z (1, 0, 1, 1, 1, 6)\n"
     ]
    }
   ],
   "source": [
    "copy_constraints = [\n",
    "    [AbsoluteCellPointer(a, 0), AbsoluteCellPointer(b, 0), AbsoluteCellPointer(c, 0)],\n",
    "    [AbsoluteCellPointer(a, 1), AbsoluteCellPointer(b, 1), AbsoluteCellPointer(c, 1)],\n",
    "    [AbsoluteCellPointer(a, 3), AbsoluteCellPointer(b, 3), AbsoluteCellPointer(c, 3)],\n",
    "    [AbsoluteCellPointer(q_l, 2), AbsoluteCellPointer(b, 2)],  # This copy constraint was not in https://hackmd.io/nQtquuk9QCGiB9EhUPgXvg\n",
    "]\n",
    "(deduplicated_m, dedup_map) = apply_copy_constraints(copy_constraints, m_plonk, constant_columns)\n",
    "(cleaned_up_m, cleaned_up_z) = reduce_witness_size(deduplicated_m, z_plonk, dedup_map)\n",
    "print(\"original\", *m_plonk, sep=\"\\n\\n\")\n",
    "print(\"deduplicated\", *deduplicated_m, sep=\"\\n\\n\")\n",
    "print(\"cleaned_up_m\", *cleaned_up_m, sep=\"\\n\\n\")\n",
    "print(\"original_z\", z_plonk)\n",
    "print(\"cleaned_up_z\", cleaned_up_z)\n",
    "assert ccs_is_satisfied(F, cleaned_up_z, cleaned_up_m, g_plonk.s(), g_plonk.c())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "492145ed-1b7f-41e5-926c-2940a6b6c741",
   "metadata": {},
   "source": [
    "# ZeroTest\n",
    "or GateCheck? I don't know the industry standard way to call it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bc4412bf-cecf-49ef-98b4-d61b77448e7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The generator of 4th roots of unity\n",
    "omega = F(9)\n",
    "subgroup_order = 4\n",
    "assert omega^subgroup_order == F(1)\n",
    "\n",
    "z_in_zerotest = R.one()\n",
    "for i in range(4):\n",
    "    z_in_zerotest *= (x - omega^i)\n",
    "\n",
    "def prove_zero_test(f):\n",
    "    q = f // z_in_zerotest\n",
    "    return q\n",
    "\n",
    "def verify_zero_test(f, q) -> bool:\n",
    "    challenge = F.random_element()\n",
    "    return f(challenge) == q(challenge) * z_in_zerotest(challenge)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "547e6758-dfae-42bc-964a-d9c5bff4726b",
   "metadata": {},
   "source": [
    "# Grand Product Argument"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "289cd48c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# The generator of 4th roots of unity\n",
    "omega = F(9)\n",
    "assert omega^4 == F(1)\n",
    "\n",
    "# We want to succinctly check that seq1 is a reordered version of seq2\n",
    "seq1 = vector(F, [0, 1, 2, 3])\n",
    "seq2 = vector(F, [3, 2, 1, 0])\n",
    "\n",
    "f = R.one()\n",
    "for a_i in seq1:\n",
    "    f = f * (x + a_i)\n",
    "\n",
    "g = R.one()\n",
    "for b_i in seq2:\n",
    "    g = g * (x + b_i)\n",
    "\n",
    "gamma = F.random_element()\n",
    "assert f(gamma) == g(gamma)\n",
    "\n",
    "\n",
    "def interpolate_at_roots_of_unity(seq):\n",
    "    points = []\n",
    "    for (i, elem) in enumerate(seq):\n",
    "        points.append((omega^i, elem))\n",
    "\n",
    "    return R.lagrange_polynomial(points)\n",
    "\n",
    "\n",
    "beta = F.random_element()\n",
    "gamma = F.random_element()\n",
    "\n",
    "# Constrain a_prime to be reordered a\n",
    "# Constrain s_prime to be reordered s\n",
    "# With just one grand product argument, by using two random challenges\n",
    "def prove_grand_product(a_prime, s_prime, a, s):\n",
    "    z_trace = [F(1)]\n",
    "\n",
    "    for i in range(subgroup_order):\n",
    "        z_next = z_trace[-1]\n",
    "        z_next *= (a(omega^i) + beta) / (a_prime(omega^i) + beta)\n",
    "        z_next *= (s(omega^i) + gamma) / (s_prime(omega^i) + gamma)\n",
    "        z_trace.append(z_next)\n",
    "\n",
    "    # When generating the trace column Z, we have to do division\n",
    "    # Sometimes we encounter a situation where we divide 0, because we shift the dividened by a random challenge\n",
    "    # In that case this code panics with ZeroDivisionError\n",
    "    # How do provers in production avoid such cases?\n",
    "\n",
    "    z_trace[0] = z_trace.pop()\n",
    "\n",
    "    z = interpolate_at_roots_of_unity(z_trace)\n",
    "    constraint1 = z(omega * x) * (a_prime(x) + beta) * (s_prime(x) + gamma) - z(x) * (a(x) + beta) * (s(x) + gamma)\n",
    "    q1 = prove_zero_test(constraint1)\n",
    "    \n",
    "    l0 = R.lagrange_polynomial([(omega^i, 1 if i == 0 else 0) for i in range(subgroup_order)])\n",
    "    constraint2 = l0(x) * (F.one() - z(x))\n",
    "    q2 = prove_zero_test(constraint2)\n",
    "    \n",
    "    return z, q1, q2\n",
    "\n",
    "def verify_grand_product(a_prime, s_prime, a, s, z, q1, q2) -> bool:\n",
    "    constraint1 = z(omega * x) * (a_prime(x) + beta) * (s_prime(x) + gamma) - z(x) * (a(x) + beta) * (s(x) + gamma)\n",
    "    l0 = R.lagrange_polynomial([(omega^i, 1 if i == 0 else 0) for i in range(subgroup_order)])\n",
    "    constraint2 = l0(x) * (F.one() - z(x))\n",
    "    return verify_zero_test(constraint1, q1) and verify_zero_test(constraint2, q2)\n",
    "\n",
    "a = interpolate_at_roots_of_unity([3, 2, 1, 0])\n",
    "a_prime = interpolate_at_roots_of_unity([0, 1, 2, 3])\n",
    "s = interpolate_at_roots_of_unity([27, 24, 40, 8])\n",
    "s_prime = interpolate_at_roots_of_unity([8, 24, 27, 40])\n",
    "proof = prove_grand_product(a_prime, s_prime, a, s)\n",
    "assert verify_grand_product(a_prime, s_prime, a, s, *proof)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b8aeb5c-4abd-4c62-9c3b-7827192950e5",
   "metadata": {},
   "source": [
    "# \"Copy this or that\" argument"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "1492f7d2-f8ba-49ef-b144-6527f993f0da",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Constrains a_prime to be s_prime but with some elements duplicated, some elements missing\n",
    "def prove_copy_this_or_that(a_prime, s_prime):\n",
    "    constraint1 = (a_prime(x) - s_prime(x)) * (a_prime(x) - a_prime(x/omega))\n",
    "    l0 = R.lagrange_polynomial([(omega^i, 1 if i == 0 else 0) for i in range(subgroup_order)])\n",
    "    constraint2 = l0 * (a_prime(x) - s_prime(x))\n",
    "    return prove_zero_test(constraint1), prove_zero_test(constraint2)\n",
    "\n",
    "def verify_copy_this_or_that(a_prime, s_prime, q1, q2) -> bool:\n",
    "    constraint1 = (a_prime(x) - s_prime(x)) * (a_prime(x) - a_prime(x/omega))\n",
    "    l0 = R.lagrange_polynomial([(omega^i, 1 if i == 0 else 0) for i in range(subgroup_order)])\n",
    "    constraint2 = l0 * (a_prime(x) - s_prime(x))\n",
    "    return verify_zero_test(constraint1, q1) and verify_zero_test(constraint2, q2)\n",
    "    \n",
    "\n",
    "s_prime = interpolate_at_roots_of_unity([0, 1, 2, 3])\n",
    "a_prime = interpolate_at_roots_of_unity([0, 0, 2, 3])\n",
    "\n",
    "q1, q2 = prove_copy_this_or_that(a_prime, s_prime)\n",
    "assert verify_copy_this_or_that(a_prime, s_prime, q1, q2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0ee6347-298d-4ef1-9be6-6f2a3726a911",
   "metadata": {},
   "source": [
    "# Halo2 lookup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "7104fc37-5bd6-4aad-b90b-23b2eed3df6e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "a_prime [8, 24, 24, 40]\n",
      "s_prime [8, 24, 7, 40]\n",
      "a_prime [8, 24, 24, 40]\n",
      "s_prime [8, 24, 7, 40]\n",
      "a_prime [9, 24, 24, 40]\n",
      "s_prime [8, 24, 7, 40]\n",
      "a_prime [9, 24, 24, 40]\n",
      "s_prime [8, 24, 7, 40]\n"
     ]
    }
   ],
   "source": [
    "def sort_s(a_trace, s_trace) -> list[F]:\n",
    "    a_sorted = sorted(a_trace)\n",
    "    s_sorted = []\n",
    "\n",
    "    # We're gonna pop an element from either of these 2\n",
    "    s_not_in_a = [s_i for s_i in s_trace if s_i not in a_trace]\n",
    "    s_in_a = [s_i for s_i in s_trace if s_i in a_trace]\n",
    "    \n",
    "    for (i, a_i) in enumerate(a_sorted):\n",
    "        if a_i in s_in_a:  # Have we encountered this a_i before?\n",
    "            # If a_i is new value we haven't seen before\n",
    "            # Copy a_i to s_i\n",
    "            s_i = s_in_a.pop(s_in_a.index(a_i))\n",
    "        else:\n",
    "            # If a_i is duplicated\n",
    "            # Copy some random unused element to s_i\n",
    "            s_i = s_not_in_a.pop()\n",
    "\n",
    "        s_sorted.append(s_i)\n",
    "\n",
    "    assert len(s_not_in_a) == 0\n",
    "    assert len(s_in_a) == 0\n",
    "    print(\"a_prime\", a_sorted)\n",
    "    print(\"s_prime\", s_sorted)\n",
    "    \n",
    "    return s_sorted\n",
    "    \n",
    "\n",
    "def prove_halo2_lookup(a_trace, s_trace) -> bool:\n",
    "    a = interpolate_at_roots_of_unity(a_trace)\n",
    "    s = interpolate_at_roots_of_unity(s_trace)\n",
    "    a_prime = interpolate_at_roots_of_unity(sorted(a_trace))\n",
    "    s_prime = interpolate_at_roots_of_unity(sort_s(a_trace, s_trace))\n",
    "    \n",
    "    z_in_grand_product, q1, q2 = prove_grand_product(a_prime, s_prime, a, s)\n",
    "    q3, q4 = prove_copy_this_or_that(a_prime, s_prime)\n",
    "    return z_in_grand_product, q1, q2, q3, q4\n",
    "\n",
    "\n",
    "def verify_halo2_lookup(a_trace, s_trace, z_in_grand_product, q1, q2, q3, q4):\n",
    "    a = interpolate_at_roots_of_unity(a_trace)\n",
    "    s = interpolate_at_roots_of_unity(s_trace)\n",
    "    a_prime = interpolate_at_roots_of_unity(sorted(a_trace))\n",
    "    s_prime = interpolate_at_roots_of_unity(sort_s(a_trace, s_trace))\n",
    "    \n",
    "    return verify_grand_product(a_prime, s_prime, a, s, z_in_grand_product, q1, q2) and verify_copy_this_or_that(a_prime, s_prime, q3, q4)\n",
    "\n",
    "\n",
    "# Valid\n",
    "lookup_table = vector(F, [7, 24, 40, 8])\n",
    "must_be_in_table = vector(F, [24, 8, 40, 24])\n",
    "proof = prove_halo2_lookup(must_be_in_table, lookup_table)\n",
    "assert verify_halo2_lookup(must_be_in_table, lookup_table, *proof)\n",
    "\n",
    "# Invalid\n",
    "lookup_table = vector(F, [7, 24, 40, 8])\n",
    "must_be_in_table = vector(F, [24, 9, 40, 24])\n",
    "proof = prove_halo2_lookup(must_be_in_table, lookup_table)\n",
    "assert not verify_halo2_lookup(must_be_in_table, lookup_table, *proof)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "738b8426-1e71-43f1-844c-b182cdc063aa",
   "metadata": {},
   "source": [
    "# CCS+"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a9f57244-bf31-4981-b846-bff8320c198b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0, 1, 2, 3, 4, 1, 2, 7, 4, 1, 2, 5, 4]\n",
      "L [1, 2, 5, 4]\n",
      "Z_o [0, 1, 6, 1]\n",
      "a_prime [0, 1, 1, 6]\n",
      "s_prime [0, 1, 3, 6]\n",
      "a_prime [0, 1, 1, 6]\n",
      "s_prime [0, 1, 3, 6]\n"
     ]
    }
   ],
   "source": [
    "print(dedup_map)\n",
    "\n",
    "column_to_lookup = c_cur.column\n",
    "L = [dedup_map[z_index(AbsoluteCellPointer(column_to_lookup, row_index), subgroup_order)] for row_index in range(subgroup_order)]\n",
    "print(\"L\", L)\n",
    "Z_o = [cleaned_up_z[o] for o in L]\n",
    "print(\"Z_o\", Z_o)\n",
    "\n",
    "T = [6, 3, 0, 1]\n",
    "proof = prove_halo2_lookup(Z_o, T)\n",
    "assert verify_halo2_lookup(Z_o, T, *proof)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4fe8690-98c7-4d30-ae70-abc0e64a0f2a",
   "metadata": {},
   "source": [
    "# Realizations\n",
    "In this code I first generate M_j with no respect to copy constraints nor lookups, and then deduplicate it. This feels ugly and redundant. Maybe I should beforehand generate a map that keeps track of which cell in the table goes to which element of Z, and then generate M_j. That way lookup implementation would look cleaner as it becomes easy to tell which element of Z has to be in the lookup table. \n",
    "\n",
    "This code only supports one custom gate, but to support multiple custom gates, we can just increase the height of M_j to constrain multiple gates to evalutate to 0. That way we don't have to sample the random challenge gamma described in [Halo2 book](https://zcash.github.io/halo2/design/proving-system/vanishing.html) that was necessary to combine multiple custom gates into one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82e65f73-0b10-423a-8944-674df394edd8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "SageMath 10.3",
   "language": "sage",
   "name": "sagemath"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}