zk-snark-circuit-arithmetization-v2

Answer

Below is the complete implementation of the R1CS→QAP arithmetization using NTT over the BN254 scalar field (254-bit prime). The code handles field arithmetic, Lagrange interpolation, polynomial division, and full QAP proof generation.

/**
 * zk-SNARK Circuit Arithmetization v2
 * R1CS → QAP conversion using NTT (Number Theoretic Transform)
 * over the BN254 scalar field (254-bit prime): 
 *   p = 21888242871839275222246405745257275088548364400416034343698204186575808495617
 *
 * Given R1CS matrices A, B, C and witness w satisfying Aw ∘ Bw = Cw,
 * computes QAP polynomials via Lagrange interpolation over a power-of-two
 * FFT domain, evaluates them, and produces H(x) = (A_sum·B_sum - C_sum)/Z(x).
 */

'use strict';

// ─── 254-bit Prime Field ──────────────────────────────────────────
const P = 21888242871839275222246405745257275088548364400416034343698204186575808495617n;

function mod(x)      { return ((BigInt(x) % P) + P) % P; }
function modAdd(a,b) { return mod(mod(a) + mod(b)); }
function modSub(a,b) { return mod(mod(a) - mod(b)); }
function modMul(a,b) { return mod(mod(a) * mod(b)); }
function modNeg(a)   { return modSub(0n, a); }

function modPow(base, exp) {
  let r = 1n, b = mod(base), e = BigInt(exp);
  while (e > 0n) { if (e & 1n) r = modMul(r, b); b = modMul(b, b); e >>= 1n; }
  return r;
}

function modInv(a) { return modPow(a, P - 2n); }  // Fermat

// ─── NTT / Polynomial Operations ─────────────────────────────────
function findPrimitiveRoot(order) {
  // g = 5 is a generator of Z_p^*
  return modPow(5n, (P - 1n) / BigInt(order));
}

function bitReverse(n, bits) {
  let rev = 0;
  for (let i = 0; i < bits; i++) { rev = (rev << 1) | (n & 1); n >>= 1; }
  return rev;
}

/** In-place iterative NTT (Cooley-Tukey) */
function ntt(a, root, inverse = false) {
  const n = a.length;
  const bits = Math.log2(n);
  if (!Number.isInteger(bits)) throw new Error(`NTT length ${n} must be power of 2`);

  for (let i = 0; i < n; i++) { const j = bitReverse(i, bits); if (i < j) [a[i], a[j]] = [a[j], a[i]]; }

  for (let len = 2; len <= n; len <<= 1) {
    const wlen = modPow(root, BigInt(n / len));
    for (let i = 0; i < n; i += len) {
      let w = 1n;
      for (let j = 0; j < len / 2; j++) {
        const u = a[i + j], v = modMul(a[i + j + len / 2], w);
        a[i + j] = modAdd(u, v);
        a[i + j + len / 2] = modSub(u, v);
        w = modMul(w, wlen);
      }
    }
  }
  if (inverse) { const nInv = modInv(BigInt(n)); for (let i = 0; i < n; i++) a[i] = modMul(a[i], nInv); }
}

function polyMul(a, b) {
  const n = a.length, m = b.length;
  let size = 1; while (size < n + m - 1) size <<= 1;
  const root = findPrimitiveRoot(size);
  const fa = a.slice().concat(new Array(size - n).fill(0n));
  const fb = b.slice().concat(new Array(size - m).fill(0n));
  ntt(fa, root, false); ntt(fb, root, false);
  for (let i = 0; i < size; i++) fa[i] = modMul(fa[i], fb[i]);
  ntt(fa, modInv(root), true);
  return fa.slice(0, n + m - 1);
}

function polyEval(coeffs, x) {
  let r = 0n; for (let i = coeffs.length - 1; i >= 0; i--) r = modAdd(modMul(r, x), coeffs[i]); return r;
}

function polySub(a, b) {
  const ml = Math.max(a.length, b.length);
  const r = new Array(ml).fill(0n);
  for (let i = 0; i < a.length; i++) r[i] = modAdd(r[i], a[i]);
  for (let i = 0; i < b.length; i++) r[i] = modSub(r[i], b[i]);
  while (r.length > 1 && r[r.length - 1] === 0n) r.pop();
  return r;
}

function polyDiv(dividend, divisor) {
  const num = dividend.slice(), den = divisor.slice();
  if (den.length === 0 || (den.length === 1 && den[0] === 0n)) throw new Error('Division by zero');
  const q = new Array(Math.max(num.length - den.length + 1, 0)).fill(0n);
  const rem = num.slice();
  const dli = modInv(den[den.length - 1]);
  while (rem.length >= den.length) {
    const dd = rem.length - den.length, s = modMul(rem[rem.length - 1], dli);
    q[dd] = s;
    for (let i = 0; i < den.length; i++) rem[dd + i] = modSub(rem[dd + i], modMul(s, den[i]));
    while (rem.length > 0 && rem[rem.length - 1] === 0n) rem.pop();
  }
  while (q.length > 1 && q[q.length - 1] === 0n) q.pop();
  return { quotient: q, remainder: rem };
}

// ─── Lagrange Interpolation ──────────────────────────────────────
function lagrangeInterpolate(domain, evals) {
  const n = domain.length;
  if (n === 0) return [];
  if (n === 1) return [evals[0]];

  // Master polynomial M(x) = ∏ (x - d_i)
  let master = [1n];
  for (let i = 0; i < n; i++) master = polyMul(master, [modNeg(domain[i]), 1n]);

  // Barycentric weights: w_i = 1 / ∏_{j≠i} (d_i - d_j)
  const weights = [];
  for (let i = 0; i < n; i++) {
    let w = 1n;
    for (let j = 0; j < n; j++) if (i !== j) w = modMul(w, modSub(domain[i], domain[j]));
    weights.push(modInv(w));
  }

  // Result = Σ_i w_i·y_i·M(x)/(x-d_i) via synthetic division
  const result = new Array(n).fill(0n);
  for (let i = 0; i < n; i++) {
    const quot = syntheticDivide(master, domain[i]);
    const sc = modMul(evals[i], weights[i]);
    for (let j = 0; j < quot.length; j++) result[j] = modAdd(result[j], modMul(quot[j], sc));
  }
  while (result.length > 1 && result[result.length - 1] === 0n) result.pop();
  return result;
}

function syntheticDivide(coeffs, root) {
  const r = new Array(coeffs.length - 1).fill(0n);
  let c = 0n;
  for (let i = coeffs.length - 1; i > 0; i--) { c = modAdd(coeffs[i], modMul(c, root)); r[i - 1] = c; }
  return r;
}

// ─── R1CS → QAP Converter ────────────────────────────────────────
class R1CStoQAP {
  constructor(A, B, C, w) {
    this.n = A.length; this.m = A[0].length;
    this.A = A; this.B = B; this.C = C;
    this.w = w.map(BigInt);
    this._verifyR1CS();
  }

  _verifyR1CS() {
    const Aw = this._matVecMul(this.A, this.w);
    const Bw = this._matVecMul(this.B, this.w);
    const Cw = this._matVecMul(this.C, this.w);
    for (let i = 0; i < this.n; i++) {
      if (modMul(Aw[i], Bw[i]) !== Cw[i])
        throw new Error(`R1CS constraint ${i} violated`);
    }
  }

  _matVecMul(mat, vec) {
    return mat.map(r => { let s = 0n; for (let j = 0; j < r.length; j++) s = modAdd(s, modMul(r[j], vec[j])); return s; });
  }

  _buildDomain() {
    let size = 1; while (size < this.n) size <<= 1;
    const root = findPrimitiveRoot(size);
    const d = []; let x = 1n;
    for (let i = 0; i < size; i++) { d.push(x); x = modMul(x, root); }
    return d;
  }

  convert() {
    const domain = this._buildDomain();
    const dn = domain.slice(0, this.n);
    const aP = [], bP = [], cP = [];
    for (let v = 0; v < this.m; v++) {
      const ae = [], be = [], ce = [];
      for (let i = 0; i < this.n; i++) { ae.push(BigInt(this.A[i][v])); be.push(BigInt(this.B[i][v])); ce.push(BigInt(this.C[i][v])); }
      aP.push(lagrangeInterpolate(dn, ae));
      bP.push(lagrangeInterpolate(dn, be));
      cP.push(lagrangeInterpolate(dn, ce));
    }
    return { aPolys: aP, bPolys: bP, cPolys: cP, domain, domainN: dn };
  }

  computeTargetPoly(dn) {
    let z = [1n]; for (const r of dn) z = polyMul(z, [modNeg(r), 1n]);
    while (z.length > 1 && z[z.length - 1] === 0n) z.pop();
    return z;
  }

  _weightedSum(polys) {
    const ml = polys.reduce((m, p) => Math.max(m, p.length), 0);
    const r = new Array(ml).fill(0n);
    for (let i = 0; i < polys.length; i++) for (let j = 0; j < polys[i].length; j++) r[j] = modAdd(r[j], modMul(polys[i][j], this.w[i]));
    while (r.length > 1 && r[r.length - 1] === 0n) r.pop();
    return r;
  }

  prove() {
    const { aPolys, bPolys, cPolys, domain, domainN } = this.convert();
    const aSum = this._weightedSum(aPolys), bSum = this._weightedSum(bPolys), cSum = this._weightedSum(cPolys);
    const ab = polyMul(aSum, bSum), abMinusC = polySub(ab, cSum);
    const zPoly = this.computeTargetPoly(domainN);
    const { quotient: hPoly, remainder } = polyDiv(abMinusC, zPoly);

    const aSumE = sEval(aPolys, this.w, domain);
    const bSumE = sEval(bPolys, this.w, domain);
    const cSumE = sEval(cPolys, this.w, domain);
    const hEvals = domain.map(x => polyEval(hPoly, x));

    return { aPolys, bPolys, cPolys, hPoly, zPoly, aSumE, bSumE, cSumE, hEvals, domain, domainN, remainder, numConstraints: this.n, numVariables: this.m, domainSize: domain.length };
  }
}

function sEval(polys, weights, domain) {
  const r = new Array(domain.length).fill(0n);
  for (let j = 0; j < polys.length; j++) {
    const w = weights[j]; if (w === 0n) continue;
    for (let i = 0; i < domain.length; i++) r[i] = modAdd(r[i], modMul(polyEval(polys[j], domain[i]), w));
  }
  return r;
}

Key implementation details: - Field arithmetic: All operations are modulo the BN254 scalar field prime (254 bits), using BigInt and Fermat's little theorem for inverses. - NTT: Iterative Cooley-Tukey with bit-reversal permutation, supporting arbitrary power-of-two domains. - Polynomial multiplication: Via NTT-based convolution for O(n log n) performance. - Lagrange interpolation: Barycentric formula with synthetic division, avoiding explicit linear system solving. - Division: Standard polynomial long division with modular leading-coefficient inversion. - Domain: Powers of a primitive root of unity with order equal to the next power of two ≥ number of constraints.