From 80d5537249d53ee7e205271173a3d2e392c29d5d Mon Sep 17 00:00:00 2001 From: kunxian xia Date: Tue, 7 Apr 2026 18:34:53 +0800 Subject: [PATCH 01/21] wip jagged_commit --- crates/mpcs/src/jagged.rs | 333 ++++++++++++++++++++++++++++++++++++++ crates/mpcs/src/lib.rs | 2 + 2 files changed, 335 insertions(+) create mode 100644 crates/mpcs/src/jagged.rs diff --git a/crates/mpcs/src/jagged.rs b/crates/mpcs/src/jagged.rs new file mode 100644 index 0000000..ddfd558 --- /dev/null +++ b/crates/mpcs/src/jagged.rs @@ -0,0 +1,333 @@ +//! Jagged PCS Commit Adaptor +//! +//! This module implements the commit protocol for the Jagged PCS as described in +//! the SP1 Jagged PCS paper () and Ceno issues #1272 / #1288. +//! +//! ## Overview +//! +//! The "Jagged PCS" reduces proof size by packing all trace polynomials from multiple +//! chips into a single "giga" multilinear polynomial `q'`: +//! +//! ```text +//! q' = bitrev(p_0) || bitrev(p_1) || ... || bitrev(p_N) +//! ``` +//! +//! where each `p_i` is a column polynomial extracted from the input trace matrices, and +//! `bitrev` is the suffix-to-prefix bit-reversal transformation. +//! +//! ## Suffix-to-Prefix Transformation +//! +//! The main sumcheck prover outputs evaluations `v_i = p_i(r[(n-s)..n])` — i.e., at the +//! **suffix** of the random challenge point. To make these evaluations compatible with the +//! jagged sumcheck (which operates on prefix-aligned polynomials), we apply a bit-reversal +//! permutation to each polynomial's evaluations: +//! +//! ```text +//! p_i'[j] = p_i[bitrev_s(j)] (for j in 0..2^s) +//! ``` +//! +//! After bit-reversal, `v_i = p_i(r[(n-s)..n]) = p_i'(reverse(r[(n-s)..n]))`. +//! +//! ## Cumulative Heights +//! +//! The cumulative height sequence `t` tracks the starting position of each polynomial in `q'`: +//! - `t[0] = 0` +//! - `t[i+1] = t[i] + h_i` where `h_i = 2^(num_vars of p_i)` is the number of evaluations +//! +//! Given a position `b` in `q'`, the verifier can locate the corresponding `(i, r)` pair via: +//! - `t[i-1] <= b < t[i]` +//! - `r = b - t[i-1]` +//! +//! The cumulative heights allow the verifier to succinctly evaluate the indicator function +//! `g(z_r, z_b, t[i-1], t[i])` needed for the jagged sumcheck. +//! +//! ## Commit Protocol +//! +//! 1. For each input matrix `M_k` (with `h_k` rows and `w_k` columns): +//! a. Extract each column as a polynomial with `h_k` evaluations. +//! b. Apply bit-reversal to the evaluations. +//! 2. Concatenate all bit-reversed polynomials: `cat = bitrev(p_0) || bitrev(p_1) || ...` +//! 3. Compute cumulative heights `t[i]`. +//! 4. Pad `cat` to the next power of two (required for MLE representation). +//! 5. Commit to the padded `cat` as a single-column matrix using the inner PCS. + +use std::iter::once; + +use crate::{Error, PolynomialCommitmentScheme}; +use ff_ext::ExtensionField; +use itertools::Itertools; +use p3::{ + matrix::{Matrix, bitrev::BitReversableMatrix}, + maybe_rayon::prelude::{ + IndexedParallelIterator, IntoParallelIterator, ParallelIterator, ParallelSliceMut, + }, +}; +use serde::{Deserialize, Serialize}; +use witness::{InstancePaddingStrategy, RowMajorMatrix}; + +/// Commitment to a jagged polynomial `q'`, together with all witness data needed +/// for opening proofs. +/// +/// Generic over the inner PCS `Pcs` so that any `PolynomialCommitmentScheme` can +/// serve as the underlying commitment engine. +pub struct JaggedCommitmentWithWitness> { + /// Commitment (with witness) to the "giga" polynomial `q'` via `Pcs`. + pub inner: Pcs::CommitmentWithWitness, + /// Cumulative height sequence `t`: + /// - `t[0] = 0` + /// - `t[i+1] = t[i] + poly_heights[i]` + /// - Length: `num_polys + 1` + pub cumulative_heights: Vec, + /// Number of evaluations `h_i = 2^(num_vars_i)` for each polynomial `p_i`. + /// Length: `num_polys`. + pub poly_heights: Vec, +} + +/// The pure commitment (without witness data) for a jagged polynomial `q'`. +/// This is what the verifier receives. +#[derive(Clone, Serialize, Deserialize)] +#[serde(bound(serialize = "", deserialize = ""))] +pub struct JaggedCommitment> { + /// Pure commitment to the underlying giga polynomial `q'`. + pub inner: Pcs::Commitment, + /// Cumulative height sequence `t` (verifier needs this to evaluate `f(b)`). + pub cumulative_heights: Vec, +} + +impl> JaggedCommitmentWithWitness { + /// Extract the pure commitment (without witness data). + pub fn to_commitment(&self) -> JaggedCommitment { + JaggedCommitment { + inner: Pcs::get_pure_commitment(&self.inner), + cumulative_heights: self.cumulative_heights.clone(), + } + } + + /// Total number of polynomials packed into `q'`. + pub fn num_polys(&self) -> usize { + self.poly_heights.len() + } + + /// Total number of evaluations in the *unpadded* concatenated polynomial + /// (= `t[num_polys]` = `cumulative_heights.last()`). + pub fn total_evaluations(&self) -> usize { + self.cumulative_heights.last().copied().unwrap_or(0) + } +} + +/// Commit to a sequence of row-major matrices using the Jagged PCS scheme. +/// +/// This function implements the commit phase described in Ceno issue #1288: +/// 1. For each matrix, bit-reverse its rows (suffix-to-prefix transformation). +/// 2. Transpose the bit-reversed matrix (row-major → column-major), so each +/// column polynomial occupies a contiguous region in memory. +/// 3. Concatenate all column polynomials: `q' = col_0 || col_1 || ...` +/// 4. Compute the cumulative height sequence `t`. +/// 5. Commit to `q'` as a single-column matrix using `Pcs::batch_commit`. +/// +/// # Arguments +/// * `pp` — Prover parameters for `Pcs`. +/// * `rmms` — Non-empty sequence of row-major matrices. This function uses each matrix's height exactly as given. +/// +/// # Errors +/// Returns `Error::InvalidPcsParam` if `rmms` is empty or all matrices are empty. +/// Any error from the inner `Pcs::batch_commit` is propagated as-is. +pub fn jagged_commit>( + pp: &Pcs::ProverParam, + rmms: Vec>, +) -> Result, Error> { + if rmms.is_empty() { + return Err(Error::InvalidPcsParam( + "jagged_commit: cannot commit to empty sequence of matrices".to_string(), + )); + } + + // --- Step 1: Compute cumulative heights --- + let mut poly_heights: Vec = Vec::new(); + for rmm in &rmms { + let num_rows = rmm.height(); + let num_cols = rmm.width(); + + if num_rows == 0 { + return Err(Error::InvalidPcsParam( + "jagged_commit: matrix has zero rows".to_string(), + )); + } + if num_cols == 0 { + return Err(Error::InvalidPcsParam( + "jagged_commit: matrix has zero columns".to_string(), + )); + } + for _ in 0..num_cols { + poly_heights.push(num_rows); + } + } + if poly_heights.is_empty() { + return Err(Error::InvalidPcsParam( + "jagged_commit: no polynomials found in input matrices".to_string(), + )); + } + // t[0] = 0, t[i+1] = t[i] + poly_heights[i] + let cumulative_heights = poly_heights + .iter() + .chain(once(&0)) + .scan(0usize, |acc, &h| { + let current = *acc; + *acc += h; + Some(current) + }) + .collect::>(); + + // --- Steps 2 & 3: Bit-reverse rows, transpose, and write to concatenated --- + let total_size = cumulative_heights.last().copied().unwrap(); + let mut concatenated: Vec = Vec::with_capacity(total_size); + // Safety: every element in `concatenated[0..total_size]` is fully written + // by the transpose loop below before it is read. + unsafe { concatenated.set_len(total_size) }; + + // `poly_idx` tracks which poly (column index in cumulative_heights) is the + // first polynomial of the current matrix. + let mut poly_idx = 0; + for rmm in &rmms { + // Step 2: Bit-reverse the rows (suffix-to-prefix transformation). + // br.values[i * n_cols + j] = original[bitrev(i)][j] + let br = rmm.as_view().bit_reverse_rows().to_row_major_matrix(); + + let n_cols = br.width(); + let n_rows = br.height(); + let n_cells = n_cols * n_rows; + + // The start position in `concatenated` for this matrix's block of polynomials. + let start = cumulative_heights[poly_idx]; + + // Step 3: Transpose — write each column j of `br` (= one polynomial) + // into its corresponding contiguous slice in `concatenated`. + (0..n_cols) + .into_par_iter() + .zip(concatenated[start..start + n_cells].par_chunks_mut(n_rows)) + .for_each(|(j, chunk)| { + br.values + .iter() + .skip(j) + .step_by(n_cols) + .zip_eq(chunk.iter_mut()) + .for_each(|(v, out)| *out = *v); + }); + + poly_idx += n_cols; + } + + // --- Step 4: Commit via the inner PCS --- + // q' is committed as a single-column matrix with height = total_size. + let giga_rmm = RowMajorMatrix::::new_by_values( + concatenated, + 1, // width = 1 (single polynomial q') + InstancePaddingStrategy::Default, + ); + + let inner = Pcs::batch_commit(pp, vec![giga_rmm])?; + + Ok(JaggedCommitmentWithWitness { + inner, + cumulative_heights, + poly_heights, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + basefold::{Basefold, BasefoldRSParams}, + test_util::setup_pcs, + }; + use ff_ext::GoldilocksExt2; + use p3::{field::FieldAlgebra, goldilocks::Goldilocks}; + + type F = Goldilocks; + type E = GoldilocksExt2; + type Pcs = Basefold; + + fn make_rmm(num_rows: usize, num_cols: usize) -> RowMajorMatrix { + let values: Vec = (0..num_rows * num_cols) + .map(|i| F::from_canonical_u64(i as u64 + 1)) + .collect(); + RowMajorMatrix::::new_by_values(values, num_cols, InstancePaddingStrategy::Default) + } + + #[test] + fn test_cumulative_heights_single_matrix() { + // 4x2 matrix → 2 polynomials with 4 evaluations each → cumulative = [0, 4, 8] + let rmm = make_rmm(4, 2); + let num_rows = rmm.height(); + let num_cols = rmm.width(); + let mut poly_heights = Vec::new(); + for _ in 0..num_cols { + poly_heights.push(num_rows); + } + let mut ch = vec![0usize]; + for &h in &poly_heights { + ch.push(ch.last().unwrap() + h); + } + assert_eq!(poly_heights, vec![4, 4]); + assert_eq!(ch, vec![0, 4, 8]); + } + + #[test] + fn test_cumulative_heights_multiple_matrices() { + // 4x1 + 8x2 → heights [4, 8, 8] → cumulative [0, 4, 12, 20] + let m1 = make_rmm(4, 1); + let m2 = make_rmm(8, 2); + let mut poly_heights: Vec = Vec::new(); + for rmm in &[m1, m2] { + for _ in 0..rmm.width() { + poly_heights.push(rmm.height()); + } + } + let mut ch = vec![0usize]; + for &h in &poly_heights { + ch.push(ch.last().unwrap() + h); + } + assert_eq!(poly_heights, vec![4, 8, 8]); + assert_eq!(ch, vec![0, 4, 12, 20]); + } + + #[test] + fn test_jagged_commit_smoke() { + // Two matrices: 4x1 and 4x2 → 3 polynomials, total 12 evals, padded to 16 + let (pp, _vp) = setup_pcs::(4); + let m1 = make_rmm(4, 1); + let m2 = make_rmm(4, 2); + + let comm = jagged_commit::(&pp, vec![m1, m2]).expect("commit should succeed"); + + assert_eq!(comm.num_polys(), 3); + assert_eq!(comm.poly_heights, vec![4, 4, 4]); + assert_eq!(comm.cumulative_heights, vec![0, 4, 8, 12]); + assert_eq!(comm.total_evaluations(), 12); + + let pure = comm.to_commitment(); + assert_eq!(pure.cumulative_heights, vec![0, 4, 8, 12]); + } + + #[test] + fn test_jagged_commit_single_poly() { + // 8x1 matrix → 1 polynomial, 8 evals, no padding needed + let (pp, _vp) = setup_pcs::(3); + let m = make_rmm(8, 1); + + let comm = jagged_commit::(&pp, vec![m]).expect("commit should succeed"); + + assert_eq!(comm.num_polys(), 1); + assert_eq!(comm.poly_heights, vec![8]); + assert_eq!(comm.cumulative_heights, vec![0, 8]); + assert_eq!(comm.total_evaluations(), 8); + } + + #[test] + fn test_jagged_commit_empty_error() { + let (pp, _vp) = setup_pcs::(4); + let result = jagged_commit::(&pp, vec![]); + assert!(matches!(result, Err(Error::InvalidPcsParam(_)))); + } +} diff --git a/crates/mpcs/src/lib.rs b/crates/mpcs/src/lib.rs index a8defea..be6380d 100644 --- a/crates/mpcs/src/lib.rs +++ b/crates/mpcs/src/lib.rs @@ -267,6 +267,8 @@ pub use basefold::{ Basefold, BasefoldCommitment, BasefoldCommitmentWithWitness, BasefoldDefault, BasefoldParams, BasefoldRSParams, BasefoldSpec, EncodingScheme, RSCode, RSCodeDefaultSpec, }; +pub mod jagged; +pub use jagged::{JaggedCommitment, JaggedCommitmentWithWitness, jagged_commit}; #[cfg(feature = "whir")] extern crate whir as whir_external; #[cfg(feature = "whir")] From 91aa59f2c9408cb0080b954f214dc97e15b8e6cf Mon Sep 17 00:00:00 2001 From: xkx Date: Tue, 14 Apr 2026 21:35:35 +0800 Subject: [PATCH 02/21] feat: jagged sumcheck (#32) * claude code impl plan * implement the jagged_sumcheck using time-space tradeoff sumcheck prover algorithm * remove debugging codes * ref to the original paper * add jagged sumcheck bench * #32 parallelize (#34) * par wip * check f(z) * g(z) matches expected evaluation * fix clippy * fix clippy: add #[cfg(test)] to test-only method and fix unused import/variable warnings Agent-Logs-Url: https://github.com/scroll-tech/gkr-backend/sessions/4a61316e-cf28-47b8-a43e-fb6ab432701e Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> * refactor test * apply functional programming style * avoid unnecessary BaseField-to-ExtensionField conversion in q_evals access Use E * BaseField multiplication directly instead of converting q_evals elements to extension field with .into() first. Co-Authored-By: Claude Opus 4.6 * replace col_row binary search with incremental ColRowIter Add ColRowIter that does one binary search at construction and O(1) per step, replacing per-element binary searches in build_m_table, bind_and_materialize, compute_claimed_sum, and final_evaluations_slow. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 * extend jagged sumcheck benchmark to cover n=25..31 * switch jagged sumcheck benchmark to BabyBearExt4 Co-Authored-By: Claude Opus 4.6 * remove jagged sumcheck plan doc Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 --- Cargo.lock | 13 + crates/mpcs/Cargo.toml | 5 + crates/mpcs/benches/jagged_sumcheck.rs | 65 +++ crates/mpcs/src/jagged.rs | 622 ++++++++++++++++++++++++- crates/mpcs/src/lib.rs | 5 +- 5 files changed, 707 insertions(+), 3 deletions(-) create mode 100644 crates/mpcs/benches/jagged_sumcheck.rs diff --git a/Cargo.lock b/Cargo.lock index 303a17c..554bd40 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -838,6 +838,7 @@ dependencies = [ "serde", "sumcheck", "tracing", + "tracing-forest", "tracing-subscriber", "transcript", "whir", @@ -1919,6 +1920,18 @@ dependencies = [ "valuable", ] +[[package]] +name = "tracing-forest" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee40835db14ddd1e3ba414292272eddde9dad04d3d4b65509656414d1c42592f" +dependencies = [ + "smallvec", + "thiserror", + "tracing", + "tracing-subscriber", +] + [[package]] name = "tracing-log" version = "0.2.0" diff --git a/crates/mpcs/Cargo.toml b/crates/mpcs/Cargo.toml index b54e03d..0962118 100644 --- a/crates/mpcs/Cargo.toml +++ b/crates/mpcs/Cargo.toml @@ -30,6 +30,7 @@ witness.workspace = true [dev-dependencies] criterion.workspace = true +tracing-forest = "0.1" [features] nightly-features = ["ff_ext/nightly-features"] @@ -54,3 +55,7 @@ name = "interpolate" harness = false name = "whir" required-features = ["whir"] + +[[bench]] +harness = false +name = "jagged_sumcheck" diff --git a/crates/mpcs/benches/jagged_sumcheck.rs b/crates/mpcs/benches/jagged_sumcheck.rs new file mode 100644 index 0000000..1390bfc --- /dev/null +++ b/crates/mpcs/benches/jagged_sumcheck.rs @@ -0,0 +1,65 @@ +use criterion::*; +use ff_ext::{BabyBearExt4, ExtensionField, FieldFrom, FromUniformBytes}; +use mpcs::{JaggedSumcheckInput, jagged_sumcheck_prove}; +use multilinear_extensions::virtual_poly::build_eq_x_r_vec; +use rand::thread_rng; +use transcript::BasicTranscript; + +type E = BabyBearExt4; +type F = ::BaseField; + +const NUM_SAMPLES: usize = 10; + +fn bench_jagged_sumcheck(c: &mut Criterion) { + let mut group = c.benchmark_group("jagged_sumcheck"); + group.sample_size(NUM_SAMPLES); + + // (num_giga_vars, num_polys, poly_height_log2) + let configs: Vec<(usize, usize, usize)> = (25..=31) + .map(|n| { + let s = 21usize; + let num_polys = 1usize << (n - s); + (n, num_polys, s) + }) + .collect(); + + for (num_giga_vars, num_polys, s) in configs { + let poly_height = 1usize << s; + let total_evals = num_polys * poly_height; + + let mut rng = thread_rng(); + + let q_evals: Vec = (0..total_evals) + .map(|i| F::from_v((i as u64 * 13 + 7) % (1 << 30))) + .collect(); + + let cumulative_heights: Vec = (0..=num_polys).map(|i| i * poly_height).collect(); + + let z_row: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); + let z_col_vars = (num_polys as f64).log2().ceil() as usize; + let z_col: Vec = (0..z_col_vars).map(|_| E::random(&mut rng)).collect(); + + let input = JaggedSumcheckInput { + q_evals: &q_evals, + num_giga_vars, + cumulative_heights: &cumulative_heights, + eq_row: build_eq_x_r_vec(&z_row), + eq_col: build_eq_x_r_vec(&z_col), + }; + + group.bench_function( + BenchmarkId::new("prove", format!("n={}", num_giga_vars)), + |b| { + b.iter(|| { + let mut transcript = BasicTranscript::::new(b"jagged_bench"); + jagged_sumcheck_prove(black_box(&input), &mut transcript) + }) + }, + ); + } + + group.finish(); +} + +criterion_group!(benches, bench_jagged_sumcheck); +criterion_main!(benches); diff --git a/crates/mpcs/src/jagged.rs b/crates/mpcs/src/jagged.rs index ddfd558..0d044cd 100644 --- a/crates/mpcs/src/jagged.rs +++ b/crates/mpcs/src/jagged.rs @@ -56,13 +56,22 @@ use std::iter::once; use crate::{Error, PolynomialCommitmentScheme}; use ff_ext::ExtensionField; use itertools::Itertools; +use multilinear_extensions::{ + mle::MultilinearExtension, util::max_usable_threads, virtual_poly::build_eq_x_r_vec, +}; use p3::{ matrix::{Matrix, bitrev::BitReversableMatrix}, maybe_rayon::prelude::{ - IndexedParallelIterator, IntoParallelIterator, ParallelIterator, ParallelSliceMut, + IndexedParallelIterator, IntoParallelIterator, ParallelIterator, ParallelSlice, + ParallelSliceMut, }, }; use serde::{Deserialize, Serialize}; +use sumcheck::{ + macros::{entered_span, exit_span}, + structs::{IOPProof, IOPProverMessage, IOPProverState}, +}; +use transcript::Transcript; use witness::{InstancePaddingStrategy, RowMajorMatrix}; /// Commitment to a jagged polynomial `q'`, together with all witness data needed @@ -183,7 +192,10 @@ pub fn jagged_commit>( let mut concatenated: Vec = Vec::with_capacity(total_size); // Safety: every element in `concatenated[0..total_size]` is fully written // by the transpose loop below before it is read. - unsafe { concatenated.set_len(total_size) }; + #[allow(clippy::uninit_vec)] + unsafe { + concatenated.set_len(total_size) + }; // `poly_idx` tracks which poly (column index in cumulative_heights) is the // first polynomial of the current matrix. @@ -234,6 +246,427 @@ pub fn jagged_commit>( }) } +// --------------------------------------------------------------------------- +// Jagged Sumcheck Prover (M-table streaming algorithm) +// --------------------------------------------------------------------------- + +// Streaming sumcheck prover using the M-table algorithm from +// "Time-Space Trade-Offs for Sumcheck" (eprint 2025/1473), Section 4. + +/// Number of streaming rounds before switching to standard sumcheck. +/// Determined by the epoch schedule: 1 + 2 + 4 + 8 = 15. +#[allow(dead_code)] +const STREAMING_ROUNDS: usize = 15; +/// Epoch sizes used in the streaming phase: j' = 1, 2, 4, 8. +const EPOCH_SIZES: [usize; 4] = [1, 2, 4, 8]; + +/// All inputs needed for the jagged sumcheck. +pub struct JaggedSumcheckInput<'a, E: ExtensionField> { + /// Giga polynomial evaluations (concatenated, bit-reversed). + pub q_evals: &'a [E::BaseField], + /// n = log2(padded_total_size). + pub num_giga_vars: usize, + /// Cumulative height sequence t[j], length num_polys + 1. + pub cumulative_heights: &'a [usize], + /// Precomputed eq table for the row evaluation point: `build_eq_x_r_vec(z_row)`. + pub eq_row: Vec, + /// Precomputed eq table for the column challenge point: `build_eq_x_r_vec(z_col)`. + pub eq_col: Vec, +} + +/// Iterator that yields `(col, row)` pairs for consecutive giga indices. +/// Uses one binary search at construction, then O(1) per step. +struct ColRowIter<'a> { + cumulative_heights: &'a [usize], + col: usize, + row: usize, + num_polys: usize, +} + +impl<'a> Iterator for ColRowIter<'a> { + type Item = (usize, usize); + + fn next(&mut self) -> Option<(usize, usize)> { + if self.col >= self.num_polys { + return None; + } + let result = (self.col, self.row); + self.row += 1; + let poly_height = self.cumulative_heights[self.col + 1] - self.cumulative_heights[self.col]; + if self.row >= poly_height { + self.row = 0; + self.col += 1; + } + Some(result) + } +} + +impl<'a, E: ExtensionField> JaggedSumcheckInput<'a, E> { + fn total_evaluations(&self) -> usize { + *self.cumulative_heights.last().unwrap_or(&0) + } + + /// Return an iterator yielding `(col, row)` for consecutive giga indices + /// starting from `start`. One binary search at construction, O(1) per step. + fn col_row_iter(&self, start: usize) -> ColRowIter<'_> { + let num_polys = self.cumulative_heights.len() - 1; + if start >= self.total_evaluations() { + return ColRowIter { + cumulative_heights: self.cumulative_heights, + col: num_polys, + row: 0, + num_polys, + }; + } + let j = self.cumulative_heights.partition_point(|&t| t <= start) - 1; + ColRowIter { + cumulative_heights: self.cumulative_heights, + col: j, + row: start - self.cumulative_heights[j], + num_polys, + } + } + + /// Brute-force computation of sum_b q'(b) * f(b). + /// O(2^n) time — only for debugging and tests. + #[cfg(test)] + fn compute_claimed_sum(&self) -> E { + self.q_evals[..self.total_evaluations()] + .iter() + .zip(self.col_row_iter(0)) + .fold(E::ZERO, |sum, (&q, (col, row))| { + sum + (self.eq_row[row] * self.eq_col[col]) * q + }) + } + + /// Brute-force MLE evaluation of q'(rb) and f(rb) at the given point. + /// O(2^n) time and memory — only for debugging and tests. + /// Returns `(q_at_point, f_at_point)`. + #[cfg(test)] + fn final_evaluations_slow(&self, point: &[E]) -> (E, E) { + let n = self.num_giga_vars; + let total_evals = self.total_evaluations(); + + // Build q' MLE (padded with zeros) and evaluate at point. + let mut q_padded: Vec = self.q_evals.to_vec(); + q_padded.resize(1 << n, Default::default()); + let q_mle = MultilinearExtension::from_evaluations_vec(n, q_padded); + let q_at_point = q_mle.evaluate(point); + + // Build f MLE from eq tables and evaluate at point. + let mut f_evals = vec![E::ZERO; 1 << n]; + for (f_eval, (col, row)) in f_evals[..total_evals].iter_mut().zip(self.col_row_iter(0)) { + *f_eval = self.eq_row[row] * self.eq_col[col]; + } + let f_mle = MultilinearExtension::from_evaluations_ext_vec(n, f_evals); + let f_at_point = f_mle.evaluate(point); + + (q_at_point, f_at_point) + } +} + +/// Run the full jagged sumcheck: streaming phase (rounds 1..K) + standard phase (rounds K+1..n). +/// +/// Returns the proof and the full list of challenges (r_1, ..., r_n). +pub fn jagged_sumcheck_prove( + input: &JaggedSumcheckInput, + transcript: &mut impl Transcript, +) -> (IOPProof, Vec) { + let n = input.num_giga_vars; + let max_degree: usize = 2; + + let mut challenges: Vec = Vec::with_capacity(n); + let mut proof_messages: Vec> = Vec::with_capacity(n); + + // Write transcript header (must match verifier's expectations). + transcript.append_message(&n.to_le_bytes()); + transcript.append_message(&max_degree.to_le_bytes()); + + // --- Streaming phase: epochs j' = 1, 2, 4, 8 --- + for &epoch_size in &EPOCH_SIZES { + // Epoch j' handles rounds j'..2j'-1. Skip if all rounds are done. + if epoch_size > n { + break; + } + + // Build M-table for this epoch. + let span = entered_span!("build_m_table", epoch = epoch_size); + let m_table = build_m_table(input, &challenges, epoch_size); + exit_span!(span); + + // Extract rounds j = epoch_size .. min(2*epoch_size - 1, n) + let span = entered_span!("compute_rounds_from_m", epoch = epoch_size); + for j in epoch_size..(2 * epoch_size).min(n + 1) { + let d = j - epoch_size; // intra-epoch offset + let intra_challenges = challenges[epoch_size - 1..epoch_size - 1 + d].to_vec(); + + let [_h0, h1, h2] = compute_round_from_m(&m_table, epoch_size, &intra_challenges); + + // Append [h(1), h(2)] to transcript and sample challenge. + transcript.append_field_element_ext(&h1); + transcript.append_field_element_ext(&h2); + let challenge = transcript + .sample_and_append_challenge(b"Internal round") + .elements; + + proof_messages.push(IOPProverMessage { + evaluations: vec![h1, h2], + }); + challenges.push(challenge); + } + exit_span!(span); + } + + // --- Phase 2: Bind and materialize, then standard sumcheck --- + let k = challenges.len(); // actual number of streaming rounds completed + if k < n { + let span = entered_span!("bind_and_materialize"); + let (q_bound, f_bound) = bind_and_materialize(input, &challenges); + exit_span!(span); + + let remaining_vars = n - k; + let q_mle = MultilinearExtension::from_evaluations_ext_vec(remaining_vars, q_bound); + let f_mle = MultilinearExtension::from_evaluations_ext_vec(remaining_vars, f_bound); + + // Use VirtualPolynomial + round-by-round proving (no extra transcript header). + use multilinear_extensions::virtual_poly::VirtualPolynomial; + use std::sync::Arc; + let q_arc = Arc::new(q_mle); + let f_arc = Arc::new(f_mle); + let vp = VirtualPolynomial::new_from_product(vec![q_arc, f_arc], E::ONE); + + let span = entered_span!("standard_sumcheck", rounds = remaining_vars); + let mut prover_state = + IOPProverState::prover_init_with_extrapolation_aux(true, vp, None, None); + let mut challenge = None; + for _ in 0..remaining_vars { + let prover_msg = + IOPProverState::prove_round_and_update_state(&mut prover_state, &challenge); + prover_msg + .evaluations + .iter() + .for_each(|e| transcript.append_field_element_ext(e)); + challenge = Some(transcript.sample_and_append_challenge(b"Internal round")); + challenges.push(challenge.unwrap().elements); + proof_messages.push(prover_msg); + } + exit_span!(span); + } + + ( + IOPProof { + proofs: proof_messages, + }, + challenges, + ) +} + +/// Build M-table for epoch j'. +/// +/// M[beta1 * 2^{j'} + beta2] = sum_b Q_bound(beta1, b) * F_bound(beta2, b) +/// +/// where: +/// - Q_bound(beta, b) = sum_{a in {0,1}^{j'-1}} eq(R, a) * q'[a || beta || b] +/// - F_bound(beta, b) = sum_{a in {0,1}^{j'-1}} eq(R, a) * f[a || beta || b] +fn build_m_table( + input: &JaggedSumcheckInput, + challenges: &[E], // R_{j'} = (r_1, ..., r_{j'-1}) + epoch_size: usize, // j' +) -> Vec { + let n = input.num_giga_vars; + let bound_vars = epoch_size - 1; // j' - 1 + + let eq_r = if bound_vars > 0 { + build_eq_x_r_vec(&challenges[..bound_vars]) + } else { + vec![E::ONE] + }; + + let beta_count = 1usize << epoch_size; // 2^{j'} + let a_count = 1usize << bound_vars; // 2^{j'-1} + let chunk_size = a_count * beta_count; // 2^{2j'-1} + // When n < 2j'-1, all variables fit in a single chunk (no b dimension). + let n_chunks = 1usize << n.saturating_sub(2 * epoch_size - 1); // 2^{max(0, n - 2j' + 1)} + + let m_size = beta_count * beta_count; // 2^{2j'} + + // Step 1: Each thread processes a batch of b-chunks, producing a local M-table. + let span = entered_span!( + "streaming_pass", + n_chunks = n_chunks, + beta_count = beta_count + ); + let indices: Vec = (0..n_chunks).collect(); + let n_threads = max_usable_threads(); + let batch_size = (n_chunks / n_threads).max(1); + let partial_tables: Vec> = indices + .par_chunks(batch_size) + .map(|batch| { + let mut local_m = vec![E::ZERO; m_size]; + let mut q_bound = vec![E::ZERO; beta_count]; + let mut f_bound = vec![E::ZERO; beta_count]; + + for &b_idx in batch { + let chunk_start = b_idx * chunk_size; + q_bound + .iter_mut() + .zip(f_bound.iter_mut()) + .enumerate() + .for_each(|(beta, (q_b, f_b))| { + let base = chunk_start + beta * a_count; + let (q_acc, f_acc) = eq_r + .iter() + .zip(input.q_evals.get(base..).unwrap_or(&[])) + .zip(input.col_row_iter(base)) + .fold( + (E::ZERO, E::ZERO), + |(q_acc, f_acc), ((&eq_r_a, &q), (col, row))| { + ( + q_acc + eq_r_a * q, + f_acc + eq_r_a * (input.eq_row[row] * input.eq_col[col]), + ) + }, + ); + *q_b = q_acc; + *f_b = f_acc; + }); + + // Outer product accumulation into local M-table. + for b1 in 0..beta_count { + if q_bound[b1] == E::ZERO { + continue; + } + for b2 in 0..beta_count { + local_m[b1 * beta_count + b2] += q_bound[b1] * f_bound[b2]; + } + } + } + local_m + }) + .collect(); + exit_span!(span); + + // Step 2: Sum partial M-tables in parallel, each thread handles a slice of cells. + let span = entered_span!("reduce_partial_tables", n_partials = partial_tables.len()); + let n_partials = partial_tables.len(); + if n_partials == 0 { + exit_span!(span); + return vec![E::ZERO; m_size]; + } + let mut m_table = partial_tables[0].clone(); + let cell_batch = (m_size / n_threads).max(1); + m_table + .par_chunks_mut(cell_batch) + .enumerate() + .for_each(|(ci, cells)| { + let start = ci * cell_batch; + for partial in &partial_tables[1..] { + for (j, cell) in cells.iter_mut().enumerate() { + *cell += partial[start + j]; + } + } + }); + exit_span!(span); + m_table +} + +/// Extract round univariate h_j(x) from M-table. +/// +/// For round j in epoch j', d = j - j' intra-epoch challenges have been collected. +/// Returns [h(0), h(1), h(2)]. +fn compute_round_from_m( + m_table: &[E], + epoch_size: usize, // j' + intra_challenges: &[E], // r_{j'}, ..., r_{j-1} (d elements) +) -> [E; 3] { + let d = intra_challenges.len(); + let beta_count = 1usize << epoch_size; + let pad_bits = epoch_size - d - 1; // number of "future" bits to sum over + let pad_count = 1usize << pad_bits; + + let eq_intra = if d > 0 { + build_eq_x_r_vec(intra_challenges) + } else { + vec![E::ONE] + }; + + let a_count = 1usize << d; // 2^d + + let mut h = [E::ZERO; 3]; + + for a in 0..a_count { + for c in 0..a_count { + let eq_weight = eq_intra[a] * eq_intra[c]; + if eq_weight == E::ZERO { + continue; + } + + // Sum over all pad bit assignments (same pad for both beta1/beta2 + // since they correspond to the same physical "future" variables). + for p in 0..pad_count { + // beta = a_bits || x_bit || pad_bits (little-endian) + // beta_val = a + x_bit * 2^d + pad * 2^{d+1} + let base1 = a + (p << (d + 1)); + let base2 = c + (p << (d + 1)); + let b1_0 = base1; // x=0 + let b1_1 = base1 + (1 << d); // x=1 + let b2_0 = base2; + let b2_1 = base2 + (1 << d); + + let m00 = m_table[b1_0 * beta_count + b2_0]; + let m10 = m_table[b1_1 * beta_count + b2_0]; + let m01 = m_table[b1_0 * beta_count + b2_1]; + let m11 = m_table[b1_1 * beta_count + b2_1]; + + h[0] += eq_weight * m00; + h[1] += eq_weight * m11; + // h(2) via bilinear: (1-2)^2*M00 + 2(1-2)*M10 + (1-2)*2*M01 + 4*M11 + h[2] += eq_weight * (m00 - m10.double() - m01.double() + m11.double().double()); + } + } + } + + h +} + +/// Bind first K variables and materialize reduced q' and f as extension-field vectors. +/// +/// q_bound[idx] = sum_{a in {0,1}^K} eq(R, a) * q'[a + idx * 2^K] +/// f_bound[idx] = sum_{a in {0,1}^K} eq(R, a) * f[a + idx * 2^K] +fn bind_and_materialize( + input: &JaggedSumcheckInput, + challenges: &[E], // R_K = (r_1, ..., r_K) +) -> (Vec, Vec) { + let n = input.num_giga_vars; + let k = challenges.len(); + let remaining_size = 1usize << (n - k); + let a_count = 1usize << k; + + let eq_r = build_eq_x_r_vec(challenges); + + // Each output index is independent — parallelize over idx. + let results: Vec<(E, E)> = (0..remaining_size) + .into_par_iter() + .map(|idx| { + let base = idx * a_count; + eq_r.iter() + .zip(input.q_evals.get(base..).unwrap_or(&[])) + .zip(input.col_row_iter(base)) + .fold( + (E::ZERO, E::ZERO), + |(q_acc, f_acc), ((&eq_r_a, &q), (col, row))| { + ( + q_acc + eq_r_a * q, + f_acc + eq_r_a * (input.eq_row[row] * input.eq_col[col]), + ) + }, + ) + }) + .collect(); + + results.into_iter().unzip() +} + #[cfg(test)] mod tests { use super::*; @@ -330,4 +763,189 @@ mod tests { let result = jagged_commit::(&pp, vec![]); assert!(matches!(result, Err(Error::InvalidPcsParam(_)))); } + + // --- Sumcheck tests --- + + use multilinear_extensions::virtual_poly::{VPAuxInfo, build_eq_x_r_vec}; + use rand::thread_rng; + use std::marker::PhantomData; + use sumcheck::structs::IOPVerifierState; + use transcript::basic::BasicTranscript; + + #[test] + fn test_jagged_sumcheck_small() { + use ff_ext::FromUniformBytes; + + let mut rng = thread_rng(); + + // 3 polynomials of height 4 (s=2), total 12 evals, padded to 16 (n=4). + let num_polys = 3usize; + let poly_height = 4usize; + let s = 2; // log2(poly_height) + let total_evals = num_polys * poly_height; + let num_giga_vars = 4; // ceil(log2(12)) = 4, 2^4 = 16 + + let q_evals: Vec = (0..total_evals) + .map(|i| F::from_canonical_u64(i as u64 + 1)) + .collect(); + + let cumulative_heights: Vec = (0..=num_polys).map(|i| i * poly_height).collect(); + + let z_row: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); + // z_col needs ceil(log2(num_polys)) = 2 variables + let z_col: Vec = (0..2).map(|_| E::random(&mut rng)).collect(); + + let input = JaggedSumcheckInput { + q_evals: &q_evals, + num_giga_vars, + cumulative_heights: &cumulative_heights, + eq_row: build_eq_x_r_vec(&z_row), + eq_col: build_eq_x_r_vec(&z_col), + }; + + let claimed_sum = input.compute_claimed_sum(); + + let mut transcript = BasicTranscript::::new(b"jagged_sumcheck_test"); + let (proof, challenges) = jagged_sumcheck_prove(&input, &mut transcript); + + assert_eq!(proof.proofs.len(), num_giga_vars); + assert_eq!(challenges.len(), num_giga_vars); + + // Verify using the standard sumcheck verifier. + let mut transcript_v = BasicTranscript::::new(b"jagged_sumcheck_test"); + let aux_info = VPAuxInfo { + max_degree: 2, + max_num_variables: num_giga_vars, + phantom: PhantomData::, + }; + let subclaim = + IOPVerifierState::::verify(claimed_sum, &proof, &aux_info, &mut transcript_v); + + // The subclaim point should match our challenges. + for (sc, ch) in subclaim.point.iter().zip(challenges.iter()) { + assert_eq!(sc.elements, *ch); + } + + // Verify the final evaluation: q'(point) * f(point) == expected_evaluation + let (q_at_point, f_at_point) = input.final_evaluations_slow(&challenges); + assert_eq!( + q_at_point * f_at_point, + subclaim.expected_evaluation, + "final evaluation mismatch" + ); + } + + #[test] + fn test_jagged_sumcheck_all_epochs() { + // n=16: exercises all 4 epochs (j'=1,2,4,8) + 1 round of standard sumcheck. + use ff_ext::FromUniformBytes; + + let mut rng = thread_rng(); + + let num_polys = 8usize; + let poly_height = 1 << 13; // 8192, s=13 + let s = 13; + let total_evals = num_polys * poly_height; // 65536 + let num_giga_vars = 16; // 2^16 = 65536 + + let q_evals: Vec = (0..total_evals) + .map(|i| F::from_canonical_u64((i as u64 * 7 + 3) % (1 << 20))) + .collect(); + + let cumulative_heights: Vec = (0..=num_polys).map(|i| i * poly_height).collect(); + + let z_row: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); + let z_col: Vec = (0..3).map(|_| E::random(&mut rng)).collect(); // ceil(log2(8))=3 + + let input = JaggedSumcheckInput { + q_evals: &q_evals, + num_giga_vars, + cumulative_heights: &cumulative_heights, + eq_row: build_eq_x_r_vec(&z_row), + eq_col: build_eq_x_r_vec(&z_col), + }; + + let claimed_sum = input.compute_claimed_sum(); + + let mut transcript = BasicTranscript::::new(b"jagged_test_16"); + let (proof, challenges) = jagged_sumcheck_prove(&input, &mut transcript); + + assert_eq!(proof.proofs.len(), num_giga_vars); + assert_eq!(challenges.len(), num_giga_vars); + + let mut transcript_v = BasicTranscript::::new(b"jagged_test_16"); + let aux_info = VPAuxInfo { + max_degree: 2, + max_num_variables: num_giga_vars, + phantom: PhantomData::, + }; + let subclaim = + IOPVerifierState::::verify(claimed_sum, &proof, &aux_info, &mut transcript_v); + + for (sc, ch) in subclaim.point.iter().zip(challenges.iter()) { + assert_eq!(sc.elements, *ch); + } + } + + #[test] + fn test_jagged_sumcheck_n25() { + // n=25: 2^25 = 33M evaluations. Exercises all epochs + 10 rounds of standard sumcheck. + use ff_ext::FromUniformBytes; + + tracing_forest::init(); + + let mut rng = thread_rng(); + + let num_polys = 1 << 10; // 1024 polynomials + let poly_height = 1 << 15; // 32768 each, s=15 + let s = 15; + let total_evals = num_polys * poly_height; // 2^25 = 33554432 + let num_giga_vars = 25; + + let q_evals: Vec = (0..total_evals) + .map(|i| F::from_canonical_u64((i as u64 * 13 + 7) % (1 << 30))) + .collect(); + + let cumulative_heights: Vec = (0..=num_polys).map(|i| i * poly_height).collect(); + + let z_row: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); + let z_col: Vec = (0..10).map(|_| E::random(&mut rng)).collect(); // ceil(log2(1024))=10 + + let input = JaggedSumcheckInput { + q_evals: &q_evals, + num_giga_vars, + cumulative_heights: &cumulative_heights, + eq_row: build_eq_x_r_vec(&z_row), + eq_col: build_eq_x_r_vec(&z_col), + }; + + let claimed_sum = input.compute_claimed_sum(); + + let mut transcript = BasicTranscript::::new(b"jagged_test_25"); + let (proof, challenges) = jagged_sumcheck_prove(&input, &mut transcript); + + assert_eq!(proof.proofs.len(), num_giga_vars); + assert_eq!(challenges.len(), num_giga_vars); + + let mut transcript_v = BasicTranscript::::new(b"jagged_test_25"); + let aux_info = VPAuxInfo { + max_degree: 2, + max_num_variables: num_giga_vars, + phantom: PhantomData::, + }; + let subclaim = + IOPVerifierState::::verify(claimed_sum, &proof, &aux_info, &mut transcript_v); + + for (sc, ch) in subclaim.point.iter().zip(challenges.iter()) { + assert_eq!(sc.elements, *ch); + } + + // Verify the final evaluation: q'(point) * f(point) == expected_evaluation + let (q_at_point, f_at_point) = input.final_evaluations_slow(&challenges); + assert_eq!( + q_at_point * f_at_point, + subclaim.expected_evaluation, + "final evaluation mismatch" + ); + } } diff --git a/crates/mpcs/src/lib.rs b/crates/mpcs/src/lib.rs index be6380d..4a4bb04 100644 --- a/crates/mpcs/src/lib.rs +++ b/crates/mpcs/src/lib.rs @@ -268,7 +268,10 @@ pub use basefold::{ BasefoldRSParams, BasefoldSpec, EncodingScheme, RSCode, RSCodeDefaultSpec, }; pub mod jagged; -pub use jagged::{JaggedCommitment, JaggedCommitmentWithWitness, jagged_commit}; +pub use jagged::{ + JaggedCommitment, JaggedCommitmentWithWitness, JaggedSumcheckInput, jagged_commit, + jagged_sumcheck_prove, +}; #[cfg(feature = "whir")] extern crate whir as whir_external; #[cfg(feature = "whir")] From adc2636905ae9ceb663dc6ab76b01c13a52d1e3f Mon Sep 17 00:00:00 2001 From: Copilot <198982749+Copilot@users.noreply.github.com> Date: Tue, 14 Apr 2026 22:00:39 +0800 Subject: [PATCH 03/21] feat: make streaming epoch schedule configurable in jagged_sumcheck_prove (#38) * claude code impl plan * implement the jagged_sumcheck using time-space tradeoff sumcheck prover algorithm * remove debugging codes * ref to the original paper * add jagged sumcheck bench * #32 parallelize (#34) * par wip * check f(z) * g(z) matches expected evaluation * fix clippy * fix clippy: add #[cfg(test)] to test-only method and fix unused import/variable warnings Agent-Logs-Url: https://github.com/scroll-tech/gkr-backend/sessions/4a61316e-cf28-47b8-a43e-fb6ab432701e Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> * refactor test * apply functional programming style * avoid unnecessary BaseField-to-ExtensionField conversion in q_evals access Use E * BaseField multiplication directly instead of converting q_evals elements to extension field with .into() first. Co-Authored-By: Claude Opus 4.6 * replace col_row binary search with incremental ColRowIter Add ColRowIter that does one binary search at construction and O(1) per step, replacing per-element binary searches in build_m_table, bind_and_materialize, compute_claimed_sum, and final_evaluations_slow. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 * extend jagged sumcheck benchmark to cover n=25..31 * switch jagged sumcheck benchmark to BabyBearExt4 Co-Authored-By: Claude Opus 4.6 * remove jagged sumcheck plan doc Co-Authored-By: Claude Opus 4.6 * Initial plan * Make EPOCH_SIZES configurable in jagged_sumcheck_prove via optional parameter Agent-Logs-Url: https://github.com/scroll-tech/gkr-backend/sessions/b18805c7-6e15-44c0-ab03-a1905068d964 Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> * Fix doc spacing and add debug_assert for epoch_sizes validation Agent-Logs-Url: https://github.com/scroll-tech/gkr-backend/sessions/b18805c7-6e15-44c0-ab03-a1905068d964 Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> --------- Co-authored-by: kunxian xia Co-authored-by: copilot-swe-agent[bot] <198982749+Copilot@users.noreply.github.com> Co-authored-by: hero78119 <3962077+hero78119@users.noreply.github.com> Co-authored-by: Claude Opus 4.6 --- crates/mpcs/benches/jagged_sumcheck.rs | 2 +- crates/mpcs/src/jagged.rs | 28 ++++++++++++++++---------- 2 files changed, 18 insertions(+), 12 deletions(-) diff --git a/crates/mpcs/benches/jagged_sumcheck.rs b/crates/mpcs/benches/jagged_sumcheck.rs index 1390bfc..4dbaed0 100644 --- a/crates/mpcs/benches/jagged_sumcheck.rs +++ b/crates/mpcs/benches/jagged_sumcheck.rs @@ -52,7 +52,7 @@ fn bench_jagged_sumcheck(c: &mut Criterion) { |b| { b.iter(|| { let mut transcript = BasicTranscript::::new(b"jagged_bench"); - jagged_sumcheck_prove(black_box(&input), &mut transcript) + jagged_sumcheck_prove(black_box(&input), &mut transcript, None) }) }, ); diff --git a/crates/mpcs/src/jagged.rs b/crates/mpcs/src/jagged.rs index 0d044cd..97e3b30 100644 --- a/crates/mpcs/src/jagged.rs +++ b/crates/mpcs/src/jagged.rs @@ -253,12 +253,10 @@ pub fn jagged_commit>( // Streaming sumcheck prover using the M-table algorithm from // "Time-Space Trade-Offs for Sumcheck" (eprint 2025/1473), Section 4. -/// Number of streaming rounds before switching to standard sumcheck. -/// Determined by the epoch schedule: 1 + 2 + 4 + 8 = 15. -#[allow(dead_code)] -const STREAMING_ROUNDS: usize = 15; -/// Epoch sizes used in the streaming phase: j' = 1, 2, 4, 8. -const EPOCH_SIZES: [usize; 4] = [1, 2, 4, 8]; +/// Default log2 of the maximum epoch size for the streaming phase. +/// Epoch sizes are `[2^0, 2^1, ..., 2^LOG2_MAX_EPOCH]` = `[1, 2, 4, 8]`, +/// covering 1 + 2 + 4 + 8 = 15 streaming rounds before switching to standard sumcheck. +const LOG2_MAX_EPOCH: u32 = 3; /// All inputs needed for the jagged sumcheck. pub struct JaggedSumcheckInput<'a, E: ExtensionField> { @@ -367,14 +365,22 @@ impl<'a, E: ExtensionField> JaggedSumcheckInput<'a, E> { /// Run the full jagged sumcheck: streaming phase (rounds 1..K) + standard phase (rounds K+1..n). /// +/// `log2_max_epoch` controls the streaming phase epoch schedule. Epoch sizes are +/// `[1, 2, 4, ..., 2^k]` where `k = log2_max_epoch`. Pass `None` to use the default +/// `LOG2_MAX_EPOCH = 3`, giving epoch sizes `[1, 2, 4, 8]` and 15 streaming rounds. +/// /// Returns the proof and the full list of challenges (r_1, ..., r_n). pub fn jagged_sumcheck_prove( input: &JaggedSumcheckInput, transcript: &mut impl Transcript, + log2_max_epoch: Option, ) -> (IOPProof, Vec) { let n = input.num_giga_vars; let max_degree: usize = 2; + let k = log2_max_epoch.unwrap_or(LOG2_MAX_EPOCH); + let epoch_sizes: Vec = (0..=k).map(|i| 1usize << i).collect(); + let mut challenges: Vec = Vec::with_capacity(n); let mut proof_messages: Vec> = Vec::with_capacity(n); @@ -382,8 +388,8 @@ pub fn jagged_sumcheck_prove( transcript.append_message(&n.to_le_bytes()); transcript.append_message(&max_degree.to_le_bytes()); - // --- Streaming phase: epochs j' = 1, 2, 4, 8 --- - for &epoch_size in &EPOCH_SIZES { + // --- Streaming phase: epochs j' = 1, 2, 4, ..., 2^k --- + for &epoch_size in &epoch_sizes { // Epoch j' handles rounds j'..2j'-1. Skip if all rounds are done. if epoch_size > n { break; @@ -806,7 +812,7 @@ mod tests { let claimed_sum = input.compute_claimed_sum(); let mut transcript = BasicTranscript::::new(b"jagged_sumcheck_test"); - let (proof, challenges) = jagged_sumcheck_prove(&input, &mut transcript); + let (proof, challenges) = jagged_sumcheck_prove(&input, &mut transcript, None); assert_eq!(proof.proofs.len(), num_giga_vars); assert_eq!(challenges.len(), num_giga_vars); @@ -868,7 +874,7 @@ mod tests { let claimed_sum = input.compute_claimed_sum(); let mut transcript = BasicTranscript::::new(b"jagged_test_16"); - let (proof, challenges) = jagged_sumcheck_prove(&input, &mut transcript); + let (proof, challenges) = jagged_sumcheck_prove(&input, &mut transcript, None); assert_eq!(proof.proofs.len(), num_giga_vars); assert_eq!(challenges.len(), num_giga_vars); @@ -922,7 +928,7 @@ mod tests { let claimed_sum = input.compute_claimed_sum(); let mut transcript = BasicTranscript::::new(b"jagged_test_25"); - let (proof, challenges) = jagged_sumcheck_prove(&input, &mut transcript); + let (proof, challenges) = jagged_sumcheck_prove(&input, &mut transcript, None); assert_eq!(proof.proofs.len(), num_giga_vars); assert_eq!(challenges.len(), num_giga_vars); From c457a173578b199da817b672b5dcff595caa8a29 Mon Sep 17 00:00:00 2001 From: xkx Date: Thu, 23 Apr 2026 20:16:19 +0800 Subject: [PATCH 04/21] Feat: evaluator for jagged indicator function (#39) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * implement evaluate_g using width-4 ROBP from Section 3.2/4 Evaluate the MLE of indicator function g(a,b,c,d) = [a+c=b AND b * implement forward and backward algorithms for evaluate_g Separate evaluate_g into two functions following the ROBP-based MLE evaluation: forward (source→sinks, MLE definition) and backward (sinks→source, Claim 4.2.1/Lemma 4.2 from jagged PCS paper). Extract shared transition_weights helper. Default to backward for future batch/symbolic evaluation support. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- crates/mpcs/src/jagged_evaluator.rs | 266 ++++++++++++++++++++++++++++ crates/mpcs/src/lib.rs | 2 + 2 files changed, 268 insertions(+) create mode 100644 crates/mpcs/src/jagged_evaluator.rs diff --git a/crates/mpcs/src/jagged_evaluator.rs b/crates/mpcs/src/jagged_evaluator.rs new file mode 100644 index 0000000..8377427 --- /dev/null +++ b/crates/mpcs/src/jagged_evaluator.rs @@ -0,0 +1,266 @@ +use ff_ext::ExtensionField; + +// --------------------------------------------------------------------------- +// Succinct evaluation of ĝ(z₁, z₂, z₃, z₄) where g(a,b,c,d) = [a+c=b ∧ b( + z1i: E, + z2i: E, + z3i: E, + z4i: E, +) -> [(usize, usize, E, E, E); 4] { + let (nz1, nz2, nz3, nz4) = (E::ONE - z1i, E::ONE - z2i, E::ONE - z3i, E::ONE - z4i); + + // eq-weights for each (a,b) and (c,d) bit combination: + // abXY = eq₁(z1ᵢ, X) · eq₁(z2ᵢ, Y) + // cdXY = eq₁(z3ᵢ, X) · eq₁(z4ᵢ, Y) + // The eq-weight for symbol (a,b,c,d) = abAB · cdCD. + let ab00 = nz1 * nz2; + let ab01 = nz1 * z2i; + let ab10 = z1i * nz2; + let ab11 = z1i * z2i; + let cd00 = nz3 * nz4; + let cd01 = nz3 * z4i; + let cd10 = z3i * nz4; + let cd11 = z3i * z4i; + + // Each entry: (carry_in, carry_out, w_same, w_lt1, w_lt0) + // w_same: symbols where b = d (lt preserved) + // w_lt1: symbols where b = 0, d = 1 (lt forced to 1) + // w_lt0: symbols where b = 1, d = 0 (lt forced to 0) + [ + // carry_in=0, carry_out=0: consistent (a,c) ∈ {(0,0),(0,1),(1,0)}; b = a⊕c + ( + 0, + 0, + ab00 * cd00 + ab01 * cd11 + ab11 * cd01, // same: (0,0,0,0),(0,1,1,1),(1,1,0,1) + ab00 * cd01, // lt=1: (0,0,0,1) + ab01 * cd10 + ab11 * cd00, // lt=0: (0,1,1,0),(1,1,0,0) + ), + // carry_in=0, carry_out=1: consistent (a,c) = (1,1); b = 0 + (0, 1, ab10 * cd10, ab10 * cd11, E::ZERO), + // carry_in=1, carry_out=0: consistent (a,c) = (0,0); b = 1 + (1, 0, ab01 * cd01, E::ZERO, ab01 * cd00), + // carry_in=1, carry_out=1: consistent (a,c) ∈ {(0,1),(1,0),(1,1)}; b = a⊕c⊕1 + ( + 1, + 1, + ab00 * cd10 + ab10 * cd00 + ab11 * cd11, // same: (0,0,1,0),(1,0,0,0),(1,1,1,1) + ab00 * cd11 + ab10 * cd01, // lt=1: (0,0,1,1),(1,0,0,1) + ab11 * cd10, // lt=0: (1,1,1,0) + ), + ] +} + +/// Evaluate ĝ(z₁, z₂, z₃, z₄) using **forward** propagation (source → sinks). +/// +/// Maintains a state vector `α` of 4 weights, where `α[carry*2 + lt]` is the +/// total eq-weight of all paths from the source to state `(carry, lt)`. At each +/// step, all 16 transitions fire simultaneously, weighted by `eq(ζᵢ, σ)`. +/// +/// Computes: `(…((e₁ᵀ · T₁) · T₂) · … · Tₙ) · u` +/// +/// This follows from the MLE definition directly: +/// ĝ(z) = Σ_{sink v} label(v) · (Σ_{paths to v} Π_i eq(ζᵢ, σᵢ)) +/// The forward pass accumulates the path-weight sums layer by layer. +pub fn evaluate_g_forward(z1: &[E], z2: &[E], z3: &[E], z4: &[E]) -> E { + let n = z1.len(); + assert_eq!(z2.len(), n); + assert_eq!(z3.len(), n); + assert_eq!(z4.len(), n); + + // state[carry * 2 + lt]: weight on ROBP state (carry, lt). + // Initial: (carry=0, lt=0) with weight 1. + let mut state = [E::ZERO; 4]; + state[0] = E::ONE; + + for i in 0..n { + let transitions = transition_weights(z1[i], z2[i], z3[i], z4[i]); + + let mut new_state = [E::ZERO; 4]; + for &(ci, co, w_same, w_lt1, w_lt0) in &transitions { + let s0 = state[ci * 2]; // weight on (carry=ci, lt=0) + let s1 = state[ci * 2 + 1]; // weight on (carry=ci, lt=1) + let total = s0 + s1; + // lt preserved: w_same keeps lt=0 → lt=0 and lt=1 → lt=1 + // lt forced to 0: w_lt0 sends both lt=0 and lt=1 → lt=0 + // lt forced to 1: w_lt1 sends both lt=0 and lt=1 → lt=1 + new_state[co * 2] += w_same * s0 + w_lt0 * total; + new_state[co * 2 + 1] += w_same * s1 + w_lt1 * total; + } + state = new_state; + } + + // Accept state: (carry=0, lt=1) → index 1. + // g requires carry=0 (exact addition a+c=b) and lt=1 (b < d). + state[1] +} + +/// Evaluate ĝ(z₁, z₂, z₃, z₄) using **backward** propagation (sinks → source). +/// +/// Processes the ROBP in reverse (Lemma 4.2, Claim 4.2.1 of the jagged PCS paper): +/// ĝ_v(ζ, z) = Σ_{σ ∈ {0,1}⁴} eq(ζ, σ) · ĝ_{Γ(v,σ)}(z) +/// +/// Maintains a 4-element vector `val` where `val[carry*2 + lt]` is ĝ evaluated +/// at the corresponding sink, for the suffix of variables processed so far. +/// Starting from sink labels, we work backward to the source. +/// +/// Computes: `e₁ᵀ · (T₁ · (T₂ · … · (Tₙ · u)))` +/// +/// This is the paper's algorithm. It naturally extends to symbolic evaluation +/// (Lemma 4.5): by not multiplying with `u` at the end, we get a result vector +/// that works for any sink labels — useful for batch evaluation in the jagged +/// assist sumcheck. +pub fn evaluate_g_backward(z1: &[E], z2: &[E], z3: &[E], z4: &[E]) -> E { + let n = z1.len(); + assert_eq!(z2.len(), n); + assert_eq!(z3.len(), n); + assert_eq!(z4.len(), n); + + // Sink labels: u[carry*2 + lt] = label of sink state (carry, lt). + // g accepts iff carry=0 AND lt=1, so only state (0,1) has label 1. + let mut val = [E::ZERO; 4]; + val[1] = E::ONE; // (carry=0, lt=1) → accept + + // Process layers in reverse: from layer n-1 down to 0. + // After processing layer i, val[v] = ĝ_v(ζᵢ, ζᵢ₊₁, …, ζₙ₋₁) for the + // suffix z[i..n], where ĝ_v is the MLE starting at state v. + for i in (0..n).rev() { + let transitions = transition_weights(z1[i], z2[i], z3[i], z4[i]); + + // For each state v at layer i, compute: + // new_val[v] = Σ_{σ} eq(ζᵢ, σ) · val[Γ(v, σ)] + // Grouped by (carry_in, carry_out) and lt effect: + // new_val[ci*2 + lt_in] += w_same * val[co*2 + lt_in] (lt preserved) + // + w_lt1 * val[co*2 + 1] (lt forced to 1) + // + w_lt0 * val[co*2 + 0] (lt forced to 0) + let mut new_val = [E::ZERO; 4]; + for &(ci, co, w_same, w_lt1, w_lt0) in &transitions { + let v0 = val[co * 2]; // val at (carry=co, lt=0) + let v1 = val[co * 2 + 1]; // val at (carry=co, lt=1) + // State (ci, lt=0): lt preserved → goes to val[co, lt=0]; + // lt forced to 1 → goes to val[co, lt=1]; + // lt forced to 0 → goes to val[co, lt=0]. + new_val[ci * 2] += w_same * v0 + w_lt1 * v1 + w_lt0 * v0; + // State (ci, lt=1): lt preserved → goes to val[co, lt=1]; + // lt forced to 1 → goes to val[co, lt=1]; + // lt forced to 0 → goes to val[co, lt=0]. + new_val[ci * 2 + 1] += w_same * v1 + w_lt1 * v1 + w_lt0 * v0; + } + val = new_val; + } + + // val[0] = ĝ_source(z) = ĝ(z₁, z₂, z₃, z₄). + // Source is state (carry=0, lt=0) → index 0. + val[0] +} + +/// Evaluate ĝ(z₁, z₂, z₃, z₄) — delegates to the backward algorithm. +pub fn evaluate_g(z1: &[E], z2: &[E], z3: &[E], z4: &[E]) -> E { + evaluate_g_backward(z1, z2, z3, z4) +} + +#[cfg(test)] +mod tests { + use ff_ext::{BabyBearExt4, FromUniformBytes}; + use multilinear_extensions::virtual_poly::build_eq_x_r_vec; + use p3::field::FieldAlgebra; + use rand::thread_rng; + + use super::{evaluate_g_backward, evaluate_g_forward}; + + type E = BabyBearExt4; + + /// Brute-force MLE evaluation of g(a,b,c,d) = [a+c=b AND b E { + let n = z1.len(); + let size = 1usize << n; + let eq1 = build_eq_x_r_vec(z1); + let eq2 = build_eq_x_r_vec(z2); + let eq3 = build_eq_x_r_vec(z3); + let eq4 = build_eq_x_r_vec(z4); + + let mut sum = E::ZERO; + for (a, eq1a) in eq1.iter().enumerate().take(size) { + for (c, eq3c) in eq3.iter().enumerate().take(size) { + let b = a + c; + if b >= size { + continue; + } + for eq4d in eq4.iter().take(size).skip(b + 1) { + sum += *eq1a * eq2[b] * *eq3c * *eq4d; + } + } + } + sum + } + + #[test] + fn test_evaluate_g_forward() { + let mut rng = thread_rng(); + for n in 1..=4 { + let z1: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + let z2: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + let z3: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + let z4: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + + let expected = evaluate_g_bruteforce(&z1, &z2, &z3, &z4); + let result = evaluate_g_forward(&z1, &z2, &z3, &z4); + assert_eq!(result, expected, "forward mismatch at n={n}"); + } + } + + #[test] + fn test_evaluate_g_backward() { + let mut rng = thread_rng(); + for n in 1..=4 { + let z1: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + let z2: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + let z3: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + let z4: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + + let expected = evaluate_g_bruteforce(&z1, &z2, &z3, &z4); + let result = evaluate_g_backward(&z1, &z2, &z3, &z4); + assert_eq!(result, expected, "backward mismatch at n={n}"); + } + } + + #[test] + fn test_forward_backward_agree() { + let mut rng = thread_rng(); + for n in 1..=6 { + let z1: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + let z2: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + let z3: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + let z4: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + + let fwd = evaluate_g_forward(&z1, &z2, &z3, &z4); + let bwd = evaluate_g_backward(&z1, &z2, &z3, &z4); + assert_eq!(fwd, bwd, "forward != backward at n={n}"); + } + } +} diff --git a/crates/mpcs/src/lib.rs b/crates/mpcs/src/lib.rs index 4a4bb04..6462825 100644 --- a/crates/mpcs/src/lib.rs +++ b/crates/mpcs/src/lib.rs @@ -272,6 +272,8 @@ pub use jagged::{ JaggedCommitment, JaggedCommitmentWithWitness, JaggedSumcheckInput, jagged_commit, jagged_sumcheck_prove, }; +pub mod jagged_evaluator; +pub use jagged_evaluator::{evaluate_g, evaluate_g_backward, evaluate_g_forward}; #[cfg(feature = "whir")] extern crate whir as whir_external; #[cfg(feature = "whir")] From 4c3a23fea1afa159a718ab8357d4e183cd1cbab1 Mon Sep 17 00:00:00 2001 From: kunxian xia Date: Fri, 24 Apr 2026 00:04:25 +0800 Subject: [PATCH 05/21] feat: implement jagged_batch_open and jagged_batch_verify Add the opening protocol for the jagged PCS: given K evaluation claims on individual column polynomials, prove they're consistent with the commitment to the giga polynomial q'. Uses the jagged sumcheck to reduce to a single evaluation of q', then opens via the inner PCS. Co-Authored-By: Claude Opus 4.6 --- crates/mpcs/src/jagged.rs | 403 +++++++++++++++++++++++++++++++++++++- crates/mpcs/src/lib.rs | 4 +- 2 files changed, 399 insertions(+), 8 deletions(-) diff --git a/crates/mpcs/src/jagged.rs b/crates/mpcs/src/jagged.rs index 97e3b30..f3bfddf 100644 --- a/crates/mpcs/src/jagged.rs +++ b/crates/mpcs/src/jagged.rs @@ -51,25 +51,26 @@ //! 4. Pad `cat` to the next power of two (required for MLE representation). //! 5. Commit to the padded `cat` as a single-column matrix using the inner PCS. -use std::iter::once; +use std::{iter::once, marker::PhantomData}; -use crate::{Error, PolynomialCommitmentScheme}; +use crate::{Error, PolynomialCommitmentScheme, jagged_evaluator::evaluate_g}; use ff_ext::ExtensionField; use itertools::Itertools; use multilinear_extensions::{ - mle::MultilinearExtension, util::max_usable_threads, virtual_poly::build_eq_x_r_vec, + mle::{FieldType, MultilinearExtension}, + util::max_usable_threads, + virtual_poly::{VPAuxInfo, build_eq_x_r_vec}, }; use p3::{ matrix::{Matrix, bitrev::BitReversableMatrix}, maybe_rayon::prelude::{ - IndexedParallelIterator, IntoParallelIterator, ParallelIterator, ParallelSlice, - ParallelSliceMut, + IntoParallelIterator, ParallelIterator, ParallelSlice, ParallelSliceMut, }, }; use serde::{Deserialize, Serialize}; use sumcheck::{ macros::{entered_span, exit_span}, - structs::{IOPProof, IOPProverMessage, IOPProverState}, + structs::{IOPProof, IOPProverMessage, IOPProverState, IOPVerifierState}, }; use transcript::Transcript; use witness::{InstancePaddingStrategy, RowMajorMatrix}; @@ -673,6 +674,240 @@ fn bind_and_materialize( results.into_iter().unzip() } +// --------------------------------------------------------------------------- +// Jagged Batch Open / Verify +// --------------------------------------------------------------------------- + +/// Proof for the jagged batch opening protocol. +/// +/// Contains a sumcheck proof (reducing K column evaluation claims to a single +/// point on q'), the evaluation q'(ρ), and an inner PCS opening proof. +#[derive(Clone, Serialize, Deserialize)] +#[serde(bound(serialize = "", deserialize = ""))] +pub struct JaggedBatchOpenProof> { + pub sumcheck_proof: IOPProof, + pub q_eval: E, + pub inner_proof: Pcs::Proof, +} + +/// Convert a `usize` to its little-endian binary representation as field elements. +fn int_to_field_bits(val: usize, num_bits: usize) -> Vec { + (0..num_bits) + .map(|i| if (val >> i) & 1 == 1 { E::ONE } else { E::ZERO }) + .collect() +} + +/// Evaluate f̂(ρ) using the ROBP evaluator. +/// +/// f̂(ρ) = Σ_{y} eq_col[y] · ĝ(z_row_padded, ρ, bits(t_y), bits(t_{y+1})) +fn compute_f_at_point( + z_row_padded: &[E], + rho_padded: &[E], + eq_col: &[E], + cumulative_heights: &[usize], + n_robp: usize, +) -> E { + let num_polys = cumulative_heights.len() - 1; + let mut f_val = E::ZERO; + for y in 0..num_polys { + if eq_col[y] == E::ZERO { + continue; + } + let t_lo = int_to_field_bits::(cumulative_heights[y], n_robp); + let t_hi = int_to_field_bits::(cumulative_heights[y + 1], n_robp); + let g_val = evaluate_g(z_row_padded, rho_padded, &t_lo, &t_hi); + f_val += eq_col[y] * g_val; + } + f_val +} + +/// Prove that evaluation claims `evals[i] = p_i(point)` are consistent with a +/// jagged commitment. +/// +/// All polynomials must have the same height. `point` is the common evaluation +/// point (a suffix `r[s..]` of the GKR challenge, length `s = log2(poly_height)`). +/// +/// The protocol: +/// 1. Batch the K column claims via a random column challenge `z_col`. +/// 2. Run the jagged sumcheck to reduce to a single evaluation of q'. +/// 3. Open q' at the sumcheck output point via the inner PCS. +pub fn jagged_batch_open>( + pp: &Pcs::ProverParam, + comm: &JaggedCommitmentWithWitness, + point: &[E], + evals: &[E], + transcript: &mut impl Transcript, +) -> Result, Error> { + let num_polys = comm.num_polys(); + if evals.len() != num_polys { + return Err(Error::InvalidPcsParam(format!( + "jagged_batch_open: expected {} evals, got {}", + num_polys, + evals.len() + ))); + } + + let s = point.len(); + let expected_height = 1usize << s; + for (i, &h) in comm.poly_heights.iter().enumerate() { + if h != expected_height { + return Err(Error::InvalidPcsParam(format!( + "jagged_batch_open: poly {} has height {}, expected {}", + i, h, expected_height + ))); + } + } + + let total_evals = comm.total_evaluations(); + let num_giga_vars = total_evals.next_power_of_two().trailing_zeros() as usize; + + // z_row = reverse(point) to account for bit-reversal in jagged_commit. + let z_row: Vec = point.iter().rev().cloned().collect(); + + // Write evals to transcript, then sample z_col. + transcript.append_field_element_exts(evals); + let num_col_vars = (num_polys.next_power_of_two().trailing_zeros() as usize).max(1); + let z_col: Vec = transcript.sample_and_append_vec(b"jagged_z_col", num_col_vars); + + let eq_col = build_eq_x_r_vec(&z_col); + + // Extract q' base-field evaluations from the inner PCS witness. + let q_mles = Pcs::get_arc_mle_witness_from_commitment(&comm.inner); + assert_eq!( + q_mles.len(), + 1, + "jagged commit produces exactly one polynomial" + ); + let q_mle = &q_mles[0]; + let q_evals_base: &[E::BaseField] = match q_mle.evaluations() { + FieldType::Base(slice) => slice, + _ => { + return Err(Error::InvalidPcsParam( + "jagged_batch_open: expected base-field evaluations for q'".into(), + )); + } + }; + + // Run the jagged sumcheck. + let eq_row = build_eq_x_r_vec(&z_row); + let input = JaggedSumcheckInput { + q_evals: q_evals_base, + num_giga_vars, + cumulative_heights: &comm.cumulative_heights, + eq_row, + eq_col, + }; + let (sumcheck_proof, challenges) = jagged_sumcheck_prove(&input, transcript, None); + + // Evaluate q'(ρ). + let q_eval = q_mle.evaluate(&challenges); + + // Write q_eval to transcript before inner PCS open. + transcript.append_field_element_ext(&q_eval); + + // Open q' at ρ via inner PCS batch_open. + let inner_proof = Pcs::batch_open( + pp, + vec![(&comm.inner, vec![(challenges, vec![q_eval])])], + transcript, + )?; + + Ok(JaggedBatchOpenProof { + sumcheck_proof, + q_eval, + inner_proof, + }) +} + +/// Verify that evaluation claims `evals[i] = p_i(point)` are consistent with a +/// jagged commitment. +pub fn jagged_batch_verify>( + vp: &Pcs::VerifierParam, + comm: &JaggedCommitment, + point: &[E], + evals: &[E], + proof: &JaggedBatchOpenProof, + transcript: &mut impl Transcript, +) -> Result<(), Error> { + let num_polys = comm.cumulative_heights.len() - 1; + if evals.len() != num_polys { + return Err(Error::InvalidPcsOpen(format!( + "jagged_batch_verify: expected {} evals, got {}", + num_polys, + evals.len() + ))); + } + + let total_evals = *comm.cumulative_heights.last().unwrap(); + let num_giga_vars = total_evals.next_power_of_two().trailing_zeros() as usize; + + // z_row = reverse(point) + let z_row: Vec = point.iter().rev().cloned().collect(); + + // Replay transcript: write evals, sample z_col. + transcript.append_field_element_exts(evals); + let num_col_vars = (num_polys.next_power_of_two().trailing_zeros() as usize).max(1); + let z_col: Vec = transcript.sample_and_append_vec(b"jagged_z_col", num_col_vars); + + // claimed_sum = Σ_j eq_col[j] · evals[j] + let eq_col = build_eq_x_r_vec(&z_col); + let claimed_sum: E = eq_col[..num_polys] + .iter() + .zip(evals.iter()) + .map(|(&eq, &ev)| eq * ev) + .sum(); + + // Verify the jagged sumcheck. + let aux_info = VPAuxInfo { + max_degree: 2, + max_num_variables: num_giga_vars, + phantom: PhantomData::, + }; + let subclaim = + IOPVerifierState::::verify(claimed_sum, &proof.sumcheck_proof, &aux_info, transcript); + let rho: Vec = subclaim.point.iter().map(|c| c.elements).collect(); + + // The ROBP needs enough bits to represent the max cumulative height (= total_evals). + // When total_evals is an exact power of 2, num_giga_vars bits can't represent it. + let n_robp = num_giga_vars + if total_evals.is_power_of_two() { 1 } else { 0 }; + + // Compute f̂(ρ) via the ROBP evaluator. + let mut z_row_padded = z_row; + z_row_padded.resize(n_robp, E::ZERO); + let mut rho_padded = rho.clone(); + rho_padded.resize(n_robp, E::ZERO); + let f_at_rho = compute_f_at_point( + &z_row_padded, + &rho_padded, + &eq_col, + &comm.cumulative_heights, + n_robp, + ); + + // Write q_eval to transcript (must match prover). + transcript.append_field_element_ext(&proof.q_eval); + + // Check subclaim: q'(ρ) · f̂(ρ) == expected_evaluation. + if proof.q_eval * f_at_rho != subclaim.expected_evaluation { + return Err(Error::InvalidPcsOpen( + "jagged_batch_verify: q_eval * f(rho) != subclaim expected evaluation".into(), + )); + } + + // Verify the inner PCS opening. + Pcs::batch_verify( + vp, + vec![( + comm.inner.clone(), + vec![(num_giga_vars, (rho, vec![proof.q_eval]))], + )], + &proof.inner_proof, + transcript, + )?; + + Ok(()) +} + #[cfg(test)] mod tests { use super::*; @@ -954,4 +1189,160 @@ mod tests { "final evaluation mismatch" ); } + + // --- Batch open/verify tests --- + + use ff_ext::FromUniformBytes; + + /// Evaluate a single column polynomial at `point` (the original, non-bit-reversed poly). + /// `col_evals` are the raw evaluations of the column (before bit-reversal). + fn eval_column_poly_at_point(col_evals: &[F], point: &[E]) -> E { + let s = point.len(); + assert_eq!(col_evals.len(), 1 << s); + let mle = MultilinearExtension::from_evaluations_vec(s, col_evals.to_vec()); + mle.evaluate(point) + } + + #[test] + fn test_jagged_batch_open_verify_small() { + let mut rng = thread_rng(); + + let num_rows = 1024usize; // s=10 + let num_cols = 3usize; + let s = 10; + let total_evals = num_rows * num_cols; + let num_giga_vars = total_evals.next_power_of_two().trailing_zeros() as usize; + + let (pp, vp) = setup_pcs::(num_giga_vars); + let rmm = make_rmm(num_rows, num_cols); + + // Extract column polynomials (before bit-reversal) for computing true evaluations. + let col_polys: Vec> = (0..num_cols) + .map(|c| { + (0..num_rows) + .map(|r| rmm.values[r * num_cols + c]) + .collect() + }) + .collect(); + + // Commit. + let mut transcript_p = BasicTranscript::::new(b"jagged_batch_test"); + let comm = jagged_commit::(&pp, vec![rmm]).expect("commit should succeed"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); + + // Random evaluation point of length s. + let point: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); + + // Compute true evaluations. + let evals: Vec = col_polys + .iter() + .map(|col| eval_column_poly_at_point(col, &point)) + .collect(); + + // Prover: batch open. + let proof = jagged_batch_open::(&pp, &comm, &point, &evals, &mut transcript_p) + .expect("batch open should succeed"); + + // Verifier: batch verify. + let mut transcript_v = BasicTranscript::::new(b"jagged_batch_test"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_v).unwrap(); + + jagged_batch_verify::( + &vp, + &comm.to_commitment(), + &point, + &evals, + &proof, + &mut transcript_v, + ) + .expect("batch verify should succeed"); + } + + #[test] + fn test_jagged_batch_open_verify_single_poly() { + let mut rng = thread_rng(); + + let num_rows = 1024usize; // s=10 + let num_cols = 1usize; + let s = 10; + let num_giga_vars = (num_rows * num_cols).next_power_of_two().trailing_zeros() as usize; + + let (pp, vp) = setup_pcs::(num_giga_vars); + let rmm = make_rmm(num_rows, num_cols); + + let col_poly: Vec = (0..num_rows).map(|r| rmm.values[r]).collect(); + + let mut transcript_p = BasicTranscript::::new(b"jagged_single"); + let comm = jagged_commit::(&pp, vec![rmm]).expect("commit"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); + + let point: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); + let evals = vec![eval_column_poly_at_point(&col_poly, &point)]; + + let proof = jagged_batch_open::(&pp, &comm, &point, &evals, &mut transcript_p) + .expect("batch open"); + + let mut transcript_v = BasicTranscript::::new(b"jagged_single"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_v).unwrap(); + + jagged_batch_verify::( + &vp, + &comm.to_commitment(), + &point, + &evals, + &proof, + &mut transcript_v, + ) + .expect("batch verify"); + } + + #[test] + fn test_jagged_batch_open_verify_soundness() { + let mut rng = thread_rng(); + + let num_rows = 1024usize; + let num_cols = 3usize; + let s = 10; + let num_giga_vars = (num_rows * num_cols).next_power_of_two().trailing_zeros() as usize; + + let (pp, vp) = setup_pcs::(num_giga_vars); + let rmm = make_rmm(num_rows, num_cols); + + let col_polys: Vec> = (0..num_cols) + .map(|c| { + (0..num_rows) + .map(|r| rmm.values[r * num_cols + c]) + .collect() + }) + .collect(); + + let mut transcript_p = BasicTranscript::::new(b"jagged_soundness"); + let comm = jagged_commit::(&pp, vec![rmm]).expect("commit"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); + + let point: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); + let mut evals: Vec = col_polys + .iter() + .map(|col| eval_column_poly_at_point(col, &point)) + .collect(); + + // Tamper with one evaluation. + evals[1] += E::ONE; + + let proof = jagged_batch_open::(&pp, &comm, &point, &evals, &mut transcript_p) + .expect("batch open with wrong evals still produces a proof"); + + let mut transcript_v = BasicTranscript::::new(b"jagged_soundness"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_v).unwrap(); + + let result = jagged_batch_verify::( + &vp, + &comm.to_commitment(), + &point, + &evals, + &proof, + &mut transcript_v, + ); + assert!(result.is_err(), "verify should reject tampered evaluations"); + } } diff --git a/crates/mpcs/src/lib.rs b/crates/mpcs/src/lib.rs index 6462825..3c68f3a 100644 --- a/crates/mpcs/src/lib.rs +++ b/crates/mpcs/src/lib.rs @@ -269,8 +269,8 @@ pub use basefold::{ }; pub mod jagged; pub use jagged::{ - JaggedCommitment, JaggedCommitmentWithWitness, JaggedSumcheckInput, jagged_commit, - jagged_sumcheck_prove, + JaggedBatchOpenProof, JaggedCommitment, JaggedCommitmentWithWitness, JaggedSumcheckInput, + jagged_batch_open, jagged_batch_verify, jagged_commit, jagged_sumcheck_prove, }; pub mod jagged_evaluator; pub use jagged_evaluator::{evaluate_g, evaluate_g_backward, evaluate_g_forward}; From a8240d06c513af307700dbb3bafbff244dc590a9 Mon Sep 17 00:00:00 2001 From: kunxian xia Date: Fri, 24 Apr 2026 11:07:05 +0800 Subject: [PATCH 06/21] update doc comments --- crates/mpcs/src/jagged.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/crates/mpcs/src/jagged.rs b/crates/mpcs/src/jagged.rs index f3bfddf..0614eb9 100644 --- a/crates/mpcs/src/jagged.rs +++ b/crates/mpcs/src/jagged.rs @@ -9,7 +9,7 @@ //! chips into a single "giga" multilinear polynomial `q'`: //! //! ```text -//! q' = bitrev(p_0) || bitrev(p_1) || ... || bitrev(p_N) +//! q' = bitrev(p_0) || bitrev(p_1) || ... || bitrev(p_{N-1}) //! ``` //! //! where each `p_i` is a column polynomial extracted from the input trace matrices, and @@ -17,7 +17,7 @@ //! //! ## Suffix-to-Prefix Transformation //! -//! The main sumcheck prover outputs evaluations `v_i = p_i(r[(n-s)..n])` — i.e., at the +//! The main sumcheck prover outputs evaluations `v_i = p_i(r[(m-s)..m])` — i.e., at the //! **suffix** of the random challenge point. To make these evaluations compatible with the //! jagged sumcheck (which operates on prefix-aligned polynomials), we apply a bit-reversal //! permutation to each polynomial's evaluations: @@ -26,7 +26,7 @@ //! p_i'[j] = p_i[bitrev_s(j)] (for j in 0..2^s) //! ``` //! -//! After bit-reversal, `v_i = p_i(r[(n-s)..n]) = p_i'(reverse(r[(n-s)..n]))`. +//! After bit-reversal, `v_i = p_i(r[(m-s)..m]) = p_i'(reverse(r[(m-s)..m]))`. //! //! ## Cumulative Heights //! @@ -35,11 +35,11 @@ //! - `t[i+1] = t[i] + h_i` where `h_i = 2^(num_vars of p_i)` is the number of evaluations //! //! Given a position `b` in `q'`, the verifier can locate the corresponding `(i, r)` pair via: -//! - `t[i-1] <= b < t[i]` -//! - `r = b - t[i-1]` +//! - `t[i] <= b < t[i+1]` +//! - `r = b - t[i]` //! //! The cumulative heights allow the verifier to succinctly evaluate the indicator function -//! `g(z_r, z_b, t[i-1], t[i])` needed for the jagged sumcheck. +//! `g(z_r, z_b, t[i], t[i+1])` needed for the jagged sumcheck. //! //! ## Commit Protocol //! From 634f7254ab15bab65524291bf5b56b3bfe8588d1 Mon Sep 17 00:00:00 2001 From: kunxian xia Date: Fri, 24 Apr 2026 14:12:47 +0800 Subject: [PATCH 07/21] polish doc comments --- crates/mpcs/src/jagged.rs | 74 ++++++++++++++++++++++++++++++++------- 1 file changed, 61 insertions(+), 13 deletions(-) diff --git a/crates/mpcs/src/jagged.rs b/crates/mpcs/src/jagged.rs index 0614eb9..3a5d49a 100644 --- a/crates/mpcs/src/jagged.rs +++ b/crates/mpcs/src/jagged.rs @@ -1,7 +1,9 @@ -//! Jagged PCS Commit Adaptor +//! Jagged PCS //! -//! This module implements the commit protocol for the Jagged PCS as described in -//! the SP1 Jagged PCS paper () and Ceno issues #1272 / #1288. +//! Implements the Jagged PCS protocol (commit, batch open, batch verify) tailored to +//! Ceno's scenario, based on the SP1 Jagged PCS paper +//! () and Ceno issue +//! [#1272](https://github.com/scroll-tech/ceno/issues/1272). //! //! ## Overview //! @@ -12,21 +14,31 @@ //! q' = bitrev(p_0) || bitrev(p_1) || ... || bitrev(p_{N-1}) //! ``` //! -//! where each `p_i` is a column polynomial extracted from the input trace matrices, and -//! `bitrev` is the suffix-to-prefix bit-reversal transformation. +//! where each `p_i` is a column polynomial extracted from the input trace matrices. +//! We cannot use the naive `q = p_0 || ... || p_{N-1}` because the jagged sumcheck +//! requires prefix-aligned evaluation points, but Ceno's main sumcheck evaluates +//! polynomials at a **suffix** of the challenge point. The `bitrev` (suffix-to-prefix +//! bit-reversal) bridges the two, making the jagged PCS applicable (see below). //! //! ## Suffix-to-Prefix Transformation //! -//! The main sumcheck prover outputs evaluations `v_i = p_i(r[(m-s)..m])` — i.e., at the -//! **suffix** of the random challenge point. To make these evaluations compatible with the -//! jagged sumcheck (which operates on prefix-aligned polynomials), we apply a bit-reversal -//! permutation to each polynomial's evaluations: +//! Each `p_i` has `s` variables (where `s` varies per polynomial; we write `s` instead +//! of `s_i` for brevity). The main sumcheck prover outputs evaluations +//! `v_i = p_i(z_r[(m-s)..m])` — i.e., at the suffix of the random challenge point. +//! To make these evaluations compatible with the jagged sumcheck (which operates on +//! prefix-aligned polynomials), we apply a bit-reversal permutation to each polynomial's +//! evaluations: //! //! ```text //! p_i'[j] = p_i[bitrev_s(j)] (for j in 0..2^s) //! ``` //! -//! After bit-reversal, `v_i = p_i(r[(m-s)..m]) = p_i'(reverse(r[(m-s)..m]))`. +//! After bit-reversal, +//! ```text +//! v_i = p_i(z_r[(m-s)..m]) +//! = p_i'(z_r'[..s]) +//! ``` +//! where `z_r' = reverse(z_r)`. //! //! ## Cumulative Heights //! @@ -34,12 +46,12 @@ //! - `t[0] = 0` //! - `t[i+1] = t[i] + h_i` where `h_i = 2^(num_vars of p_i)` is the number of evaluations //! -//! Given a position `b` in `q'`, the verifier can locate the corresponding `(i, r)` pair via: +//! Given a position `b` in `q'`, the inverse mapping `inv(b) = (i, r)` is defined by: //! - `t[i] <= b < t[i+1]` //! - `r = b - t[i]` //! -//! The cumulative heights allow the verifier to succinctly evaluate the indicator function -//! `g(z_r, z_b, t[i], t[i+1])` needed for the jagged sumcheck. +//! Using the indicator `g(r, b, t[i], t[i+1]) = [r + t[i] = b ∧ b < t[i+1]]`, this has +//! the closed form: `inv(b) = Σ_{i,r} (i, r) · g(r, b, t[i], t[i+1])`. //! //! ## Commit Protocol //! @@ -50,6 +62,42 @@ //! 3. Compute cumulative heights `t[i]`. //! 4. Pad `cat` to the next power of two (required for MLE representation). //! 5. Commit to the padded `cat` as a single-column matrix using the inner PCS. +//! +//! ## Batch Open Protocol +//! +//! For notational simplicity, we write `p_i` and `z_r` instead of `p_i'` and `z_r'` +//! (the bit-reversed variants) throughout this section. +//! +//! Each column opening `v_i = p_i(z_r[..s])` requires its own sumcheck. We batch all K +//! openings into one using `eq(z_c, ·)` weights (soundness loss: `log2(K) / |E|`): +//! +//! ```text +//! v = Σ_i eq(z_c, i) · p_i(z_r[..s]) +//! = Σ_i eq(z_c, i) · Σ_r eq(z_r, r) · p_i(r) +//! = Σ_{i,r} f(i, r) · p_i(r) where f(i, r) = eq(z_c, i) · eq(z_r, r) +//! ``` +//! +//! We apply `inv(b)` to rewrite `f` in terms of the giga-index `b`: +//! +//! ```text +//! f(inv(b)) = Σ_{i,r} f(i, r) · g(r, b, t[i], t[i+1]) +//! = Σ_i eq(z_c, i) · ĝ(z_r, b, t[i], t[i+1]) +//! ``` +//! +//! where `ĝ(z_r, b, ·) = Σ_r eq(z_r, r) · g(r, b, ·)` absorbs the row weight. +//! +//! Since summation over (i,r) is equivalent to summation over b, we can rewrite the batch opening claim as +//! +//! ```text +//! v = Σ_b q'(b) · f(inv(b)) +//! = Σ_b q'(b) · Σ_i eq(z_c, i) · ĝ(z_r, b, t[i], t[i+1]) +//! ``` +//! +//! Defining `h(b) = f(inv(b)) = Σ_i eq(z_c, i) · ĝ(z_r, b, t[i], t[i+1])` (multilinear in `b`): +//! +//! The sumcheck reduces this to a single opening of `q'(ρ)` plus a verifier check of +//! `h(ρ)`. The ROBP makes `ĝ` efficiently evaluable, so the verifier computes +//! `h(ρ) = Σ_i eq(z_c, i) · ĝ(z_r, ρ, t[i], t[i+1])` in `O(K·n)` time. use std::{iter::once, marker::PhantomData}; From ef452293f566336688d4323cec8771003f20daeb Mon Sep 17 00:00:00 2001 From: kunxian xia Date: Fri, 24 Apr 2026 14:27:14 +0800 Subject: [PATCH 08/21] polish doc comments --- crates/mpcs/src/jagged.rs | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/crates/mpcs/src/jagged.rs b/crates/mpcs/src/jagged.rs index 3a5d49a..cf8d588 100644 --- a/crates/mpcs/src/jagged.rs +++ b/crates/mpcs/src/jagged.rs @@ -23,8 +23,9 @@ //! ## Suffix-to-Prefix Transformation //! //! Each `p_i` has `s` variables (where `s` varies per polynomial; we write `s` instead -//! of `s_i` for brevity). The main sumcheck prover outputs evaluations -//! `v_i = p_i(z_r[(m-s)..m])` — i.e., at the suffix of the random challenge point. +//! of `s_i` for brevity). Let `m` be the length of `z_r`. The main sumcheck prover +//! outputs evaluations `v_i = p_i(z_r[(m-s)..m])` — i.e., at the suffix of the +//! random challenge point. //! To make these evaluations compatible with the jagged sumcheck (which operates on //! prefix-aligned polynomials), we apply a bit-reversal permutation to each polynomial's //! evaluations: @@ -33,7 +34,7 @@ //! p_i'[j] = p_i[bitrev_s(j)] (for j in 0..2^s) //! ``` //! -//! After bit-reversal, +//! After bit-reversal, //! ```text //! v_i = p_i(z_r[(m-s)..m]) //! = p_i'(z_r'[..s]) @@ -69,11 +70,11 @@ //! (the bit-reversed variants) throughout this section. //! //! Each column opening `v_i = p_i(z_r[..s])` requires its own sumcheck. We batch all K -//! openings into one using `eq(z_c, ·)` weights (soundness loss: `log2(K) / |E|`): +//! openings into one using `eq(z_c, ·)` weights (soundness loss: `log2(N) / |E|`): //! //! ```text //! v = Σ_i eq(z_c, i) · p_i(z_r[..s]) -//! = Σ_i eq(z_c, i) · Σ_r eq(z_r, r) · p_i(r) +//! = Σ_{i,r} eq(z_c, i) · eq(z_r, r) · p_i(r) //! = Σ_{i,r} f(i, r) · p_i(r) where f(i, r) = eq(z_c, i) · eq(z_r, r) //! ``` //! From 20f8dc528d81b5c2decee824655af6dcaaac5c39 Mon Sep 17 00:00:00 2001 From: kunxian xia Date: Fri, 24 Apr 2026 14:36:29 +0800 Subject: [PATCH 09/21] reorg --- .../evaluator.rs} | 0 crates/mpcs/src/{jagged.rs => jagged/mod.rs} | 637 +----------------- crates/mpcs/src/jagged/sumcheck.rs | 623 +++++++++++++++++ crates/mpcs/src/lib.rs | 5 +- 4 files changed, 639 insertions(+), 626 deletions(-) rename crates/mpcs/src/{jagged_evaluator.rs => jagged/evaluator.rs} (100%) rename crates/mpcs/src/{jagged.rs => jagged/mod.rs} (54%) create mode 100644 crates/mpcs/src/jagged/sumcheck.rs diff --git a/crates/mpcs/src/jagged_evaluator.rs b/crates/mpcs/src/jagged/evaluator.rs similarity index 100% rename from crates/mpcs/src/jagged_evaluator.rs rename to crates/mpcs/src/jagged/evaluator.rs diff --git a/crates/mpcs/src/jagged.rs b/crates/mpcs/src/jagged/mod.rs similarity index 54% rename from crates/mpcs/src/jagged.rs rename to crates/mpcs/src/jagged/mod.rs index cf8d588..bd201f8 100644 --- a/crates/mpcs/src/jagged.rs +++ b/crates/mpcs/src/jagged/mod.rs @@ -36,7 +36,7 @@ //! //! After bit-reversal, //! ```text -//! v_i = p_i(z_r[(m-s)..m]) +//! v_i = p_i(z_r[(m-s)..m]) //! = p_i'(z_r'[..s]) //! ``` //! where `z_r' = reverse(z_r)`. @@ -100,27 +100,27 @@ //! `h(ρ)`. The ROBP makes `ĝ` efficiently evaluable, so the verifier computes //! `h(ρ) = Σ_i eq(z_c, i) · ĝ(z_r, ρ, t[i], t[i+1])` in `O(K·n)` time. +pub mod evaluator; +pub mod sumcheck; + +pub use evaluator::{evaluate_g, evaluate_g_backward, evaluate_g_forward}; +pub use sumcheck::{JaggedSumcheckInput, jagged_sumcheck_prove}; + use std::{iter::once, marker::PhantomData}; -use crate::{Error, PolynomialCommitmentScheme, jagged_evaluator::evaluate_g}; +use crate::{Error, PolynomialCommitmentScheme}; +use ::sumcheck::structs::{IOPProof, IOPVerifierState}; use ff_ext::ExtensionField; use itertools::Itertools; use multilinear_extensions::{ - mle::{FieldType, MultilinearExtension}, - util::max_usable_threads, + mle::FieldType, virtual_poly::{VPAuxInfo, build_eq_x_r_vec}, }; use p3::{ matrix::{Matrix, bitrev::BitReversableMatrix}, - maybe_rayon::prelude::{ - IntoParallelIterator, ParallelIterator, ParallelSlice, ParallelSliceMut, - }, + maybe_rayon::prelude::{IntoParallelIterator, ParallelSliceMut}, }; use serde::{Deserialize, Serialize}; -use sumcheck::{ - macros::{entered_span, exit_span}, - structs::{IOPProof, IOPProverMessage, IOPProverState, IOPVerifierState}, -}; use transcript::Transcript; use witness::{InstancePaddingStrategy, RowMajorMatrix}; @@ -296,433 +296,6 @@ pub fn jagged_commit>( }) } -// --------------------------------------------------------------------------- -// Jagged Sumcheck Prover (M-table streaming algorithm) -// --------------------------------------------------------------------------- - -// Streaming sumcheck prover using the M-table algorithm from -// "Time-Space Trade-Offs for Sumcheck" (eprint 2025/1473), Section 4. - -/// Default log2 of the maximum epoch size for the streaming phase. -/// Epoch sizes are `[2^0, 2^1, ..., 2^LOG2_MAX_EPOCH]` = `[1, 2, 4, 8]`, -/// covering 1 + 2 + 4 + 8 = 15 streaming rounds before switching to standard sumcheck. -const LOG2_MAX_EPOCH: u32 = 3; - -/// All inputs needed for the jagged sumcheck. -pub struct JaggedSumcheckInput<'a, E: ExtensionField> { - /// Giga polynomial evaluations (concatenated, bit-reversed). - pub q_evals: &'a [E::BaseField], - /// n = log2(padded_total_size). - pub num_giga_vars: usize, - /// Cumulative height sequence t[j], length num_polys + 1. - pub cumulative_heights: &'a [usize], - /// Precomputed eq table for the row evaluation point: `build_eq_x_r_vec(z_row)`. - pub eq_row: Vec, - /// Precomputed eq table for the column challenge point: `build_eq_x_r_vec(z_col)`. - pub eq_col: Vec, -} - -/// Iterator that yields `(col, row)` pairs for consecutive giga indices. -/// Uses one binary search at construction, then O(1) per step. -struct ColRowIter<'a> { - cumulative_heights: &'a [usize], - col: usize, - row: usize, - num_polys: usize, -} - -impl<'a> Iterator for ColRowIter<'a> { - type Item = (usize, usize); - - fn next(&mut self) -> Option<(usize, usize)> { - if self.col >= self.num_polys { - return None; - } - let result = (self.col, self.row); - self.row += 1; - let poly_height = self.cumulative_heights[self.col + 1] - self.cumulative_heights[self.col]; - if self.row >= poly_height { - self.row = 0; - self.col += 1; - } - Some(result) - } -} - -impl<'a, E: ExtensionField> JaggedSumcheckInput<'a, E> { - fn total_evaluations(&self) -> usize { - *self.cumulative_heights.last().unwrap_or(&0) - } - - /// Return an iterator yielding `(col, row)` for consecutive giga indices - /// starting from `start`. One binary search at construction, O(1) per step. - fn col_row_iter(&self, start: usize) -> ColRowIter<'_> { - let num_polys = self.cumulative_heights.len() - 1; - if start >= self.total_evaluations() { - return ColRowIter { - cumulative_heights: self.cumulative_heights, - col: num_polys, - row: 0, - num_polys, - }; - } - let j = self.cumulative_heights.partition_point(|&t| t <= start) - 1; - ColRowIter { - cumulative_heights: self.cumulative_heights, - col: j, - row: start - self.cumulative_heights[j], - num_polys, - } - } - - /// Brute-force computation of sum_b q'(b) * f(b). - /// O(2^n) time — only for debugging and tests. - #[cfg(test)] - fn compute_claimed_sum(&self) -> E { - self.q_evals[..self.total_evaluations()] - .iter() - .zip(self.col_row_iter(0)) - .fold(E::ZERO, |sum, (&q, (col, row))| { - sum + (self.eq_row[row] * self.eq_col[col]) * q - }) - } - - /// Brute-force MLE evaluation of q'(rb) and f(rb) at the given point. - /// O(2^n) time and memory — only for debugging and tests. - /// Returns `(q_at_point, f_at_point)`. - #[cfg(test)] - fn final_evaluations_slow(&self, point: &[E]) -> (E, E) { - let n = self.num_giga_vars; - let total_evals = self.total_evaluations(); - - // Build q' MLE (padded with zeros) and evaluate at point. - let mut q_padded: Vec = self.q_evals.to_vec(); - q_padded.resize(1 << n, Default::default()); - let q_mle = MultilinearExtension::from_evaluations_vec(n, q_padded); - let q_at_point = q_mle.evaluate(point); - - // Build f MLE from eq tables and evaluate at point. - let mut f_evals = vec![E::ZERO; 1 << n]; - for (f_eval, (col, row)) in f_evals[..total_evals].iter_mut().zip(self.col_row_iter(0)) { - *f_eval = self.eq_row[row] * self.eq_col[col]; - } - let f_mle = MultilinearExtension::from_evaluations_ext_vec(n, f_evals); - let f_at_point = f_mle.evaluate(point); - - (q_at_point, f_at_point) - } -} - -/// Run the full jagged sumcheck: streaming phase (rounds 1..K) + standard phase (rounds K+1..n). -/// -/// `log2_max_epoch` controls the streaming phase epoch schedule. Epoch sizes are -/// `[1, 2, 4, ..., 2^k]` where `k = log2_max_epoch`. Pass `None` to use the default -/// `LOG2_MAX_EPOCH = 3`, giving epoch sizes `[1, 2, 4, 8]` and 15 streaming rounds. -/// -/// Returns the proof and the full list of challenges (r_1, ..., r_n). -pub fn jagged_sumcheck_prove( - input: &JaggedSumcheckInput, - transcript: &mut impl Transcript, - log2_max_epoch: Option, -) -> (IOPProof, Vec) { - let n = input.num_giga_vars; - let max_degree: usize = 2; - - let k = log2_max_epoch.unwrap_or(LOG2_MAX_EPOCH); - let epoch_sizes: Vec = (0..=k).map(|i| 1usize << i).collect(); - - let mut challenges: Vec = Vec::with_capacity(n); - let mut proof_messages: Vec> = Vec::with_capacity(n); - - // Write transcript header (must match verifier's expectations). - transcript.append_message(&n.to_le_bytes()); - transcript.append_message(&max_degree.to_le_bytes()); - - // --- Streaming phase: epochs j' = 1, 2, 4, ..., 2^k --- - for &epoch_size in &epoch_sizes { - // Epoch j' handles rounds j'..2j'-1. Skip if all rounds are done. - if epoch_size > n { - break; - } - - // Build M-table for this epoch. - let span = entered_span!("build_m_table", epoch = epoch_size); - let m_table = build_m_table(input, &challenges, epoch_size); - exit_span!(span); - - // Extract rounds j = epoch_size .. min(2*epoch_size - 1, n) - let span = entered_span!("compute_rounds_from_m", epoch = epoch_size); - for j in epoch_size..(2 * epoch_size).min(n + 1) { - let d = j - epoch_size; // intra-epoch offset - let intra_challenges = challenges[epoch_size - 1..epoch_size - 1 + d].to_vec(); - - let [_h0, h1, h2] = compute_round_from_m(&m_table, epoch_size, &intra_challenges); - - // Append [h(1), h(2)] to transcript and sample challenge. - transcript.append_field_element_ext(&h1); - transcript.append_field_element_ext(&h2); - let challenge = transcript - .sample_and_append_challenge(b"Internal round") - .elements; - - proof_messages.push(IOPProverMessage { - evaluations: vec![h1, h2], - }); - challenges.push(challenge); - } - exit_span!(span); - } - - // --- Phase 2: Bind and materialize, then standard sumcheck --- - let k = challenges.len(); // actual number of streaming rounds completed - if k < n { - let span = entered_span!("bind_and_materialize"); - let (q_bound, f_bound) = bind_and_materialize(input, &challenges); - exit_span!(span); - - let remaining_vars = n - k; - let q_mle = MultilinearExtension::from_evaluations_ext_vec(remaining_vars, q_bound); - let f_mle = MultilinearExtension::from_evaluations_ext_vec(remaining_vars, f_bound); - - // Use VirtualPolynomial + round-by-round proving (no extra transcript header). - use multilinear_extensions::virtual_poly::VirtualPolynomial; - use std::sync::Arc; - let q_arc = Arc::new(q_mle); - let f_arc = Arc::new(f_mle); - let vp = VirtualPolynomial::new_from_product(vec![q_arc, f_arc], E::ONE); - - let span = entered_span!("standard_sumcheck", rounds = remaining_vars); - let mut prover_state = - IOPProverState::prover_init_with_extrapolation_aux(true, vp, None, None); - let mut challenge = None; - for _ in 0..remaining_vars { - let prover_msg = - IOPProverState::prove_round_and_update_state(&mut prover_state, &challenge); - prover_msg - .evaluations - .iter() - .for_each(|e| transcript.append_field_element_ext(e)); - challenge = Some(transcript.sample_and_append_challenge(b"Internal round")); - challenges.push(challenge.unwrap().elements); - proof_messages.push(prover_msg); - } - exit_span!(span); - } - - ( - IOPProof { - proofs: proof_messages, - }, - challenges, - ) -} - -/// Build M-table for epoch j'. -/// -/// M[beta1 * 2^{j'} + beta2] = sum_b Q_bound(beta1, b) * F_bound(beta2, b) -/// -/// where: -/// - Q_bound(beta, b) = sum_{a in {0,1}^{j'-1}} eq(R, a) * q'[a || beta || b] -/// - F_bound(beta, b) = sum_{a in {0,1}^{j'-1}} eq(R, a) * f[a || beta || b] -fn build_m_table( - input: &JaggedSumcheckInput, - challenges: &[E], // R_{j'} = (r_1, ..., r_{j'-1}) - epoch_size: usize, // j' -) -> Vec { - let n = input.num_giga_vars; - let bound_vars = epoch_size - 1; // j' - 1 - - let eq_r = if bound_vars > 0 { - build_eq_x_r_vec(&challenges[..bound_vars]) - } else { - vec![E::ONE] - }; - - let beta_count = 1usize << epoch_size; // 2^{j'} - let a_count = 1usize << bound_vars; // 2^{j'-1} - let chunk_size = a_count * beta_count; // 2^{2j'-1} - // When n < 2j'-1, all variables fit in a single chunk (no b dimension). - let n_chunks = 1usize << n.saturating_sub(2 * epoch_size - 1); // 2^{max(0, n - 2j' + 1)} - - let m_size = beta_count * beta_count; // 2^{2j'} - - // Step 1: Each thread processes a batch of b-chunks, producing a local M-table. - let span = entered_span!( - "streaming_pass", - n_chunks = n_chunks, - beta_count = beta_count - ); - let indices: Vec = (0..n_chunks).collect(); - let n_threads = max_usable_threads(); - let batch_size = (n_chunks / n_threads).max(1); - let partial_tables: Vec> = indices - .par_chunks(batch_size) - .map(|batch| { - let mut local_m = vec![E::ZERO; m_size]; - let mut q_bound = vec![E::ZERO; beta_count]; - let mut f_bound = vec![E::ZERO; beta_count]; - - for &b_idx in batch { - let chunk_start = b_idx * chunk_size; - q_bound - .iter_mut() - .zip(f_bound.iter_mut()) - .enumerate() - .for_each(|(beta, (q_b, f_b))| { - let base = chunk_start + beta * a_count; - let (q_acc, f_acc) = eq_r - .iter() - .zip(input.q_evals.get(base..).unwrap_or(&[])) - .zip(input.col_row_iter(base)) - .fold( - (E::ZERO, E::ZERO), - |(q_acc, f_acc), ((&eq_r_a, &q), (col, row))| { - ( - q_acc + eq_r_a * q, - f_acc + eq_r_a * (input.eq_row[row] * input.eq_col[col]), - ) - }, - ); - *q_b = q_acc; - *f_b = f_acc; - }); - - // Outer product accumulation into local M-table. - for b1 in 0..beta_count { - if q_bound[b1] == E::ZERO { - continue; - } - for b2 in 0..beta_count { - local_m[b1 * beta_count + b2] += q_bound[b1] * f_bound[b2]; - } - } - } - local_m - }) - .collect(); - exit_span!(span); - - // Step 2: Sum partial M-tables in parallel, each thread handles a slice of cells. - let span = entered_span!("reduce_partial_tables", n_partials = partial_tables.len()); - let n_partials = partial_tables.len(); - if n_partials == 0 { - exit_span!(span); - return vec![E::ZERO; m_size]; - } - let mut m_table = partial_tables[0].clone(); - let cell_batch = (m_size / n_threads).max(1); - m_table - .par_chunks_mut(cell_batch) - .enumerate() - .for_each(|(ci, cells)| { - let start = ci * cell_batch; - for partial in &partial_tables[1..] { - for (j, cell) in cells.iter_mut().enumerate() { - *cell += partial[start + j]; - } - } - }); - exit_span!(span); - m_table -} - -/// Extract round univariate h_j(x) from M-table. -/// -/// For round j in epoch j', d = j - j' intra-epoch challenges have been collected. -/// Returns [h(0), h(1), h(2)]. -fn compute_round_from_m( - m_table: &[E], - epoch_size: usize, // j' - intra_challenges: &[E], // r_{j'}, ..., r_{j-1} (d elements) -) -> [E; 3] { - let d = intra_challenges.len(); - let beta_count = 1usize << epoch_size; - let pad_bits = epoch_size - d - 1; // number of "future" bits to sum over - let pad_count = 1usize << pad_bits; - - let eq_intra = if d > 0 { - build_eq_x_r_vec(intra_challenges) - } else { - vec![E::ONE] - }; - - let a_count = 1usize << d; // 2^d - - let mut h = [E::ZERO; 3]; - - for a in 0..a_count { - for c in 0..a_count { - let eq_weight = eq_intra[a] * eq_intra[c]; - if eq_weight == E::ZERO { - continue; - } - - // Sum over all pad bit assignments (same pad for both beta1/beta2 - // since they correspond to the same physical "future" variables). - for p in 0..pad_count { - // beta = a_bits || x_bit || pad_bits (little-endian) - // beta_val = a + x_bit * 2^d + pad * 2^{d+1} - let base1 = a + (p << (d + 1)); - let base2 = c + (p << (d + 1)); - let b1_0 = base1; // x=0 - let b1_1 = base1 + (1 << d); // x=1 - let b2_0 = base2; - let b2_1 = base2 + (1 << d); - - let m00 = m_table[b1_0 * beta_count + b2_0]; - let m10 = m_table[b1_1 * beta_count + b2_0]; - let m01 = m_table[b1_0 * beta_count + b2_1]; - let m11 = m_table[b1_1 * beta_count + b2_1]; - - h[0] += eq_weight * m00; - h[1] += eq_weight * m11; - // h(2) via bilinear: (1-2)^2*M00 + 2(1-2)*M10 + (1-2)*2*M01 + 4*M11 - h[2] += eq_weight * (m00 - m10.double() - m01.double() + m11.double().double()); - } - } - } - - h -} - -/// Bind first K variables and materialize reduced q' and f as extension-field vectors. -/// -/// q_bound[idx] = sum_{a in {0,1}^K} eq(R, a) * q'[a + idx * 2^K] -/// f_bound[idx] = sum_{a in {0,1}^K} eq(R, a) * f[a + idx * 2^K] -fn bind_and_materialize( - input: &JaggedSumcheckInput, - challenges: &[E], // R_K = (r_1, ..., r_K) -) -> (Vec, Vec) { - let n = input.num_giga_vars; - let k = challenges.len(); - let remaining_size = 1usize << (n - k); - let a_count = 1usize << k; - - let eq_r = build_eq_x_r_vec(challenges); - - // Each output index is independent — parallelize over idx. - let results: Vec<(E, E)> = (0..remaining_size) - .into_par_iter() - .map(|idx| { - let base = idx * a_count; - eq_r.iter() - .zip(input.q_evals.get(base..).unwrap_or(&[])) - .zip(input.col_row_iter(base)) - .fold( - (E::ZERO, E::ZERO), - |(q_acc, f_acc), ((&eq_r_a, &q), (col, row))| { - ( - q_acc + eq_r_a * q, - f_acc + eq_r_a * (input.eq_row[row] * input.eq_col[col]), - ) - }, - ) - }) - .collect(); - - results.into_iter().unzip() -} - // --------------------------------------------------------------------------- // Jagged Batch Open / Verify // --------------------------------------------------------------------------- @@ -965,6 +538,7 @@ mod tests { test_util::setup_pcs, }; use ff_ext::GoldilocksExt2; + use multilinear_extensions::mle::MultilinearExtension; use p3::{field::FieldAlgebra, goldilocks::Goldilocks}; type F = Goldilocks; @@ -1054,194 +628,11 @@ mod tests { assert!(matches!(result, Err(Error::InvalidPcsParam(_)))); } - // --- Sumcheck tests --- - - use multilinear_extensions::virtual_poly::{VPAuxInfo, build_eq_x_r_vec}; - use rand::thread_rng; - use std::marker::PhantomData; - use sumcheck::structs::IOPVerifierState; - use transcript::basic::BasicTranscript; - - #[test] - fn test_jagged_sumcheck_small() { - use ff_ext::FromUniformBytes; - - let mut rng = thread_rng(); - - // 3 polynomials of height 4 (s=2), total 12 evals, padded to 16 (n=4). - let num_polys = 3usize; - let poly_height = 4usize; - let s = 2; // log2(poly_height) - let total_evals = num_polys * poly_height; - let num_giga_vars = 4; // ceil(log2(12)) = 4, 2^4 = 16 - - let q_evals: Vec = (0..total_evals) - .map(|i| F::from_canonical_u64(i as u64 + 1)) - .collect(); - - let cumulative_heights: Vec = (0..=num_polys).map(|i| i * poly_height).collect(); - - let z_row: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); - // z_col needs ceil(log2(num_polys)) = 2 variables - let z_col: Vec = (0..2).map(|_| E::random(&mut rng)).collect(); - - let input = JaggedSumcheckInput { - q_evals: &q_evals, - num_giga_vars, - cumulative_heights: &cumulative_heights, - eq_row: build_eq_x_r_vec(&z_row), - eq_col: build_eq_x_r_vec(&z_col), - }; - - let claimed_sum = input.compute_claimed_sum(); - - let mut transcript = BasicTranscript::::new(b"jagged_sumcheck_test"); - let (proof, challenges) = jagged_sumcheck_prove(&input, &mut transcript, None); - - assert_eq!(proof.proofs.len(), num_giga_vars); - assert_eq!(challenges.len(), num_giga_vars); - - // Verify using the standard sumcheck verifier. - let mut transcript_v = BasicTranscript::::new(b"jagged_sumcheck_test"); - let aux_info = VPAuxInfo { - max_degree: 2, - max_num_variables: num_giga_vars, - phantom: PhantomData::, - }; - let subclaim = - IOPVerifierState::::verify(claimed_sum, &proof, &aux_info, &mut transcript_v); - - // The subclaim point should match our challenges. - for (sc, ch) in subclaim.point.iter().zip(challenges.iter()) { - assert_eq!(sc.elements, *ch); - } - - // Verify the final evaluation: q'(point) * f(point) == expected_evaluation - let (q_at_point, f_at_point) = input.final_evaluations_slow(&challenges); - assert_eq!( - q_at_point * f_at_point, - subclaim.expected_evaluation, - "final evaluation mismatch" - ); - } - - #[test] - fn test_jagged_sumcheck_all_epochs() { - // n=16: exercises all 4 epochs (j'=1,2,4,8) + 1 round of standard sumcheck. - use ff_ext::FromUniformBytes; - - let mut rng = thread_rng(); - - let num_polys = 8usize; - let poly_height = 1 << 13; // 8192, s=13 - let s = 13; - let total_evals = num_polys * poly_height; // 65536 - let num_giga_vars = 16; // 2^16 = 65536 - - let q_evals: Vec = (0..total_evals) - .map(|i| F::from_canonical_u64((i as u64 * 7 + 3) % (1 << 20))) - .collect(); - - let cumulative_heights: Vec = (0..=num_polys).map(|i| i * poly_height).collect(); - - let z_row: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); - let z_col: Vec = (0..3).map(|_| E::random(&mut rng)).collect(); // ceil(log2(8))=3 - - let input = JaggedSumcheckInput { - q_evals: &q_evals, - num_giga_vars, - cumulative_heights: &cumulative_heights, - eq_row: build_eq_x_r_vec(&z_row), - eq_col: build_eq_x_r_vec(&z_col), - }; - - let claimed_sum = input.compute_claimed_sum(); - - let mut transcript = BasicTranscript::::new(b"jagged_test_16"); - let (proof, challenges) = jagged_sumcheck_prove(&input, &mut transcript, None); - - assert_eq!(proof.proofs.len(), num_giga_vars); - assert_eq!(challenges.len(), num_giga_vars); - - let mut transcript_v = BasicTranscript::::new(b"jagged_test_16"); - let aux_info = VPAuxInfo { - max_degree: 2, - max_num_variables: num_giga_vars, - phantom: PhantomData::, - }; - let subclaim = - IOPVerifierState::::verify(claimed_sum, &proof, &aux_info, &mut transcript_v); - - for (sc, ch) in subclaim.point.iter().zip(challenges.iter()) { - assert_eq!(sc.elements, *ch); - } - } - - #[test] - fn test_jagged_sumcheck_n25() { - // n=25: 2^25 = 33M evaluations. Exercises all epochs + 10 rounds of standard sumcheck. - use ff_ext::FromUniformBytes; - - tracing_forest::init(); - - let mut rng = thread_rng(); - - let num_polys = 1 << 10; // 1024 polynomials - let poly_height = 1 << 15; // 32768 each, s=15 - let s = 15; - let total_evals = num_polys * poly_height; // 2^25 = 33554432 - let num_giga_vars = 25; - - let q_evals: Vec = (0..total_evals) - .map(|i| F::from_canonical_u64((i as u64 * 13 + 7) % (1 << 30))) - .collect(); - - let cumulative_heights: Vec = (0..=num_polys).map(|i| i * poly_height).collect(); - - let z_row: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); - let z_col: Vec = (0..10).map(|_| E::random(&mut rng)).collect(); // ceil(log2(1024))=10 - - let input = JaggedSumcheckInput { - q_evals: &q_evals, - num_giga_vars, - cumulative_heights: &cumulative_heights, - eq_row: build_eq_x_r_vec(&z_row), - eq_col: build_eq_x_r_vec(&z_col), - }; - - let claimed_sum = input.compute_claimed_sum(); - - let mut transcript = BasicTranscript::::new(b"jagged_test_25"); - let (proof, challenges) = jagged_sumcheck_prove(&input, &mut transcript, None); - - assert_eq!(proof.proofs.len(), num_giga_vars); - assert_eq!(challenges.len(), num_giga_vars); - - let mut transcript_v = BasicTranscript::::new(b"jagged_test_25"); - let aux_info = VPAuxInfo { - max_degree: 2, - max_num_variables: num_giga_vars, - phantom: PhantomData::, - }; - let subclaim = - IOPVerifierState::::verify(claimed_sum, &proof, &aux_info, &mut transcript_v); - - for (sc, ch) in subclaim.point.iter().zip(challenges.iter()) { - assert_eq!(sc.elements, *ch); - } - - // Verify the final evaluation: q'(point) * f(point) == expected_evaluation - let (q_at_point, f_at_point) = input.final_evaluations_slow(&challenges); - assert_eq!( - q_at_point * f_at_point, - subclaim.expected_evaluation, - "final evaluation mismatch" - ); - } - // --- Batch open/verify tests --- use ff_ext::FromUniformBytes; + use rand::thread_rng; + use transcript::basic::BasicTranscript; /// Evaluate a single column polynomial at `point` (the original, non-bit-reversed poly). /// `col_evals` are the raw evaluations of the column (before bit-reversal). diff --git a/crates/mpcs/src/jagged/sumcheck.rs b/crates/mpcs/src/jagged/sumcheck.rs new file mode 100644 index 0000000..4c5b57f --- /dev/null +++ b/crates/mpcs/src/jagged/sumcheck.rs @@ -0,0 +1,623 @@ +//! Jagged Sumcheck Prover +//! +//! Streaming sumcheck prover using the M-table algorithm from +//! "Time-Space Trade-Offs for Sumcheck" (eprint 2025/1473), Section 4. + +use ff_ext::ExtensionField; +use multilinear_extensions::{ + mle::MultilinearExtension, util::max_usable_threads, virtual_poly::build_eq_x_r_vec, +}; +use p3::maybe_rayon::prelude::{ + IntoParallelIterator, ParallelIterator, ParallelSlice, ParallelSliceMut, +}; +use sumcheck::{ + macros::{entered_span, exit_span}, + structs::{IOPProof, IOPProverMessage, IOPProverState}, +}; +use transcript::Transcript; + +/// Default log2 of the maximum epoch size for the streaming phase. +/// Epoch sizes are `[2^0, 2^1, ..., 2^LOG2_MAX_EPOCH]` = `[1, 2, 4, 8]`, +/// covering 1 + 2 + 4 + 8 = 15 streaming rounds before switching to standard sumcheck. +const LOG2_MAX_EPOCH: u32 = 3; + +/// All inputs needed for the jagged sumcheck. +pub struct JaggedSumcheckInput<'a, E: ExtensionField> { + /// Giga polynomial evaluations (concatenated, bit-reversed). + pub q_evals: &'a [E::BaseField], + /// n = log2(padded_total_size). + pub num_giga_vars: usize, + /// Cumulative height sequence t[j], length num_polys + 1. + pub cumulative_heights: &'a [usize], + /// Precomputed eq table for the row evaluation point: `build_eq_x_r_vec(z_row)`. + pub eq_row: Vec, + /// Precomputed eq table for the column challenge point: `build_eq_x_r_vec(z_col)`. + pub eq_col: Vec, +} + +/// Iterator that yields `(col, row)` pairs for consecutive giga indices. +/// Uses one binary search at construction, then O(1) per step. +struct ColRowIter<'a> { + cumulative_heights: &'a [usize], + col: usize, + row: usize, + num_polys: usize, +} + +impl<'a> Iterator for ColRowIter<'a> { + type Item = (usize, usize); + + fn next(&mut self) -> Option<(usize, usize)> { + if self.col >= self.num_polys { + return None; + } + let result = (self.col, self.row); + self.row += 1; + let poly_height = self.cumulative_heights[self.col + 1] - self.cumulative_heights[self.col]; + if self.row >= poly_height { + self.row = 0; + self.col += 1; + } + Some(result) + } +} + +impl<'a, E: ExtensionField> JaggedSumcheckInput<'a, E> { + fn total_evaluations(&self) -> usize { + *self.cumulative_heights.last().unwrap_or(&0) + } + + /// Return an iterator yielding `(col, row)` for consecutive giga indices + /// starting from `start`. One binary search at construction, O(1) per step. + fn col_row_iter(&self, start: usize) -> ColRowIter<'_> { + let num_polys = self.cumulative_heights.len() - 1; + if start >= self.total_evaluations() { + return ColRowIter { + cumulative_heights: self.cumulative_heights, + col: num_polys, + row: 0, + num_polys, + }; + } + let j = self.cumulative_heights.partition_point(|&t| t <= start) - 1; + ColRowIter { + cumulative_heights: self.cumulative_heights, + col: j, + row: start - self.cumulative_heights[j], + num_polys, + } + } + + /// Brute-force computation of sum_b q'(b) * f(b). + /// O(2^n) time — only for debugging and tests. + #[cfg(test)] + fn compute_claimed_sum(&self) -> E { + self.q_evals[..self.total_evaluations()] + .iter() + .zip(self.col_row_iter(0)) + .fold(E::ZERO, |sum, (&q, (col, row))| { + sum + (self.eq_row[row] * self.eq_col[col]) * q + }) + } + + /// Brute-force MLE evaluation of q'(rb) and f(rb) at the given point. + /// O(2^n) time and memory — only for debugging and tests. + /// Returns `(q_at_point, f_at_point)`. + #[cfg(test)] + fn final_evaluations_slow(&self, point: &[E]) -> (E, E) { + let n = self.num_giga_vars; + let total_evals = self.total_evaluations(); + + // Build q' MLE (padded with zeros) and evaluate at point. + let mut q_padded: Vec = self.q_evals.to_vec(); + q_padded.resize(1 << n, Default::default()); + let q_mle = MultilinearExtension::from_evaluations_vec(n, q_padded); + let q_at_point = q_mle.evaluate(point); + + // Build f MLE from eq tables and evaluate at point. + let mut f_evals = vec![E::ZERO; 1 << n]; + for (f_eval, (col, row)) in f_evals[..total_evals].iter_mut().zip(self.col_row_iter(0)) { + *f_eval = self.eq_row[row] * self.eq_col[col]; + } + let f_mle = MultilinearExtension::from_evaluations_ext_vec(n, f_evals); + let f_at_point = f_mle.evaluate(point); + + (q_at_point, f_at_point) + } +} + +/// Run the full jagged sumcheck: streaming phase (rounds 1..K) + standard phase (rounds K+1..n). +/// +/// `log2_max_epoch` controls the streaming phase epoch schedule. Epoch sizes are +/// `[1, 2, 4, ..., 2^k]` where `k = log2_max_epoch`. Pass `None` to use the default +/// `LOG2_MAX_EPOCH = 3`, giving epoch sizes `[1, 2, 4, 8]` and 15 streaming rounds. +/// +/// Returns the proof and the full list of challenges (r_1, ..., r_n). +pub fn jagged_sumcheck_prove( + input: &JaggedSumcheckInput, + transcript: &mut impl Transcript, + log2_max_epoch: Option, +) -> (IOPProof, Vec) { + let n = input.num_giga_vars; + let max_degree: usize = 2; + + let k = log2_max_epoch.unwrap_or(LOG2_MAX_EPOCH); + let epoch_sizes: Vec = (0..=k).map(|i| 1usize << i).collect(); + + let mut challenges: Vec = Vec::with_capacity(n); + let mut proof_messages: Vec> = Vec::with_capacity(n); + + // Write transcript header (must match verifier's expectations). + transcript.append_message(&n.to_le_bytes()); + transcript.append_message(&max_degree.to_le_bytes()); + + // --- Streaming phase: epochs j' = 1, 2, 4, ..., 2^k --- + for &epoch_size in &epoch_sizes { + // Epoch j' handles rounds j'..2j'-1. Skip if all rounds are done. + if epoch_size > n { + break; + } + + // Build M-table for this epoch. + let span = entered_span!("build_m_table", epoch = epoch_size); + let m_table = build_m_table(input, &challenges, epoch_size); + exit_span!(span); + + // Extract rounds j = epoch_size .. min(2*epoch_size - 1, n) + let span = entered_span!("compute_rounds_from_m", epoch = epoch_size); + for j in epoch_size..(2 * epoch_size).min(n + 1) { + let d = j - epoch_size; // intra-epoch offset + let intra_challenges = challenges[epoch_size - 1..epoch_size - 1 + d].to_vec(); + + let [_h0, h1, h2] = compute_round_from_m(&m_table, epoch_size, &intra_challenges); + + // Append [h(1), h(2)] to transcript and sample challenge. + transcript.append_field_element_ext(&h1); + transcript.append_field_element_ext(&h2); + let challenge = transcript + .sample_and_append_challenge(b"Internal round") + .elements; + + proof_messages.push(IOPProverMessage { + evaluations: vec![h1, h2], + }); + challenges.push(challenge); + } + exit_span!(span); + } + + // --- Phase 2: Bind and materialize, then standard sumcheck --- + let k = challenges.len(); // actual number of streaming rounds completed + if k < n { + let span = entered_span!("bind_and_materialize"); + let (q_bound, f_bound) = bind_and_materialize(input, &challenges); + exit_span!(span); + + let remaining_vars = n - k; + let q_mle = MultilinearExtension::from_evaluations_ext_vec(remaining_vars, q_bound); + let f_mle = MultilinearExtension::from_evaluations_ext_vec(remaining_vars, f_bound); + + // Use VirtualPolynomial + round-by-round proving (no extra transcript header). + use multilinear_extensions::virtual_poly::VirtualPolynomial; + use std::sync::Arc; + let q_arc = Arc::new(q_mle); + let f_arc = Arc::new(f_mle); + let vp = VirtualPolynomial::new_from_product(vec![q_arc, f_arc], E::ONE); + + let span = entered_span!("standard_sumcheck", rounds = remaining_vars); + let mut prover_state = + IOPProverState::prover_init_with_extrapolation_aux(true, vp, None, None); + let mut challenge = None; + for _ in 0..remaining_vars { + let prover_msg = + IOPProverState::prove_round_and_update_state(&mut prover_state, &challenge); + prover_msg + .evaluations + .iter() + .for_each(|e| transcript.append_field_element_ext(e)); + challenge = Some(transcript.sample_and_append_challenge(b"Internal round")); + challenges.push(challenge.unwrap().elements); + proof_messages.push(prover_msg); + } + exit_span!(span); + } + + ( + IOPProof { + proofs: proof_messages, + }, + challenges, + ) +} + +/// Build M-table for epoch j'. +/// +/// M[beta1 * 2^{j'} + beta2] = sum_b Q_bound(beta1, b) * F_bound(beta2, b) +/// +/// where: +/// - Q_bound(beta, b) = sum_{a in {0,1}^{j'-1}} eq(R, a) * q'[a || beta || b] +/// - F_bound(beta, b) = sum_{a in {0,1}^{j'-1}} eq(R, a) * f[a || beta || b] +fn build_m_table( + input: &JaggedSumcheckInput, + challenges: &[E], // R_{j'} = (r_1, ..., r_{j'-1}) + epoch_size: usize, // j' +) -> Vec { + let n = input.num_giga_vars; + let bound_vars = epoch_size - 1; // j' - 1 + + let eq_r = if bound_vars > 0 { + build_eq_x_r_vec(&challenges[..bound_vars]) + } else { + vec![E::ONE] + }; + + let beta_count = 1usize << epoch_size; // 2^{j'} + let a_count = 1usize << bound_vars; // 2^{j'-1} + let chunk_size = a_count * beta_count; // 2^{2j'-1} + // When n < 2j'-1, all variables fit in a single chunk (no b dimension). + let n_chunks = 1usize << n.saturating_sub(2 * epoch_size - 1); // 2^{max(0, n - 2j' + 1)} + + let m_size = beta_count * beta_count; // 2^{2j'} + + // Step 1: Each thread processes a batch of b-chunks, producing a local M-table. + let span = entered_span!( + "streaming_pass", + n_chunks = n_chunks, + beta_count = beta_count + ); + let indices: Vec = (0..n_chunks).collect(); + let n_threads = max_usable_threads(); + let batch_size = (n_chunks / n_threads).max(1); + let partial_tables: Vec> = indices + .par_chunks(batch_size) + .map(|batch| { + let mut local_m = vec![E::ZERO; m_size]; + let mut q_bound = vec![E::ZERO; beta_count]; + let mut f_bound = vec![E::ZERO; beta_count]; + + for &b_idx in batch { + let chunk_start = b_idx * chunk_size; + q_bound + .iter_mut() + .zip(f_bound.iter_mut()) + .enumerate() + .for_each(|(beta, (q_b, f_b))| { + let base = chunk_start + beta * a_count; + let (q_acc, f_acc) = eq_r + .iter() + .zip(input.q_evals.get(base..).unwrap_or(&[])) + .zip(input.col_row_iter(base)) + .fold( + (E::ZERO, E::ZERO), + |(q_acc, f_acc), ((&eq_r_a, &q), (col, row))| { + ( + q_acc + eq_r_a * q, + f_acc + eq_r_a * (input.eq_row[row] * input.eq_col[col]), + ) + }, + ); + *q_b = q_acc; + *f_b = f_acc; + }); + + // Outer product accumulation into local M-table. + for b1 in 0..beta_count { + if q_bound[b1] == E::ZERO { + continue; + } + for b2 in 0..beta_count { + local_m[b1 * beta_count + b2] += q_bound[b1] * f_bound[b2]; + } + } + } + local_m + }) + .collect(); + exit_span!(span); + + // Step 2: Sum partial M-tables in parallel, each thread handles a slice of cells. + let span = entered_span!("reduce_partial_tables", n_partials = partial_tables.len()); + let n_partials = partial_tables.len(); + if n_partials == 0 { + exit_span!(span); + return vec![E::ZERO; m_size]; + } + let mut m_table = partial_tables[0].clone(); + let cell_batch = (m_size / n_threads).max(1); + m_table + .par_chunks_mut(cell_batch) + .enumerate() + .for_each(|(ci, cells)| { + let start = ci * cell_batch; + for partial in &partial_tables[1..] { + for (j, cell) in cells.iter_mut().enumerate() { + *cell += partial[start + j]; + } + } + }); + exit_span!(span); + m_table +} + +/// Extract round univariate h_j(x) from M-table. +/// +/// For round j in epoch j', d = j - j' intra-epoch challenges have been collected. +/// Returns [h(0), h(1), h(2)]. +fn compute_round_from_m( + m_table: &[E], + epoch_size: usize, // j' + intra_challenges: &[E], // r_{j'}, ..., r_{j-1} (d elements) +) -> [E; 3] { + let d = intra_challenges.len(); + let beta_count = 1usize << epoch_size; + let pad_bits = epoch_size - d - 1; // number of "future" bits to sum over + let pad_count = 1usize << pad_bits; + + let eq_intra = if d > 0 { + build_eq_x_r_vec(intra_challenges) + } else { + vec![E::ONE] + }; + + let a_count = 1usize << d; // 2^d + + let mut h = [E::ZERO; 3]; + + for a in 0..a_count { + for c in 0..a_count { + let eq_weight = eq_intra[a] * eq_intra[c]; + if eq_weight == E::ZERO { + continue; + } + + // Sum over all pad bit assignments (same pad for both beta1/beta2 + // since they correspond to the same physical "future" variables). + for p in 0..pad_count { + // beta = a_bits || x_bit || pad_bits (little-endian) + // beta_val = a + x_bit * 2^d + pad * 2^{d+1} + let base1 = a + (p << (d + 1)); + let base2 = c + (p << (d + 1)); + let b1_0 = base1; // x=0 + let b1_1 = base1 + (1 << d); // x=1 + let b2_0 = base2; + let b2_1 = base2 + (1 << d); + + let m00 = m_table[b1_0 * beta_count + b2_0]; + let m10 = m_table[b1_1 * beta_count + b2_0]; + let m01 = m_table[b1_0 * beta_count + b2_1]; + let m11 = m_table[b1_1 * beta_count + b2_1]; + + h[0] += eq_weight * m00; + h[1] += eq_weight * m11; + // h(2) via bilinear: (1-2)^2*M00 + 2(1-2)*M10 + (1-2)*2*M01 + 4*M11 + h[2] += eq_weight * (m00 - m10.double() - m01.double() + m11.double().double()); + } + } + } + + h +} + +/// Bind first K variables and materialize reduced q' and f as extension-field vectors. +/// +/// q_bound[idx] = sum_{a in {0,1}^K} eq(R, a) * q'[a + idx * 2^K] +/// f_bound[idx] = sum_{a in {0,1}^K} eq(R, a) * f[a + idx * 2^K] +fn bind_and_materialize( + input: &JaggedSumcheckInput, + challenges: &[E], // R_K = (r_1, ..., r_K) +) -> (Vec, Vec) { + let n = input.num_giga_vars; + let k = challenges.len(); + let remaining_size = 1usize << (n - k); + let a_count = 1usize << k; + + let eq_r = build_eq_x_r_vec(challenges); + + // Each output index is independent — parallelize over idx. + let results: Vec<(E, E)> = (0..remaining_size) + .into_par_iter() + .map(|idx| { + let base = idx * a_count; + eq_r.iter() + .zip(input.q_evals.get(base..).unwrap_or(&[])) + .zip(input.col_row_iter(base)) + .fold( + (E::ZERO, E::ZERO), + |(q_acc, f_acc), ((&eq_r_a, &q), (col, row))| { + ( + q_acc + eq_r_a * q, + f_acc + eq_r_a * (input.eq_row[row] * input.eq_col[col]), + ) + }, + ) + }) + .collect(); + + results.into_iter().unzip() +} + +#[cfg(test)] +mod tests { + use super::*; + use ff_ext::{FromUniformBytes, GoldilocksExt2}; + use multilinear_extensions::virtual_poly::{VPAuxInfo, build_eq_x_r_vec}; + use p3::{field::FieldAlgebra, goldilocks::Goldilocks}; + use rand::thread_rng; + use std::marker::PhantomData; + use sumcheck::structs::IOPVerifierState; + use transcript::basic::BasicTranscript; + + type F = Goldilocks; + type E = GoldilocksExt2; + + #[test] + fn test_jagged_sumcheck_small() { + let mut rng = thread_rng(); + + // 3 polynomials of height 4 (s=2), total 12 evals, padded to 16 (n=4). + let num_polys = 3usize; + let poly_height = 4usize; + let s = 2; // log2(poly_height) + let total_evals = num_polys * poly_height; + let num_giga_vars = 4; // ceil(log2(12)) = 4, 2^4 = 16 + + let q_evals: Vec = (0..total_evals) + .map(|i| F::from_canonical_u64(i as u64 + 1)) + .collect(); + + let cumulative_heights: Vec = (0..=num_polys).map(|i| i * poly_height).collect(); + + let z_row: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); + // z_col needs ceil(log2(num_polys)) = 2 variables + let z_col: Vec = (0..2).map(|_| E::random(&mut rng)).collect(); + + let input = JaggedSumcheckInput { + q_evals: &q_evals, + num_giga_vars, + cumulative_heights: &cumulative_heights, + eq_row: build_eq_x_r_vec(&z_row), + eq_col: build_eq_x_r_vec(&z_col), + }; + + let claimed_sum = input.compute_claimed_sum(); + + let mut transcript = BasicTranscript::::new(b"jagged_sumcheck_test"); + let (proof, challenges) = jagged_sumcheck_prove(&input, &mut transcript, None); + + assert_eq!(proof.proofs.len(), num_giga_vars); + assert_eq!(challenges.len(), num_giga_vars); + + // Verify using the standard sumcheck verifier. + let mut transcript_v = BasicTranscript::::new(b"jagged_sumcheck_test"); + let aux_info = VPAuxInfo { + max_degree: 2, + max_num_variables: num_giga_vars, + phantom: PhantomData::, + }; + let subclaim = + IOPVerifierState::::verify(claimed_sum, &proof, &aux_info, &mut transcript_v); + + // The subclaim point should match our challenges. + for (sc, ch) in subclaim.point.iter().zip(challenges.iter()) { + assert_eq!(sc.elements, *ch); + } + + // Verify the final evaluation: q'(point) * f(point) == expected_evaluation + let (q_at_point, f_at_point) = input.final_evaluations_slow(&challenges); + assert_eq!( + q_at_point * f_at_point, + subclaim.expected_evaluation, + "final evaluation mismatch" + ); + } + + #[test] + fn test_jagged_sumcheck_all_epochs() { + // n=16: exercises all 4 epochs (j'=1,2,4,8) + 1 round of standard sumcheck. + let mut rng = thread_rng(); + + let num_polys = 8usize; + let poly_height = 1 << 13; // 8192, s=13 + let s = 13; + let total_evals = num_polys * poly_height; // 65536 + let num_giga_vars = 16; // 2^16 = 65536 + + let q_evals: Vec = (0..total_evals) + .map(|i| F::from_canonical_u64((i as u64 * 7 + 3) % (1 << 20))) + .collect(); + + let cumulative_heights: Vec = (0..=num_polys).map(|i| i * poly_height).collect(); + + let z_row: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); + let z_col: Vec = (0..3).map(|_| E::random(&mut rng)).collect(); // ceil(log2(8))=3 + + let input = JaggedSumcheckInput { + q_evals: &q_evals, + num_giga_vars, + cumulative_heights: &cumulative_heights, + eq_row: build_eq_x_r_vec(&z_row), + eq_col: build_eq_x_r_vec(&z_col), + }; + + let claimed_sum = input.compute_claimed_sum(); + + let mut transcript = BasicTranscript::::new(b"jagged_test_16"); + let (proof, challenges) = jagged_sumcheck_prove(&input, &mut transcript, None); + + assert_eq!(proof.proofs.len(), num_giga_vars); + assert_eq!(challenges.len(), num_giga_vars); + + let mut transcript_v = BasicTranscript::::new(b"jagged_test_16"); + let aux_info = VPAuxInfo { + max_degree: 2, + max_num_variables: num_giga_vars, + phantom: PhantomData::, + }; + let subclaim = + IOPVerifierState::::verify(claimed_sum, &proof, &aux_info, &mut transcript_v); + + for (sc, ch) in subclaim.point.iter().zip(challenges.iter()) { + assert_eq!(sc.elements, *ch); + } + } + + #[test] + fn test_jagged_sumcheck_n25() { + // n=25: 2^25 = 33M evaluations. Exercises all epochs + 10 rounds of standard sumcheck. + tracing_forest::init(); + + let mut rng = thread_rng(); + + let num_polys = 1 << 10; // 1024 polynomials + let poly_height = 1 << 15; // 32768 each, s=15 + let s = 15; + let total_evals = num_polys * poly_height; // 2^25 = 33554432 + let num_giga_vars = 25; + + let q_evals: Vec = (0..total_evals) + .map(|i| F::from_canonical_u64((i as u64 * 13 + 7) % (1 << 30))) + .collect(); + + let cumulative_heights: Vec = (0..=num_polys).map(|i| i * poly_height).collect(); + + let z_row: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); + let z_col: Vec = (0..10).map(|_| E::random(&mut rng)).collect(); // ceil(log2(1024))=10 + + let input = JaggedSumcheckInput { + q_evals: &q_evals, + num_giga_vars, + cumulative_heights: &cumulative_heights, + eq_row: build_eq_x_r_vec(&z_row), + eq_col: build_eq_x_r_vec(&z_col), + }; + + let claimed_sum = input.compute_claimed_sum(); + + let mut transcript = BasicTranscript::::new(b"jagged_test_25"); + let (proof, challenges) = jagged_sumcheck_prove(&input, &mut transcript, None); + + assert_eq!(proof.proofs.len(), num_giga_vars); + assert_eq!(challenges.len(), num_giga_vars); + + let mut transcript_v = BasicTranscript::::new(b"jagged_test_25"); + let aux_info = VPAuxInfo { + max_degree: 2, + max_num_variables: num_giga_vars, + phantom: PhantomData::, + }; + let subclaim = + IOPVerifierState::::verify(claimed_sum, &proof, &aux_info, &mut transcript_v); + + for (sc, ch) in subclaim.point.iter().zip(challenges.iter()) { + assert_eq!(sc.elements, *ch); + } + + // Verify the final evaluation: q'(point) * f(point) == expected_evaluation + let (q_at_point, f_at_point) = input.final_evaluations_slow(&challenges); + assert_eq!( + q_at_point * f_at_point, + subclaim.expected_evaluation, + "final evaluation mismatch" + ); + } +} diff --git a/crates/mpcs/src/lib.rs b/crates/mpcs/src/lib.rs index 3c68f3a..cfbb5fb 100644 --- a/crates/mpcs/src/lib.rs +++ b/crates/mpcs/src/lib.rs @@ -270,10 +270,9 @@ pub use basefold::{ pub mod jagged; pub use jagged::{ JaggedBatchOpenProof, JaggedCommitment, JaggedCommitmentWithWitness, JaggedSumcheckInput, - jagged_batch_open, jagged_batch_verify, jagged_commit, jagged_sumcheck_prove, + evaluate_g, evaluate_g_backward, evaluate_g_forward, jagged_batch_open, jagged_batch_verify, + jagged_commit, jagged_sumcheck_prove, }; -pub mod jagged_evaluator; -pub use jagged_evaluator::{evaluate_g, evaluate_g_backward, evaluate_g_forward}; #[cfg(feature = "whir")] extern crate whir as whir_external; #[cfg(feature = "whir")] From a5e0b3e2897d1d5783617635e67b1055cc421925 Mon Sep 17 00:00:00 2001 From: kunxian xia Date: Fri, 24 Apr 2026 14:51:04 +0800 Subject: [PATCH 10/21] move types to a separate file --- crates/mpcs/src/jagged/mod.rs | 75 ++------------------------------- crates/mpcs/src/jagged/types.rs | 73 ++++++++++++++++++++++++++++++++ 2 files changed, 77 insertions(+), 71 deletions(-) create mode 100644 crates/mpcs/src/jagged/types.rs diff --git a/crates/mpcs/src/jagged/mod.rs b/crates/mpcs/src/jagged/mod.rs index bd201f8..60d19a7 100644 --- a/crates/mpcs/src/jagged/mod.rs +++ b/crates/mpcs/src/jagged/mod.rs @@ -102,14 +102,16 @@ pub mod evaluator; pub mod sumcheck; +mod types; pub use evaluator::{evaluate_g, evaluate_g_backward, evaluate_g_forward}; pub use sumcheck::{JaggedSumcheckInput, jagged_sumcheck_prove}; +pub use types::{JaggedBatchOpenProof, JaggedCommitment, JaggedCommitmentWithWitness}; use std::{iter::once, marker::PhantomData}; use crate::{Error, PolynomialCommitmentScheme}; -use ::sumcheck::structs::{IOPProof, IOPVerifierState}; +use ::sumcheck::structs::IOPVerifierState; use ff_ext::ExtensionField; use itertools::Itertools; use multilinear_extensions::{ @@ -120,60 +122,10 @@ use p3::{ matrix::{Matrix, bitrev::BitReversableMatrix}, maybe_rayon::prelude::{IntoParallelIterator, ParallelSliceMut}, }; -use serde::{Deserialize, Serialize}; use transcript::Transcript; +use types::int_to_field_bits; use witness::{InstancePaddingStrategy, RowMajorMatrix}; -/// Commitment to a jagged polynomial `q'`, together with all witness data needed -/// for opening proofs. -/// -/// Generic over the inner PCS `Pcs` so that any `PolynomialCommitmentScheme` can -/// serve as the underlying commitment engine. -pub struct JaggedCommitmentWithWitness> { - /// Commitment (with witness) to the "giga" polynomial `q'` via `Pcs`. - pub inner: Pcs::CommitmentWithWitness, - /// Cumulative height sequence `t`: - /// - `t[0] = 0` - /// - `t[i+1] = t[i] + poly_heights[i]` - /// - Length: `num_polys + 1` - pub cumulative_heights: Vec, - /// Number of evaluations `h_i = 2^(num_vars_i)` for each polynomial `p_i`. - /// Length: `num_polys`. - pub poly_heights: Vec, -} - -/// The pure commitment (without witness data) for a jagged polynomial `q'`. -/// This is what the verifier receives. -#[derive(Clone, Serialize, Deserialize)] -#[serde(bound(serialize = "", deserialize = ""))] -pub struct JaggedCommitment> { - /// Pure commitment to the underlying giga polynomial `q'`. - pub inner: Pcs::Commitment, - /// Cumulative height sequence `t` (verifier needs this to evaluate `f(b)`). - pub cumulative_heights: Vec, -} - -impl> JaggedCommitmentWithWitness { - /// Extract the pure commitment (without witness data). - pub fn to_commitment(&self) -> JaggedCommitment { - JaggedCommitment { - inner: Pcs::get_pure_commitment(&self.inner), - cumulative_heights: self.cumulative_heights.clone(), - } - } - - /// Total number of polynomials packed into `q'`. - pub fn num_polys(&self) -> usize { - self.poly_heights.len() - } - - /// Total number of evaluations in the *unpadded* concatenated polynomial - /// (= `t[num_polys]` = `cumulative_heights.last()`). - pub fn total_evaluations(&self) -> usize { - self.cumulative_heights.last().copied().unwrap_or(0) - } -} - /// Commit to a sequence of row-major matrices using the Jagged PCS scheme. /// /// This function implements the commit phase described in Ceno issue #1288: @@ -300,25 +252,6 @@ pub fn jagged_commit>( // Jagged Batch Open / Verify // --------------------------------------------------------------------------- -/// Proof for the jagged batch opening protocol. -/// -/// Contains a sumcheck proof (reducing K column evaluation claims to a single -/// point on q'), the evaluation q'(ρ), and an inner PCS opening proof. -#[derive(Clone, Serialize, Deserialize)] -#[serde(bound(serialize = "", deserialize = ""))] -pub struct JaggedBatchOpenProof> { - pub sumcheck_proof: IOPProof, - pub q_eval: E, - pub inner_proof: Pcs::Proof, -} - -/// Convert a `usize` to its little-endian binary representation as field elements. -fn int_to_field_bits(val: usize, num_bits: usize) -> Vec { - (0..num_bits) - .map(|i| if (val >> i) & 1 == 1 { E::ONE } else { E::ZERO }) - .collect() -} - /// Evaluate f̂(ρ) using the ROBP evaluator. /// /// f̂(ρ) = Σ_{y} eq_col[y] · ĝ(z_row_padded, ρ, bits(t_y), bits(t_{y+1})) diff --git a/crates/mpcs/src/jagged/types.rs b/crates/mpcs/src/jagged/types.rs new file mode 100644 index 0000000..0b5606c --- /dev/null +++ b/crates/mpcs/src/jagged/types.rs @@ -0,0 +1,73 @@ +use crate::PolynomialCommitmentScheme; +use ::sumcheck::structs::IOPProof; +use ff_ext::ExtensionField; +use serde::{Deserialize, Serialize}; + +/// Commitment to a jagged polynomial `q'`, together with all witness data needed +/// for opening proofs. +/// +/// Generic over the inner PCS `Pcs` so that any `PolynomialCommitmentScheme` can +/// serve as the underlying commitment engine. +pub struct JaggedCommitmentWithWitness> { + /// Commitment (with witness) to the "giga" polynomial `q'` via `Pcs`. + pub inner: Pcs::CommitmentWithWitness, + /// Cumulative height sequence `t`: + /// - `t[0] = 0` + /// - `t[i+1] = t[i] + poly_heights[i]` + /// - Length: `num_polys + 1` + pub cumulative_heights: Vec, + /// Number of evaluations `h_i = 2^(num_vars_i)` for each polynomial `p_i`. + /// Length: `num_polys`. + pub poly_heights: Vec, +} + +/// The pure commitment (without witness data) for a jagged polynomial `q'`. +/// This is what the verifier receives. +#[derive(Clone, Serialize, Deserialize)] +#[serde(bound(serialize = "", deserialize = ""))] +pub struct JaggedCommitment> { + /// Pure commitment to the underlying giga polynomial `q'`. + pub inner: Pcs::Commitment, + /// Cumulative height sequence `t` (verifier needs this to evaluate `f(b)`). + pub cumulative_heights: Vec, +} + +impl> JaggedCommitmentWithWitness { + /// Extract the pure commitment (without witness data). + pub fn to_commitment(&self) -> JaggedCommitment { + JaggedCommitment { + inner: Pcs::get_pure_commitment(&self.inner), + cumulative_heights: self.cumulative_heights.clone(), + } + } + + /// Total number of polynomials packed into `q'`. + pub fn num_polys(&self) -> usize { + self.poly_heights.len() + } + + /// Total number of evaluations in the *unpadded* concatenated polynomial + /// (= `t[num_polys]` = `cumulative_heights.last()`). + pub fn total_evaluations(&self) -> usize { + self.cumulative_heights.last().copied().unwrap_or(0) + } +} + +/// Proof for the jagged batch opening protocol. +/// +/// Contains a sumcheck proof (reducing K column evaluation claims to a single +/// point on q'), the evaluation q'(ρ), and an inner PCS opening proof. +#[derive(Clone, Serialize, Deserialize)] +#[serde(bound(serialize = "", deserialize = ""))] +pub struct JaggedBatchOpenProof> { + pub sumcheck_proof: IOPProof, + pub q_eval: E, + pub inner_proof: Pcs::Proof, +} + +/// Convert a `usize` to its little-endian binary representation as field elements. +pub(super) fn int_to_field_bits(val: usize, num_bits: usize) -> Vec { + (0..num_bits) + .map(|i| if (val >> i) & 1 == 1 { E::ONE } else { E::ZERO }) + .collect() +} From a0031e1425681eaf26713d138b61e0673a16012a Mon Sep 17 00:00:00 2001 From: kunxian xia Date: Fri, 24 Apr 2026 14:54:25 +0800 Subject: [PATCH 11/21] fix: add missing rayon trait imports for --all-features clippy Co-Authored-By: Claude Opus 4.6 --- crates/mpcs/src/jagged/mod.rs | 4 +++- crates/mpcs/src/jagged/sumcheck.rs | 3 ++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/crates/mpcs/src/jagged/mod.rs b/crates/mpcs/src/jagged/mod.rs index 60d19a7..8f90ae5 100644 --- a/crates/mpcs/src/jagged/mod.rs +++ b/crates/mpcs/src/jagged/mod.rs @@ -120,7 +120,9 @@ use multilinear_extensions::{ }; use p3::{ matrix::{Matrix, bitrev::BitReversableMatrix}, - maybe_rayon::prelude::{IntoParallelIterator, ParallelSliceMut}, + maybe_rayon::prelude::{ + IndexedParallelIterator, IntoParallelIterator, ParallelIterator, ParallelSliceMut, + }, }; use transcript::Transcript; use types::int_to_field_bits; diff --git a/crates/mpcs/src/jagged/sumcheck.rs b/crates/mpcs/src/jagged/sumcheck.rs index 4c5b57f..0ebe66f 100644 --- a/crates/mpcs/src/jagged/sumcheck.rs +++ b/crates/mpcs/src/jagged/sumcheck.rs @@ -8,7 +8,8 @@ use multilinear_extensions::{ mle::MultilinearExtension, util::max_usable_threads, virtual_poly::build_eq_x_r_vec, }; use p3::maybe_rayon::prelude::{ - IntoParallelIterator, ParallelIterator, ParallelSlice, ParallelSliceMut, + IndexedParallelIterator, IntoParallelIterator, ParallelIterator, ParallelSlice, + ParallelSliceMut, }; use sumcheck::{ macros::{entered_span, exit_span}, From 7df243c0dc405a5a4b8f807807fb17d7c6a54764 Mon Sep 17 00:00:00 2001 From: kunxian xia Date: Fri, 24 Apr 2026 17:27:25 +0800 Subject: [PATCH 12/21] update doc comments: document correction factors for different-height polynomials The zero-padding interpretation (p_i^pad) explains why C_i = eq(z_r[s_i..], 0) is needed when batching polynomials of different heights with a single eq table. Co-Authored-By: Claude Opus 4.6 --- crates/mpcs/src/jagged/mod.rs | 150 ++++++++++++++++++++++------------ 1 file changed, 97 insertions(+), 53 deletions(-) diff --git a/crates/mpcs/src/jagged/mod.rs b/crates/mpcs/src/jagged/mod.rs index 8f90ae5..a1b30b6 100644 --- a/crates/mpcs/src/jagged/mod.rs +++ b/crates/mpcs/src/jagged/mod.rs @@ -22,22 +22,21 @@ //! //! ## Suffix-to-Prefix Transformation //! -//! Each `p_i` has `s` variables (where `s` varies per polynomial; we write `s` instead -//! of `s_i` for brevity). Let `m` be the length of `z_r`. The main sumcheck prover -//! outputs evaluations `v_i = p_i(z_r[(m-s)..m])` — i.e., at the suffix of the +//! Each `p_i` has `s_i` variables and let `m` be the length of `z_r`. The main sumcheck prover +//! outputs evaluations `v_i = p_i(z_r[(m-s_i)..m])` — i.e., at the suffix of the //! random challenge point. //! To make these evaluations compatible with the jagged sumcheck (which operates on //! prefix-aligned polynomials), we apply a bit-reversal permutation to each polynomial's //! evaluations: //! //! ```text -//! p_i'[j] = p_i[bitrev_s(j)] (for j in 0..2^s) +//! p_i'[j] = p_i[bitrev_s_i(j)] (for j in 0..2^s_i) //! ``` //! //! After bit-reversal, //! ```text -//! v_i = p_i(z_r[(m-s)..m]) -//! = p_i'(z_r'[..s]) +//! v_i = p_i(z_r[(m-s_i)..m]) +//! = p_i'(z_r'[..s_i]) //! ``` //! where `z_r' = reverse(z_r)`. //! @@ -45,7 +44,8 @@ //! //! The cumulative height sequence `t` tracks the starting position of each polynomial in `q'`: //! - `t[0] = 0` -//! - `t[i+1] = t[i] + h_i` where `h_i = 2^(num_vars of p_i)` is the number of evaluations +//! - `t[i+1] = t[i] + h_i` where `h_i` is the number of evaluations of `p_i` +//! (`s_i = ceil_log2(h_i)` variables, so `h_i <= 2^{s_i}`) //! //! Given a position `b` in `q'`, the inverse mapping `inv(b) = (i, r)` is defined by: //! - `t[i] <= b < t[i+1]` @@ -67,17 +67,45 @@ //! ## Batch Open Protocol //! //! For notational simplicity, we write `p_i` and `z_r` instead of `p_i'` and `z_r'` -//! (the bit-reversed variants) throughout this section. +//! (the bit-reversed variants) throughout this section. Let `s_i` denote the number +//! of variables of `p_i` and let `m = max(s_i)`. //! -//! Each column opening `v_i = p_i(z_r[..s])` requires its own sumcheck. We batch all K -//! openings into one using `eq(z_c, ·)` weights (soundness loss: `log2(N) / |E|`): +//! Each column opening `v_i = p_i(z_r[..s_i])` requires its own sumcheck. We batch +//! all K openings into one using `eq(z_c, ·)` weights (soundness loss: `log2(N) / |E|`). +//! +//! ### Correction factors for different-height polynomials +//! +//! In the giga polynomial `q'`, each `p_i` occupies only `h_i` slots (padded to +//! `2^{s_i}` for MLE representation). When `s_i < m`, this is equivalent to +//! zero-padding `p_i` to `m` variables: //! //! ```text -//! v = Σ_i eq(z_c, i) · p_i(z_r[..s]) +//! p_i^pad(r, b) = p_i(r) if b = 0 (r ∈ {0,1}^{s_i}, b ∈ {0,1}^{m - s_i}) +//! 0 if b ≠ 0 +//! ``` +//! +//! The MLE of this zero-padded polynomial evaluates to: +//! +//! ```text +//! p_i^pad(z_r) = eq(z_r[s_i..], 0) · p_i(z_r[..s_i]) +//! ``` +//! +//! where `eq(z_r[s_i..], 0) = Π_{j=s_i}^{m-1} (1 - z_r[j])` is the correction +//! factor `C_i` arising from the zero-padded positions. The batched claim becomes: +//! +//! ```text +//! v = Σ_i eq(z_c, i) · C_i · p_i(z_r[..s_i]) //! = Σ_{i,r} eq(z_c, i) · eq(z_r, r) · p_i(r) //! = Σ_{i,r} f(i, r) · p_i(r) where f(i, r) = eq(z_c, i) · eq(z_r, r) //! ``` //! +//! Here `eq(z_r, r)` uses the full `z_r` of length `m`. For `r < 2^{s_i}`, the +//! high bits of `r` are zero, so `eq(z_r, r) = eq(z_r[..s_i], r) · C_i`. +//! This means a single precomputed eq table of size `2^m` naturally incorporates +//! the correction factors — no per-polynomial tables are needed. +//! +//! ### Rewriting via the inverse mapping +//! //! We apply `inv(b)` to rewrite `f` in terms of the giga-index `b`: //! //! ```text @@ -116,6 +144,7 @@ use ff_ext::ExtensionField; use itertools::Itertools; use multilinear_extensions::{ mle::FieldType, + util::ceil_log2, virtual_poly::{VPAuxInfo, build_eq_x_r_vec}, }; use p3::{ @@ -278,11 +307,12 @@ fn compute_f_at_point( f_val } -/// Prove that evaluation claims `evals[i] = p_i(point)` are consistent with a +/// Prove that evaluation claims `evals[i] = p_i(point_i)` are consistent with a /// jagged commitment. /// -/// All polynomials must have the same height. `point` is the common evaluation -/// point (a suffix `r[s..]` of the GKR challenge, length `s = log2(poly_height)`). +/// Polynomials may have different heights. `point` is the evaluation point of +/// length `max_s = max(log2(h_i))`. Polynomial `i` with `s_i` variables is +/// evaluated at the suffix `point[(max_s - s_i)..]`. /// /// The protocol: /// 1. Batch the K column claims via a random column challenge `z_col`. @@ -304,26 +334,29 @@ pub fn jagged_batch_open>( ))); } - let s = point.len(); - let expected_height = 1usize << s; - for (i, &h) in comm.poly_heights.iter().enumerate() { - if h != expected_height { - return Err(Error::InvalidPcsParam(format!( - "jagged_batch_open: poly {} has height {}, expected {}", - i, h, expected_height - ))); - } + let max_s = comm + .poly_heights + .iter() + .map(|&h| ceil_log2(h)) + .max() + .unwrap_or(0); + if point.len() != max_s { + return Err(Error::InvalidPcsParam(format!( + "jagged_batch_open: point length {} != max poly log-height {}", + point.len(), + max_s + ))); } let total_evals = comm.total_evaluations(); - let num_giga_vars = total_evals.next_power_of_two().trailing_zeros() as usize; + let num_giga_vars = ceil_log2(total_evals); // z_row = reverse(point) to account for bit-reversal in jagged_commit. let z_row: Vec = point.iter().rev().cloned().collect(); // Write evals to transcript, then sample z_col. transcript.append_field_element_exts(evals); - let num_col_vars = (num_polys.next_power_of_two().trailing_zeros() as usize).max(1); + let num_col_vars = ceil_log2(num_polys).max(1); let z_col: Vec = transcript.sample_and_append_vec(b"jagged_z_col", num_col_vars); let eq_col = build_eq_x_r_vec(&z_col); @@ -376,8 +409,11 @@ pub fn jagged_batch_open>( }) } -/// Verify that evaluation claims `evals[i] = p_i(point)` are consistent with a +/// Verify that evaluation claims `evals[i] = p_i(point_i)` are consistent with a /// jagged commitment. +/// +/// Polynomials may have different heights. `point` has length `max_s`. Polynomial +/// `i` with `s_i` variables is evaluated at the suffix `point[(max_s - s_i)..]`. pub fn jagged_batch_verify>( vp: &Pcs::VerifierParam, comm: &JaggedCommitment, @@ -396,22 +432,31 @@ pub fn jagged_batch_verify } let total_evals = *comm.cumulative_heights.last().unwrap(); - let num_giga_vars = total_evals.next_power_of_two().trailing_zeros() as usize; + let num_giga_vars = ceil_log2(total_evals); + let max_s = point.len(); // z_row = reverse(point) let z_row: Vec = point.iter().rev().cloned().collect(); // Replay transcript: write evals, sample z_col. transcript.append_field_element_exts(evals); - let num_col_vars = (num_polys.next_power_of_two().trailing_zeros() as usize).max(1); + let num_col_vars = ceil_log2(num_polys).max(1); let z_col: Vec = transcript.sample_and_append_vec(b"jagged_z_col", num_col_vars); - // claimed_sum = Σ_j eq_col[j] · evals[j] + // claimed_sum = Σ_i eq_col[i] · C_i · evals[i] + // where C_i = eq(z_row[s_i..], 0) = Π_{j=s_i}^{max_s-1} (1 - z_row[j]) is the + // correction factor from zero-padding p_i to max_s variables (see module docs). let eq_col = build_eq_x_r_vec(&z_col); - let claimed_sum: E = eq_col[..num_polys] - .iter() - .zip(evals.iter()) - .map(|(&eq, &ev)| eq * ev) + let mut suffix_prod = vec![E::ONE; max_s + 1]; + for j in (0..max_s).rev() { + suffix_prod[j] = suffix_prod[j + 1] * (E::ONE - z_row[j]); + } + let claimed_sum: E = (0..num_polys) + .map(|i| { + let h_i = comm.cumulative_heights[i + 1] - comm.cumulative_heights[i]; + let s_i = ceil_log2(h_i); + eq_col[i] * suffix_prod[s_i] * evals[i] + }) .sum(); // Verify the jagged sumcheck. @@ -582,36 +627,35 @@ mod tests { fn test_jagged_batch_open_verify_small() { let mut rng = thread_rng(); - let num_rows = 1024usize; // s=10 - let num_cols = 3usize; - let s = 10; - let total_evals = num_rows * num_cols; - let num_giga_vars = total_evals.next_power_of_two().trailing_zeros() as usize; + // 3 matrices with different heights: 2^10, 2^11, 2^9 (each 1 column). + let log_heights = [10usize, 11, 9]; + let heights: Vec = log_heights.iter().map(|&s| 1 << s).collect(); + let max_s = *log_heights.iter().max().unwrap(); + let total_evals: usize = heights.iter().sum(); + let num_giga_vars = ceil_log2(total_evals); let (pp, vp) = setup_pcs::(num_giga_vars); - let rmm = make_rmm(num_rows, num_cols); + let rmms: Vec<_> = heights.iter().map(|&h| make_rmm(h, 1)).collect(); - // Extract column polynomials (before bit-reversal) for computing true evaluations. - let col_polys: Vec> = (0..num_cols) - .map(|c| { - (0..num_rows) - .map(|r| rmm.values[r * num_cols + c]) - .collect() - }) + // Extract each column polynomial (before bit-reversal). + let col_polys: Vec> = rmms + .iter() + .map(|rmm| (0..rmm.height()).map(|r| rmm.values[r]).collect()) .collect(); // Commit. let mut transcript_p = BasicTranscript::::new(b"jagged_batch_test"); - let comm = jagged_commit::(&pp, vec![rmm]).expect("commit should succeed"); + let comm = jagged_commit::(&pp, rmms).expect("commit should succeed"); Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); - // Random evaluation point of length s. - let point: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); + // Random evaluation point of length max_s. + // Poly i with s_i variables is evaluated at the suffix point[(max_s - s_i)..]. + let point: Vec = (0..max_s).map(|_| E::random(&mut rng)).collect(); - // Compute true evaluations. let evals: Vec = col_polys .iter() - .map(|col| eval_column_poly_at_point(col, &point)) + .zip(log_heights.iter()) + .map(|(col, &s_i)| eval_column_poly_at_point(col, &point[(max_s - s_i)..])) .collect(); // Prover: batch open. @@ -640,7 +684,7 @@ mod tests { let num_rows = 1024usize; // s=10 let num_cols = 1usize; let s = 10; - let num_giga_vars = (num_rows * num_cols).next_power_of_two().trailing_zeros() as usize; + let num_giga_vars = ceil_log2(num_rows * num_cols); let (pp, vp) = setup_pcs::(num_giga_vars); let rmm = make_rmm(num_rows, num_cols); @@ -678,7 +722,7 @@ mod tests { let num_rows = 1024usize; let num_cols = 3usize; let s = 10; - let num_giga_vars = (num_rows * num_cols).next_power_of_two().trailing_zeros() as usize; + let num_giga_vars = ceil_log2(num_rows * num_cols); let (pp, vp) = setup_pcs::(num_giga_vars); let rmm = make_rmm(num_rows, num_cols); From ed522afa10033c8eac92d5e51b5774b4812ebdc1 Mon Sep 17 00:00:00 2001 From: kunxian xia Date: Fri, 24 Apr 2026 17:33:27 +0800 Subject: [PATCH 13/21] refactor: rename generic parameter Pcs to InnerPcs for readability Co-Authored-By: Claude Opus 4.6 --- crates/mpcs/src/jagged/mod.rs | 36 ++++++++++++++++----------------- crates/mpcs/src/jagged/types.rs | 24 ++++++++++++---------- 2 files changed, 31 insertions(+), 29 deletions(-) diff --git a/crates/mpcs/src/jagged/mod.rs b/crates/mpcs/src/jagged/mod.rs index a1b30b6..15695eb 100644 --- a/crates/mpcs/src/jagged/mod.rs +++ b/crates/mpcs/src/jagged/mod.rs @@ -165,19 +165,19 @@ use witness::{InstancePaddingStrategy, RowMajorMatrix}; /// column polynomial occupies a contiguous region in memory. /// 3. Concatenate all column polynomials: `q' = col_0 || col_1 || ...` /// 4. Compute the cumulative height sequence `t`. -/// 5. Commit to `q'` as a single-column matrix using `Pcs::batch_commit`. +/// 5. Commit to `q'` as a single-column matrix using `InnerPcs::batch_commit`. /// /// # Arguments -/// * `pp` — Prover parameters for `Pcs`. +/// * `pp` — Prover parameters for `InnerPcs`. /// * `rmms` — Non-empty sequence of row-major matrices. This function uses each matrix's height exactly as given. /// /// # Errors /// Returns `Error::InvalidPcsParam` if `rmms` is empty or all matrices are empty. -/// Any error from the inner `Pcs::batch_commit` is propagated as-is. -pub fn jagged_commit>( - pp: &Pcs::ProverParam, +/// Any error from the inner `InnerPcs::batch_commit` is propagated as-is. +pub fn jagged_commit>( + pp: &InnerPcs::ProverParam, rmms: Vec>, -) -> Result, Error> { +) -> Result, Error> { if rmms.is_empty() { return Err(Error::InvalidPcsParam( "jagged_commit: cannot commit to empty sequence of matrices".to_string(), @@ -270,7 +270,7 @@ pub fn jagged_commit>( InstancePaddingStrategy::Default, ); - let inner = Pcs::batch_commit(pp, vec![giga_rmm])?; + let inner = InnerPcs::batch_commit(pp, vec![giga_rmm])?; Ok(JaggedCommitmentWithWitness { inner, @@ -318,13 +318,13 @@ fn compute_f_at_point( /// 1. Batch the K column claims via a random column challenge `z_col`. /// 2. Run the jagged sumcheck to reduce to a single evaluation of q'. /// 3. Open q' at the sumcheck output point via the inner PCS. -pub fn jagged_batch_open>( - pp: &Pcs::ProverParam, - comm: &JaggedCommitmentWithWitness, +pub fn jagged_batch_open>( + pp: &InnerPcs::ProverParam, + comm: &JaggedCommitmentWithWitness, point: &[E], evals: &[E], transcript: &mut impl Transcript, -) -> Result, Error> { +) -> Result, Error> { let num_polys = comm.num_polys(); if evals.len() != num_polys { return Err(Error::InvalidPcsParam(format!( @@ -362,7 +362,7 @@ pub fn jagged_batch_open>( let eq_col = build_eq_x_r_vec(&z_col); // Extract q' base-field evaluations from the inner PCS witness. - let q_mles = Pcs::get_arc_mle_witness_from_commitment(&comm.inner); + let q_mles = InnerPcs::get_arc_mle_witness_from_commitment(&comm.inner); assert_eq!( q_mles.len(), 1, @@ -396,7 +396,7 @@ pub fn jagged_batch_open>( transcript.append_field_element_ext(&q_eval); // Open q' at ρ via inner PCS batch_open. - let inner_proof = Pcs::batch_open( + let inner_proof = InnerPcs::batch_open( pp, vec![(&comm.inner, vec![(challenges, vec![q_eval])])], transcript, @@ -414,12 +414,12 @@ pub fn jagged_batch_open>( /// /// Polynomials may have different heights. `point` has length `max_s`. Polynomial /// `i` with `s_i` variables is evaluated at the suffix `point[(max_s - s_i)..]`. -pub fn jagged_batch_verify>( - vp: &Pcs::VerifierParam, - comm: &JaggedCommitment, +pub fn jagged_batch_verify>( + vp: &InnerPcs::VerifierParam, + comm: &JaggedCommitment, point: &[E], evals: &[E], - proof: &JaggedBatchOpenProof, + proof: &JaggedBatchOpenProof, transcript: &mut impl Transcript, ) -> Result<(), Error> { let num_polys = comm.cumulative_heights.len() - 1; @@ -497,7 +497,7 @@ pub fn jagged_batch_verify } // Verify the inner PCS opening. - Pcs::batch_verify( + InnerPcs::batch_verify( vp, vec![( comm.inner.clone(), diff --git a/crates/mpcs/src/jagged/types.rs b/crates/mpcs/src/jagged/types.rs index 0b5606c..bc2381e 100644 --- a/crates/mpcs/src/jagged/types.rs +++ b/crates/mpcs/src/jagged/types.rs @@ -6,11 +6,11 @@ use serde::{Deserialize, Serialize}; /// Commitment to a jagged polynomial `q'`, together with all witness data needed /// for opening proofs. /// -/// Generic over the inner PCS `Pcs` so that any `PolynomialCommitmentScheme` can +/// Generic over the inner PCS `InnerPcs` so that any `PolynomialCommitmentScheme` can /// serve as the underlying commitment engine. -pub struct JaggedCommitmentWithWitness> { - /// Commitment (with witness) to the "giga" polynomial `q'` via `Pcs`. - pub inner: Pcs::CommitmentWithWitness, +pub struct JaggedCommitmentWithWitness> { + /// Commitment (with witness) to the "giga" polynomial `q'` via `InnerPcs`. + pub inner: InnerPcs::CommitmentWithWitness, /// Cumulative height sequence `t`: /// - `t[0] = 0` /// - `t[i+1] = t[i] + poly_heights[i]` @@ -25,18 +25,20 @@ pub struct JaggedCommitmentWithWitness> { +pub struct JaggedCommitment> { /// Pure commitment to the underlying giga polynomial `q'`. - pub inner: Pcs::Commitment, + pub inner: InnerPcs::Commitment, /// Cumulative height sequence `t` (verifier needs this to evaluate `f(b)`). pub cumulative_heights: Vec, } -impl> JaggedCommitmentWithWitness { +impl> + JaggedCommitmentWithWitness +{ /// Extract the pure commitment (without witness data). - pub fn to_commitment(&self) -> JaggedCommitment { + pub fn to_commitment(&self) -> JaggedCommitment { JaggedCommitment { - inner: Pcs::get_pure_commitment(&self.inner), + inner: InnerPcs::get_pure_commitment(&self.inner), cumulative_heights: self.cumulative_heights.clone(), } } @@ -59,10 +61,10 @@ impl> JaggedCommitmentWith /// point on q'), the evaluation q'(ρ), and an inner PCS opening proof. #[derive(Clone, Serialize, Deserialize)] #[serde(bound(serialize = "", deserialize = ""))] -pub struct JaggedBatchOpenProof> { +pub struct JaggedBatchOpenProof> { pub sumcheck_proof: IOPProof, pub q_eval: E, - pub inner_proof: Pcs::Proof, + pub inner_proof: InnerPcs::Proof, } /// Convert a `usize` to its little-endian binary representation as field elements. From 352ebd2d59a7b15719423b5197239fa9694abefc Mon Sep 17 00:00:00 2001 From: kunxian xia Date: Fri, 24 Apr 2026 21:18:15 +0800 Subject: [PATCH 14/21] docs: document 2^{s_i} block sizes and zero-padding cost vs SP1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Clarify that each polynomial in q' occupies a full 2^{s_i} block (not the original h_i entries) because bit-reversal scatters the zero-padding throughout the block. Add note comparing with SP1's jagged PCS which uses raw h_i, highlighting the Σ(2^{s_i} - h_i) extra zeros as the cost of suffix-to-prefix bit-reversal. Co-Authored-By: Claude Opus 4.6 --- crates/mpcs/src/jagged/mod.rs | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/crates/mpcs/src/jagged/mod.rs b/crates/mpcs/src/jagged/mod.rs index 15695eb..c62118f 100644 --- a/crates/mpcs/src/jagged/mod.rs +++ b/crates/mpcs/src/jagged/mod.rs @@ -42,10 +42,19 @@ //! //! ## Cumulative Heights //! +//! Each polynomial `p_i` has `s_i = ceil_log2(h_i)` variables, where `h_i` is the +//! original (unpadded) number of evaluations. Before bit-reversal, `p_i` is +//! zero-padded to `2^{s_i}` entries. **After bit-reversal, these zeros are scattered +//! throughout the block** (not contiguous at the end), so each polynomial occupies +//! a full `2^{s_i}`-entry block in `q'`. +//! +//! **Note:** Unlike SP1's jagged PCS (which uses raw `h_i`), our `q'` contains +//! `Σ (2^{s_i} - h_i)` extra zeros from this padding — the cost of the +//! suffix-to-prefix bit-reversal required by Ceno's main sumcheck. +//! //! The cumulative height sequence `t` tracks the starting position of each polynomial in `q'`: //! - `t[0] = 0` -//! - `t[i+1] = t[i] + h_i` where `h_i` is the number of evaluations of `p_i` -//! (`s_i = ceil_log2(h_i)` variables, so `h_i <= 2^{s_i}`) +//! - `t[i+1] = t[i] + 2^{s_i}` (the padded height, not the original `h_i`) //! //! Given a position `b` in `q'`, the inverse mapping `inv(b) = (i, r)` is defined by: //! - `t[i] <= b < t[i+1]` @@ -57,8 +66,10 @@ //! ## Commit Protocol //! //! 1. For each input matrix `M_k` (with `h_k` rows and `w_k` columns): -//! a. Extract each column as a polynomial with `h_k` evaluations. -//! b. Apply bit-reversal to the evaluations. +//! a. Zero-pad `M_k` to `2^{s_k}` rows (where `s_k = ceil_log2(h_k)`). +//! b. Apply bit-reversal to the rows. Note: after bit-reversal the zero-padding +//! is scattered throughout the block, not contiguous at the end. +//! c. Extract each column as a polynomial with `2^{s_k}` evaluations. //! 2. Concatenate all bit-reversed polynomials: `cat = bitrev(p_0) || bitrev(p_1) || ...` //! 3. Compute cumulative heights `t[i]`. //! 4. Pad `cat` to the next power of two (required for MLE representation). From a14b87bff6e0d3fd61cf40a332cb9c6afb472e3f Mon Sep 17 00:00:00 2001 From: xkx Date: Tue, 28 Apr 2026 19:13:18 +0800 Subject: [PATCH 15/21] feat: assist sumcheck for jagged PCS (#42) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * feat: implement assist sumcheck to reduce K ROBP evaluations to one Implements the assist sumcheck protocol (Lemma 5.1 of eprint 2025/917) which batches K indicator function evaluations into a single opening, using an interleaved variable ordering and forward-backward ROBP state decomposition. Co-Authored-By: Claude Opus 4.6 * perf: parallelize backward precomputation and fuse round accumulations - Parallelize the O(K × n_robp) backward vector precomputation with rayon - Derive bwd_sum_d from bwd_sum in O(1) instead of a second O(K) pass - Fuse the two per-step weight updates into one - Add assist_sumcheck benchmark (K=100/500/1000) K=1000 drops from 14.1ms to 4.6ms (3.1x speedup with --features parallel). Co-Authored-By: Claude Opus 4.6 * docs: annotate key MLE telescoping identity in assist sumcheck Co-Authored-By: Claude Opus 4.6 * bench: update assist sumcheck benchmark to K=1000,2000,4000,8000 Co-Authored-By: Claude Opus 4.6 * refactor: derive ROBP transition matrices from raw transition table Separate the ROBP state machine logic from the eq-weighting algebra. The raw transition table ROBP_TRANSITION encodes the state machine directly, and symbol_transition_matrices derives M_i^{(c,d)} via the closed formula Σ_{a,b} eq₁(z1,a)·eq₁(z2,b)·[transition(from,(a,b,c,d))=to]. Also document why z₁/z₂ are bound first (reducing alphabet from {0,1}⁴ to {0,1}²) and z₃/z₄ interleaved to match ROBP step order. Co-Authored-By: Claude Opus 4.6 * fix: suppress needless_range_loop clippy lint and apply fmt Co-Authored-By: Claude Opus 4.6 * bench: add end-to-end jagged PCS benchmark (commit, open, verify) Co-Authored-By: Claude Opus 4.6 * fix: suppress clippy needless_range_loop in evaluator tests, add K=16000 bench Co-Authored-By: Claude Opus 4.6 * perf: transpose bwd/c_bits/d_bits to step-major layout for cache-friendly access Reorganize data structures from poly-major bwd[y][i] to step-major bwd[i][y], and similarly for c_bits and d_bits. The main loop now scans contiguous memory when accumulating over K polynomials, reducing cache misses at large K (~38% improvement at K=16000). Backward precomputation uses split_at_mut + par_iter for safe parallel writes. Also adds the round polynomial formula (§2.3, Eq. 4) as a code comment to clarify the bwd_sum accumulation logic. Co-Authored-By: Claude Opus 4.6 * perf: parallelize assist sumcheck round-polynomial computation and weight update Thread-local bwd_sum buckets + per-thread p(0),p(1),p(2) evaluation eliminates the sequential O(K) bottleneck in the main step loop. Each thread builds local buckets, computes local round polynomial values, and only the scalars are summed across threads. Benchmark (BabyBearExt4, --features parallel, vs HEAD): K=1000: 7.8ms → 6.8ms (-14%) K=2000: 13.5ms → 10.2ms (-24%) K=4000: 25.6ms → 16.2ms (-38%) K=8000: 50.7ms → 28.0ms (-46%) K=16000: 100.5ms → 53.6ms (-48%) Co-Authored-By: Claude Opus 4.6 * perf: eliminate bwd_sum merge by reusing per-thread buckets in round 2i+1 Instead of merging thread-local bwd_sums after round 2i, cache them and absorb alpha directly into each thread's local bwd_sum for round 2i+1. Only scalar p(0),p(1),p(2) values are summed across threads in both rounds. Co-Authored-By: Claude Opus 4.6 * fix: gate IndexedParallelIterator import behind parallel feature, remove tautological debug_assert Co-Authored-By: Claude Opus 4.6 * fmt * fix: use wildcard maybe_rayon import, remove unused direct rayon dep Wildcard `use p3::maybe_rayon::prelude::*` avoids unused-import warnings caused by Cargo feature unification: `p3/parallel` can be activated by other workspace crates even when `mpcs/parallel` is off, making a `#[cfg(feature = "parallel")]` gate on IndexedParallelIterator incorrect. Also remove `dep:rayon` from mpcs — all parallelism goes through `p3::maybe_rayon`, so the direct rayon dependency was unnecessary. Co-Authored-By: Claude Opus 4.6 --------- Co-authored-by: Claude Opus 4.6 --- Cargo.lock | 1 - crates/mpcs/Cargo.toml | 7 +- crates/mpcs/benches/jagged_pcs.rs | 169 ++++++++ crates/mpcs/benches/jagged_sumcheck.rs | 46 ++- crates/mpcs/src/jagged/assist.rs | 534 +++++++++++++++++++++++++ crates/mpcs/src/jagged/evaluator.rs | 209 +++++++++- crates/mpcs/src/jagged/mod.rs | 89 ++++- crates/mpcs/src/jagged/sumcheck.rs | 5 +- crates/mpcs/src/jagged/types.rs | 5 +- crates/mpcs/src/lib.rs | 4 +- 10 files changed, 1035 insertions(+), 34 deletions(-) create mode 100644 crates/mpcs/benches/jagged_pcs.rs create mode 100644 crates/mpcs/src/jagged/assist.rs diff --git a/Cargo.lock b/Cargo.lock index 554bd40..a4a5bea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -834,7 +834,6 @@ dependencies = [ "p3", "rand", "rand_chacha", - "rayon", "serde", "sumcheck", "tracing", diff --git a/crates/mpcs/Cargo.toml b/crates/mpcs/Cargo.toml index 0962118..fe997f9 100644 --- a/crates/mpcs/Cargo.toml +++ b/crates/mpcs/Cargo.toml @@ -20,7 +20,6 @@ num-integer = "0.1" p3.workspace = true rand.workspace = true rand_chacha.workspace = true -rayon = { workspace = true, optional = true } serde.workspace = true sumcheck.workspace = true tracing.workspace = true @@ -34,7 +33,7 @@ tracing-forest = "0.1" [features] nightly-features = ["ff_ext/nightly-features"] -parallel = ["p3/parallel", "dep:rayon"] +parallel = ["p3/parallel"] whir = ["dep:whir"] print-trace = ["whir/print-trace"] sanity-check = [] @@ -59,3 +58,7 @@ required-features = ["whir"] [[bench]] harness = false name = "jagged_sumcheck" + +[[bench]] +harness = false +name = "jagged_pcs" diff --git a/crates/mpcs/benches/jagged_pcs.rs b/crates/mpcs/benches/jagged_pcs.rs new file mode 100644 index 0000000..ba44ea4 --- /dev/null +++ b/crates/mpcs/benches/jagged_pcs.rs @@ -0,0 +1,169 @@ +use std::time::Duration; + +use criterion::*; +use ff_ext::{FromUniformBytes, GoldilocksExt2}; +use mpcs::{ + Basefold, BasefoldRSParams, PolynomialCommitmentScheme, SecurityLevel, jagged_batch_open, + jagged_commit, +}; +use multilinear_extensions::{mle::MultilinearExtension, util::ceil_log2}; +use p3::{field::FieldAlgebra, goldilocks::Goldilocks, matrix::Matrix}; +use rand::{Rng, thread_rng}; +use transcript::BasicTranscript; +use witness::{InstancePaddingStrategy, RowMajorMatrix}; + +type E = GoldilocksExt2; +type F = Goldilocks; +type Pcs = Basefold; + +const NUM_SAMPLES: usize = 10; +const NUM_MATRICES: usize = 30; +const NUM_COLS: usize = 32; + +fn make_rmm(num_rows: usize, num_cols: usize) -> RowMajorMatrix { + let values: Vec = (0..num_rows * num_cols) + .map(|i| F::from_canonical_u32(((i as u64 * 13 + 7) % (1 << 30)) as u32)) + .collect(); + RowMajorMatrix::::new_by_values(values, num_cols, InstancePaddingStrategy::Default) +} + +fn sample_heights(rng: &mut impl Rng, num_matrices: usize) -> Vec { + let log_range: Vec = (16..=22).collect(); + (0..num_matrices) + .map(|_| 1usize << log_range[rng.gen_range(0..log_range.len())]) + .collect() +} + +fn eval_column_poly_at_point(col_evals: &[F], point: &[E]) -> E { + let s = point.len(); + assert_eq!(col_evals.len(), 1 << s); + let mle = MultilinearExtension::from_evaluations_vec(s, col_evals.to_vec()); + mle.evaluate(point) +} + +fn bench_jagged_pcs(c: &mut Criterion) { + let mut group = c.benchmark_group("jagged_pcs"); + group.sample_size(NUM_SAMPLES); + + let mut rng = thread_rng(); + + let heights = sample_heights(&mut rng, NUM_MATRICES); + println!( + "Matrix heights (log2): {:?}", + heights.iter().map(|h| ceil_log2(*h)).collect::>() + ); + + let rmms: Vec<_> = heights.iter().map(|&h| make_rmm(h, NUM_COLS)).collect(); + + let log_heights: Vec = rmms.iter().map(|rmm| ceil_log2(rmm.height())).collect(); + let max_s = *log_heights.iter().max().unwrap(); + let total_evals: usize = rmms.iter().map(|rmm| rmm.height() * rmm.width()).sum(); + let num_giga_vars = ceil_log2(total_evals); + let num_polys = NUM_MATRICES * NUM_COLS; + + println!( + "num_matrices={}, num_cols={}, num_polys={}, total_evals={}, num_giga_vars={}, max_s={}", + NUM_MATRICES, NUM_COLS, num_polys, total_evals, num_giga_vars, max_s + ); + + let poly_size = 1usize << num_giga_vars; + let param = Pcs::setup(poly_size, SecurityLevel::Conjecture100bits).unwrap(); + let (pp, vp) = Pcs::trim(param, poly_size).unwrap(); + + // --- Bench commit --- + let comm = jagged_commit::(&pp, rmms.clone()).expect("commit failed"); + + group.bench_function( + BenchmarkId::new("commit", format!("{}x{}", NUM_MATRICES, NUM_COLS)), + |b| { + b.iter_custom(|iters| { + let mut time = Duration::new(0, 0); + for _ in 0..iters { + let rmms_clone = rmms.clone(); + let instant = std::time::Instant::now(); + let _ = jagged_commit::(&pp, rmms_clone).unwrap(); + time += instant.elapsed(); + } + time + }) + }, + ); + + // --- Prepare evaluation data --- + // Extract column polynomials (before bit-reversal) for computing true evaluations. + let col_polys: Vec> = rmms + .iter() + .flat_map(|rmm| { + let h = rmm.height(); + let w = rmm.width(); + (0..w).map(move |c| (0..h).map(|r| rmm.values[r * w + c]).collect()) + }) + .collect(); + + let point: Vec = (0..max_s).map(|_| E::random(&mut rng)).collect(); + + let evals: Vec = col_polys + .iter() + .zip( + log_heights + .iter() + .flat_map(|&s| std::iter::repeat_n(s, NUM_COLS)), + ) + .map(|(col, s_i)| eval_column_poly_at_point(col, &point[(max_s - s_i)..])) + .collect(); + + // --- Bench batch open --- + group.bench_function( + BenchmarkId::new("batch_open", format!("{}x{}", NUM_MATRICES, NUM_COLS)), + |b| { + b.iter_batched( + || { + let mut transcript = BasicTranscript::::new(b"jagged_bench"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript).unwrap(); + transcript + }, + |mut transcript| { + jagged_batch_open::(&pp, &comm, &point, &evals, &mut transcript) + .unwrap(); + }, + BatchSize::SmallInput, + ); + }, + ); + + // --- Bench batch verify --- + let mut transcript_p = BasicTranscript::::new(b"jagged_bench"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); + let proof = jagged_batch_open::(&pp, &comm, &point, &evals, &mut transcript_p).unwrap(); + let pure_comm = comm.to_commitment(); + + group.bench_function( + BenchmarkId::new("batch_verify", format!("{}x{}", NUM_MATRICES, NUM_COLS)), + |b| { + b.iter_batched( + || { + let mut transcript = BasicTranscript::::new(b"jagged_bench"); + Pcs::write_commitment(&pure_comm.inner, &mut transcript).unwrap(); + transcript + }, + |mut transcript| { + mpcs::jagged_batch_verify::( + &vp, + &pure_comm, + &point, + &evals, + &proof, + &mut transcript, + ) + .unwrap(); + }, + BatchSize::SmallInput, + ); + }, + ); + + group.finish(); +} + +criterion_group!(benches, bench_jagged_pcs); +criterion_main!(benches); diff --git a/crates/mpcs/benches/jagged_sumcheck.rs b/crates/mpcs/benches/jagged_sumcheck.rs index 4dbaed0..a6efe4e 100644 --- a/crates/mpcs/benches/jagged_sumcheck.rs +++ b/crates/mpcs/benches/jagged_sumcheck.rs @@ -1,7 +1,7 @@ use criterion::*; use ff_ext::{BabyBearExt4, ExtensionField, FieldFrom, FromUniformBytes}; -use mpcs::{JaggedSumcheckInput, jagged_sumcheck_prove}; -use multilinear_extensions::virtual_poly::build_eq_x_r_vec; +use mpcs::{JaggedSumcheckInput, assist_sumcheck_prove, jagged_sumcheck_prove}; +use multilinear_extensions::{util::ceil_log2, virtual_poly::build_eq_x_r_vec}; use rand::thread_rng; use transcript::BasicTranscript; @@ -61,5 +61,45 @@ fn bench_jagged_sumcheck(c: &mut Criterion) { group.finish(); } -criterion_group!(benches, bench_jagged_sumcheck); +fn bench_assist_sumcheck(c: &mut Criterion) { + let mut group = c.benchmark_group("assist_sumcheck"); + group.sample_size(NUM_SAMPLES); + + let poly_height_log2 = 20usize; + let poly_height = 1usize << poly_height_log2; + + for num_polys in [1000, 2000, 4000, 8000, 16000] { + let mut rng = thread_rng(); + + let total_evals = num_polys * poly_height; + let num_giga_vars = ceil_log2(total_evals); + let n_robp = num_giga_vars + if total_evals.is_power_of_two() { 1 } else { 0 }; + + let cumulative_heights: Vec = (0..=num_polys).map(|i| i * poly_height).collect(); + + let z_row: Vec = (0..n_robp).map(|_| E::random(&mut rng)).collect(); + let rho: Vec = (0..n_robp).map(|_| E::random(&mut rng)).collect(); + let num_col_vars = ceil_log2(num_polys).max(1); + let z_col: Vec = (0..num_col_vars).map(|_| E::random(&mut rng)).collect(); + let eq_col = build_eq_x_r_vec(&z_col); + + group.bench_function(BenchmarkId::new("prove", format!("K={}", num_polys)), |b| { + b.iter(|| { + let mut transcript = BasicTranscript::::new(b"assist_bench"); + assist_sumcheck_prove( + black_box(&z_row), + black_box(&rho), + black_box(&eq_col), + black_box(&cumulative_heights), + black_box(n_robp), + &mut transcript, + ) + }) + }); + } + + group.finish(); +} + +criterion_group!(benches, bench_jagged_sumcheck, bench_assist_sumcheck); criterion_main!(benches); diff --git a/crates/mpcs/src/jagged/assist.rs b/crates/mpcs/src/jagged/assist.rs new file mode 100644 index 0000000..51c5e3c --- /dev/null +++ b/crates/mpcs/src/jagged/assist.rs @@ -0,0 +1,534 @@ +//! Jagged Assist Sumcheck +//! +//! Reduces K evaluations of the indicator function ĝ (one per column) to a +//! single evaluation, via the batch-proving protocol of §5 (Lemma 5.1) of the +//! jagged PCS paper. +//! +//! The sumcheck operates on P(b) = h(b) · Q(b) where: +//! h(z₃, z₄) = ĝ(z_row, ρ, z₃, z₄) (multilinear in 2n variables) +//! Q(b) = Σ_y eq_col[y] · eq(b, x_y) (x_y are Boolean evaluation points) +//! +//! Variables are interleaved as (z₃[0], z₄[0], z₃[1], z₄[1], …) so that +//! each pair of consecutive sumcheck rounds maps to one ROBP step. + +use ff_ext::ExtensionField; +use multilinear_extensions::util::max_usable_threads; +use p3::maybe_rayon::prelude::*; +use sumcheck::structs::{IOPProof, IOPProverMessage}; +use transcript::Transcript; + +use super::evaluator::{ + ROBP_WIDTH, StateVec, TransitionMatrix, dot4, mat_vec_mul, sink_labels, source_vec, + symbol_transition_matrices, vec_mat_mul, +}; + +/// Run the assist sumcheck prover. +/// +/// Proves that `claimed_sum = Σ_y eq_col[y] · ĝ(z_row, ρ, bits(t_y), bits(t_{y+1}))`. +/// +/// Returns the proof (2·n_robp rounds) and the full challenge vector. +pub fn assist_sumcheck_prove( + z_row_padded: &[E], + rho_padded: &[E], + eq_col: &[E], + cumulative_heights: &[usize], + n_robp: usize, + transcript: &mut impl Transcript, +) -> (IOPProof, Vec) { + let num_polys = cumulative_heights.len() - 1; + let n_vars = 2 * n_robp; + let max_degree: usize = 2; + + // Write transcript header (must match verifier). + transcript.append_message(&n_vars.to_le_bytes()); + transcript.append_message(&max_degree.to_le_bytes()); + + // Precompute per-step symbol matrices. + let step_mats: Vec<[TransitionMatrix; 4]> = (0..n_robp) + .map(|i| symbol_transition_matrices(z_row_padded[i], rho_padded[i])) + .collect(); + + // Extract Boolean bits in step-major layout: c_bits[i][y], d_bits[i][y]. + // c_bits[i][y] = bit_i(t_y), d_bits[i][y] = bit_i(t_{y+1}) + let mut c_bits = vec![vec![0usize; num_polys]; n_robp]; + let mut d_bits = vec![vec![0usize; num_polys]; n_robp]; + for i in 0..n_robp { + for y in 0..num_polys { + c_bits[i][y] = (cumulative_heights[y] >> i) & 1; + d_bits[i][y] = (cumulative_heights[y + 1] >> i) & 1; + } + } + + // Precompute backward vectors bwd[i][y] (step-major layout). + // + // bwd[i][y] is a 4-element state vector representing the ROBP suffix + // product from step i to the sinks, for polynomial y's Boolean bits: + // + // bwd[n][y] = u = [0, 1, 0, 0] (sink labels: accept at carry=0, lt=1) + // bwd[i][y] = M_i^{(c_y[i], d_y[i])} · bwd[i+1][y], i = n-1, …, 0 + // + // where c_y[i] = bit_i(t_y), d_y[i] = bit_i(t_{y+1}), and M_i^{(c,d)} + // is the per-symbol transition matrix with z_row[i], ρ[i] baked in. + // + // The full ROBP evaluation for polynomial y equals: + // ĝ(z_row, ρ, bits(t_y), bits(t_{y+1})) = e_0^T · bwd[0][y] + // + // During the sumcheck, the prover decomposes ĝ into fwd · bwd. + // M_i^{(c,d)} extends to field arguments via MLE: M_i^{(α,β)} means + // Σ_{c,d} eq₁(α,c)·eq₁(β,d)·M_i^{(c,d)}. With this notation: + // round 2i: h_y(λ) = fwd · M_i^{(λ, d_y[i])} · bwd[i+1][y] + // round 2i+1: h_y(λ) = fwd · M_i^{(α, λ)} · bwd[i+1][y] + // where fwd absorbs bound challenges and bwd provides the Boolean suffix. + let u = sink_labels::(); + let mut bwd = vec![vec![[E::ZERO; ROBP_WIDTH]; num_polys]; n_robp + 1]; + for y in 0..num_polys { + bwd[n_robp][y] = u; + } + for i in (0..n_robp).rev() { + let (left, right) = bwd.split_at_mut(i + 1); + let dst = &mut left[i]; + let src = &right[0]; + dst.into_par_iter().enumerate().for_each(|(y, dst_y)| { + let cd = c_bits[i][y] * 2 + d_bits[i][y]; + *dst_y = mat_vec_mul(&step_mats[i][cd], &src[y]); + }); + } + + // Initialize weights and forward vector. + let mut weights: Vec = eq_col[..num_polys].to_vec(); + let mut fwd: StateVec = source_vec(); + + let mut challenges: Vec = Vec::with_capacity(n_vars); + let mut proof_messages: Vec> = Vec::with_capacity(n_vars); + + let n_threads = max_usable_threads(); + let poly_indices: Vec = (0..num_polys).collect(); + + for i in 0..n_robp { + // Precompute fwd · M^{(c,d)} for all 4 symbols. + let r_cd: [StateVec; 4] = std::array::from_fn(|cd| vec_mat_mul(&fwd, &step_mats[i][cd])); + + // ---- Round 2i: bind z₃[i] ---- + // + // Round polynomial for round 2i: + // p_{2i}(λ) = Σ_y w_y · eq₁(λ, c_y[i]) · fwd · M_i^{(λ, d_y[i])} · bwd[i+1][y] + // + // Precompute R[c*2+d] = fwd · M_i^{(c,d)} as row-vectors for Boolean (c,d). + // Then fwd · M_i^{(λ,d)} = (1-λ)·R[d] + λ·R[2+d] is also a row-vector, + // so each term reduces to a dot product with bwd — no matrix-vector multiply. + // + // Expanding M_i^{(λ,d)} = Σ_{c'} eq₁(λ,c')·M_i^{(c',d)} in the bucketed form: + // p_{2i}(λ) = Σ_{c,c',d} eq₁(λ,c) · eq₁(λ,c') · R[c'*2+d] · bwd_sum[c*2+d] + // where c indexes the Q bucket and c' indexes the M expansion (independent). + // + // Bucket bwd vectors by (c,d): + // bwd_sum[c*2+d] = Σ_{y: c_y[i]=c, d_y[i]=d} w_y · bwd[i+1][y] + // + // This reduces evaluation at each λ to 4 dot products against bwd_sum. + // + // Parallelization: partition K polynomials across threads. Each thread + // builds local bwd_sum, computes local (p0, p1, p2), then we sum + // the scalars across threads. We also merge bwd_sum for round 2i+1. + let row2_d0: StateVec = std::array::from_fn(|j| r_cd[2][j].double() - r_cd[0][j]); + let row2_d1: StateVec = std::array::from_fn(|j| r_cd[3][j].double() - r_cd[1][j]); + + let batch_size = (num_polys / n_threads).max(1); + let partials: Vec<([E; 3], [[E; ROBP_WIDTH]; 4])> = poly_indices + .par_chunks(batch_size) + .map(|chunk| { + let mut local_bwd_sum = [[E::ZERO; ROBP_WIDTH]; 4]; + for &y in chunk { + let cd = c_bits[i][y] * 2 + d_bits[i][y]; + let w = weights[y]; + for s in 0..ROBP_WIDTH { + local_bwd_sum[cd][s] += w * bwd[i + 1][y][s]; + } + } + let lp0 = dot4(&r_cd[0], &local_bwd_sum[0]) + dot4(&r_cd[1], &local_bwd_sum[1]); + let lp1 = dot4(&r_cd[2], &local_bwd_sum[2]) + dot4(&r_cd[3], &local_bwd_sum[3]); + let tc0 = dot4(&row2_d0, &local_bwd_sum[0]) + dot4(&row2_d1, &local_bwd_sum[1]); + let tc1 = dot4(&row2_d0, &local_bwd_sum[2]) + dot4(&row2_d1, &local_bwd_sum[3]); + let lp2 = tc1.double() - tc0; + ([lp0, lp1, lp2], local_bwd_sum) + }) + .collect(); + + let (mut p0, mut p1, mut p2) = (E::ZERO, E::ZERO, E::ZERO); + for (p_local, _) in &partials { + p0 += p_local[0]; + p1 += p_local[1]; + p2 += p_local[2]; + } + + transcript.append_field_element_ext(&p1); + transcript.append_field_element_ext(&p2); + let alpha = transcript + .sample_and_append_challenge(b"Internal round") + .elements; + challenges.push(alpha); + proof_messages.push(IOPProverMessage { + evaluations: vec![p1, p2], + }); + + // ---- Round 2i+1: bind z₄[i] ---- + // + // Round polynomial for round 2i+1: + // p_{2i+1}(λ) = Σ_y w_y · eq₁(λ, d_y[i]) · fwd · M_i^{(α, λ)} · bwd[i+1][y] + // where w_y logically includes eq₁(α, c_y[i]) from round 2i. + // + // Reuse cached per-thread bwd_sums from round 2i: each absorbs α to + // get local bwd_sum_d, computes local p(0), p(1), p(2), then we sum. + let na = E::ONE - alpha; + // fwd · M_i^{(α, λ)} at λ=0: (1-α)·R_{0,0} + α·R_{1,0} + let row_at_0: StateVec = std::array::from_fn(|j| na * r_cd[0][j] + alpha * r_cd[2][j]); + // at λ=1: (1-α)·R_{0,1} + α·R_{1,1} + let row_at_1: StateVec = std::array::from_fn(|j| na * r_cd[1][j] + alpha * r_cd[3][j]); + // at λ=2: 2·row_at_1 - row_at_0 + let row_at_2: StateVec = std::array::from_fn(|j| row_at_1[j].double() - row_at_0[j]); + + let (mut p0, mut p1, mut p2) = (E::ZERO, E::ZERO, E::ZERO); + for (_, local_bwd_sum) in &partials { + let local_bwd_sum_d: [[E; ROBP_WIDTH]; 2] = [ + std::array::from_fn(|s| na * local_bwd_sum[0][s] + alpha * local_bwd_sum[2][s]), + std::array::from_fn(|s| na * local_bwd_sum[1][s] + alpha * local_bwd_sum[3][s]), + ]; + p0 += dot4(&row_at_0, &local_bwd_sum_d[0]); + p1 += dot4(&row_at_1, &local_bwd_sum_d[1]); + p2 += term_d1_d0_combine( + dot4(&row_at_2, &local_bwd_sum_d[0]), + dot4(&row_at_2, &local_bwd_sum_d[1]), + ); + } + + transcript.append_field_element_ext(&p1); + transcript.append_field_element_ext(&p2); + let beta = transcript + .sample_and_append_challenge(b"Internal round") + .elements; + challenges.push(beta); + proof_messages.push(IOPProverMessage { + evaluations: vec![p1, p2], + }); + + // Fused weight update: w_y *= eq₁(α, c_y[i]) · eq₁(β, d_y[i]) + let nb = E::ONE - beta; + let eq_cd = [na * nb, na * beta, alpha * nb, alpha * beta]; + weights + .par_chunks_mut(batch_size) + .enumerate() + .for_each(|(chunk_idx, w_chunk)| { + let start = chunk_idx * batch_size; + for (j, w) in w_chunk.iter_mut().enumerate() { + let y = start + j; + let cd = c_bits[i][y] * 2 + d_bits[i][y]; + *w *= eq_cd[cd]; + } + }); + + // Update forward vector: fwd ← fwd · M_i^{(α, β)} + fwd = std::array::from_fn(|j| { + eq_cd[0] * r_cd[0][j] + + eq_cd[1] * r_cd[1][j] + + eq_cd[2] * r_cd[2][j] + + eq_cd[3] * r_cd[3][j] + }); + } + + ( + IOPProof { + proofs: proof_messages, + }, + challenges, + ) +} + +/// eq₁(2, d=0) = -1, eq₁(2, d=1) = 2. +#[inline] +fn term_d1_d0_combine(dot_d0: E, dot_d1: E) -> E { + dot_d1.double() - dot_d0 +} + +/// Compute Q(ρ*) = Σ_y eq_col[y] · eq(ρ*, x_y) where x_y are the interleaved +/// Boolean evaluation points. +/// +/// `assist_point` is the interleaved challenge point from the assist sumcheck, +/// of length 2·n_robp. +pub fn compute_q_at_assist_point( + assist_point: &[E], + eq_col: &[E], + cumulative_heights: &[usize], + n_robp: usize, +) -> E { + let num_polys = cumulative_heights.len() - 1; + let mut q_val = E::ZERO; + for y in 0..num_polys { + let mut prod = eq_col[y]; + if prod == E::ZERO { + continue; + } + for i in 0..n_robp { + let c_bit = (cumulative_heights[y] >> i) & 1; + let d_bit = (cumulative_heights[y + 1] >> i) & 1; + let z3_val = assist_point[2 * i]; + let z4_val = assist_point[2 * i + 1]; + prod *= if c_bit == 1 { z3_val } else { E::ONE - z3_val }; + prod *= if d_bit == 1 { z4_val } else { E::ONE - z4_val }; + } + q_val += prod; + } + q_val +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::jagged::{ + evaluator::{evaluate_g, evaluate_g_forward}, + types::int_to_field_bits, + }; + use ff_ext::{FromUniformBytes, GoldilocksExt2}; + use multilinear_extensions::{util::ceil_log2, virtual_poly::build_eq_x_r_vec}; + use p3::field::FieldAlgebra; + use rand::thread_rng; + use std::marker::PhantomData; + use sumcheck::structs::IOPVerifierState; + use transcript::basic::BasicTranscript; + + type E = GoldilocksExt2; + + /// Brute-force f̂(ρ) = Σ_y eq_col[y] · ĝ(z_row, ρ, bits(t_y), bits(t_{y+1})) + fn compute_f_at_point_slow( + z_row: &[E], + rho: &[E], + eq_col: &[E], + cumulative_heights: &[usize], + n_robp: usize, + ) -> E { + let num_polys = cumulative_heights.len() - 1; + let mut val = E::ZERO; + for y in 0..num_polys { + let t_lo = int_to_field_bits::(cumulative_heights[y], n_robp); + let t_hi = int_to_field_bits::(cumulative_heights[y + 1], n_robp); + val += eq_col[y] * evaluate_g(z_row, rho, &t_lo, &t_hi); + } + val + } + + #[test] + fn test_assist_sumcheck_small() { + let mut rng = thread_rng(); + + // 4 polynomials with heights [4, 8, 4, 8] → cumulative [0, 4, 12, 16, 24] + let cumulative_heights = vec![0, 4, 12, 16, 24]; + let num_polys = cumulative_heights.len() - 1; + let total_evals = *cumulative_heights.last().unwrap(); + let num_giga_vars = ceil_log2(total_evals); + let n_robp = num_giga_vars + if total_evals.is_power_of_two() { 1 } else { 0 }; + + let z_row: Vec = (0..n_robp).map(|_| E::random(&mut rng)).collect(); + let rho: Vec = (0..n_robp).map(|_| E::random(&mut rng)).collect(); + let num_col_vars = ceil_log2(num_polys).max(1); + let z_col: Vec = (0..num_col_vars).map(|_| E::random(&mut rng)).collect(); + let eq_col = build_eq_x_r_vec(&z_col); + + let claimed_sum = + compute_f_at_point_slow(&z_row, &rho, &eq_col, &cumulative_heights, n_robp); + + // Run prover. + let mut transcript_p = BasicTranscript::::new(b"assist_test"); + let (proof, challenges) = assist_sumcheck_prove( + &z_row, + &rho, + &eq_col, + &cumulative_heights, + n_robp, + &mut transcript_p, + ); + + let n_vars = 2 * n_robp; + assert_eq!(proof.proofs.len(), n_vars); + assert_eq!(challenges.len(), n_vars); + + // Verify using standard sumcheck verifier. + let mut transcript_v = BasicTranscript::::new(b"assist_test"); + let aux_info = multilinear_extensions::virtual_poly::VPAuxInfo { + max_degree: 2, + max_num_variables: n_vars, + phantom: PhantomData::, + }; + let subclaim = + IOPVerifierState::::verify(claimed_sum, &proof, &aux_info, &mut transcript_v); + + // Check challenges match. + for (sc, ch) in subclaim.point.iter().zip(challenges.iter()) { + assert_eq!(sc.elements, *ch); + } + + // De-interleave the assist point: (z3[0], z4[0], z3[1], z4[1], ...) + let rho_star_c: Vec = (0..n_robp).map(|i| challenges[2 * i]).collect(); + let rho_star_d: Vec = (0..n_robp).map(|i| challenges[2 * i + 1]).collect(); + + // h(ρ*) = ĝ(z_row, ρ, ρ*_c, ρ*_d) + let h_val = evaluate_g(&z_row, &rho, &rho_star_c, &rho_star_d); + + // Q(ρ*) = Σ_y eq_col[y] · eq(ρ*, x_y) + let q_val = compute_q_at_assist_point(&challenges, &eq_col, &cumulative_heights, n_robp); + + assert_eq!( + h_val * q_val, + subclaim.expected_evaluation, + "h(ρ*) * Q(ρ*) != subclaim" + ); + } + + #[test] + fn test_assist_sumcheck_single_poly() { + let mut rng = thread_rng(); + + let cumulative_heights = vec![0, 8]; + let total_evals = 8; + let num_giga_vars = ceil_log2(total_evals); + let n_robp = num_giga_vars + if total_evals.is_power_of_two() { 1 } else { 0 }; + + let z_row: Vec = (0..n_robp).map(|_| E::random(&mut rng)).collect(); + let rho: Vec = (0..n_robp).map(|_| E::random(&mut rng)).collect(); + let eq_col = vec![E::ONE]; // single poly + + let claimed_sum = + compute_f_at_point_slow(&z_row, &rho, &eq_col, &cumulative_heights, n_robp); + + let mut transcript_p = BasicTranscript::::new(b"assist_single"); + let (proof, challenges) = assist_sumcheck_prove( + &z_row, + &rho, + &eq_col, + &cumulative_heights, + n_robp, + &mut transcript_p, + ); + + let n_vars = 2 * n_robp; + let mut transcript_v = BasicTranscript::::new(b"assist_single"); + let aux_info = multilinear_extensions::virtual_poly::VPAuxInfo { + max_degree: 2, + max_num_variables: n_vars, + phantom: PhantomData::, + }; + let subclaim = + IOPVerifierState::::verify(claimed_sum, &proof, &aux_info, &mut transcript_v); + + let rho_star_c: Vec = (0..n_robp).map(|i| challenges[2 * i]).collect(); + let rho_star_d: Vec = (0..n_robp).map(|i| challenges[2 * i + 1]).collect(); + let h_val = evaluate_g(&z_row, &rho, &rho_star_c, &rho_star_d); + let q_val = compute_q_at_assist_point(&challenges, &eq_col, &cumulative_heights, n_robp); + + assert_eq!(h_val * q_val, subclaim.expected_evaluation); + } + + #[test] + fn test_assist_sumcheck_many_polys() { + let mut rng = thread_rng(); + + // 16 polynomials, each height 32 → cumulative [0, 32, 64, ..., 512] + let num_polys = 16; + let poly_height = 32usize; + let cumulative_heights: Vec = (0..=num_polys).map(|i| i * poly_height).collect(); + let total_evals = *cumulative_heights.last().unwrap(); + let num_giga_vars = ceil_log2(total_evals); + let n_robp = num_giga_vars + if total_evals.is_power_of_two() { 1 } else { 0 }; + + let z_row: Vec = (0..n_robp).map(|_| E::random(&mut rng)).collect(); + let rho: Vec = (0..n_robp).map(|_| E::random(&mut rng)).collect(); + let num_col_vars = ceil_log2(num_polys).max(1); + let z_col: Vec = (0..num_col_vars).map(|_| E::random(&mut rng)).collect(); + let eq_col = build_eq_x_r_vec(&z_col); + + let claimed_sum = + compute_f_at_point_slow(&z_row, &rho, &eq_col, &cumulative_heights, n_robp); + + let mut transcript_p = BasicTranscript::::new(b"assist_many"); + let (proof, challenges) = assist_sumcheck_prove( + &z_row, + &rho, + &eq_col, + &cumulative_heights, + n_robp, + &mut transcript_p, + ); + + let n_vars = 2 * n_robp; + let mut transcript_v = BasicTranscript::::new(b"assist_many"); + let aux_info = multilinear_extensions::virtual_poly::VPAuxInfo { + max_degree: 2, + max_num_variables: n_vars, + phantom: PhantomData::, + }; + let subclaim = + IOPVerifierState::::verify(claimed_sum, &proof, &aux_info, &mut transcript_v); + + for (sc, ch) in subclaim.point.iter().zip(challenges.iter()) { + assert_eq!(sc.elements, *ch); + } + + let rho_star_c: Vec = (0..n_robp).map(|i| challenges[2 * i]).collect(); + let rho_star_d: Vec = (0..n_robp).map(|i| challenges[2 * i + 1]).collect(); + let h_val = evaluate_g(&z_row, &rho, &rho_star_c, &rho_star_d); + let q_val = compute_q_at_assist_point(&challenges, &eq_col, &cumulative_heights, n_robp); + + assert_eq!(h_val * q_val, subclaim.expected_evaluation); + } + + /// Also verify using the forward evaluator for cross-validation. + #[test] + fn test_assist_forward_backward_consistency() { + let mut rng = thread_rng(); + + let cumulative_heights = vec![0, 3, 7, 10]; + let num_polys = 3; + let total_evals = 10; + let num_giga_vars = ceil_log2(total_evals); + let n_robp = num_giga_vars + if total_evals.is_power_of_two() { 1 } else { 0 }; + + let z_row: Vec = (0..n_robp).map(|_| E::random(&mut rng)).collect(); + let rho: Vec = (0..n_robp).map(|_| E::random(&mut rng)).collect(); + let eq_col = build_eq_x_r_vec(&vec![E::random(&mut rng); ceil_log2(num_polys).max(1)]); + + // Compute f̂(ρ) using both forward and backward evaluators. + let mut f_fwd = E::ZERO; + let mut f_bwd = E::ZERO; + for y in 0..num_polys { + let t_lo = int_to_field_bits::(cumulative_heights[y], n_robp); + let t_hi = int_to_field_bits::(cumulative_heights[y + 1], n_robp); + f_fwd += eq_col[y] * evaluate_g_forward(&z_row, &rho, &t_lo, &t_hi); + f_bwd += eq_col[y] * evaluate_g(&z_row, &rho, &t_lo, &t_hi); + } + assert_eq!(f_fwd, f_bwd, "forward/backward f̂ disagree"); + + // Run assist sumcheck and verify. + let mut transcript_p = BasicTranscript::::new(b"assist_fb"); + let (proof, challenges) = assist_sumcheck_prove( + &z_row, + &rho, + &eq_col, + &cumulative_heights, + n_robp, + &mut transcript_p, + ); + + let n_vars = 2 * n_robp; + let mut transcript_v = BasicTranscript::::new(b"assist_fb"); + let aux_info = multilinear_extensions::virtual_poly::VPAuxInfo { + max_degree: 2, + max_num_variables: n_vars, + phantom: PhantomData::, + }; + let subclaim = IOPVerifierState::::verify(f_bwd, &proof, &aux_info, &mut transcript_v); + + let rho_star_c: Vec = (0..n_robp).map(|i| challenges[2 * i]).collect(); + let rho_star_d: Vec = (0..n_robp).map(|i| challenges[2 * i + 1]).collect(); + let h_val = evaluate_g(&z_row, &rho, &rho_star_c, &rho_star_d); + let q_val = compute_q_at_assist_point(&challenges, &eq_col, &cumulative_heights, n_robp); + + assert_eq!(h_val * q_val, subclaim.expected_evaluation); + } +} diff --git a/crates/mpcs/src/jagged/evaluator.rs b/crates/mpcs/src/jagged/evaluator.rs index 8377427..5ff1246 100644 --- a/crates/mpcs/src/jagged/evaluator.rs +++ b/crates/mpcs/src/jagged/evaluator.rs @@ -12,8 +12,125 @@ use ff_ext::ExtensionField; // where Tᵢ = Σ_{σ ∈ {0,1}⁴} eq(ζᵢ, σ) · M^(σ) is the eq-weighted // transition matrix at step i, e₁ is the initial state vector, and u is // the sink label vector. +// +// Since multilinear variables can be bound in any order, we fix z₁ and z₂ +// first (to z_row and ρ), which reduces the per-step alphabet from {0,1}⁴ +// to {0,1}² and gives the 4 matrices M_i^{(c,d)}. The remaining variables +// z₃, z₄ are then interleaved as (z₃[0], z₄[0], z₃[1], …) to match the +// ROBP step order, enabling the forward/backward decomposition. // --------------------------------------------------------------------------- +pub const ROBP_WIDTH: usize = 4; + +pub type StateVec = [E; ROBP_WIDTH]; +pub type TransitionMatrix = [[E; ROBP_WIDTH]; ROBP_WIDTH]; + +/// Sink label vector: accept at state (carry=0, lt=1) = index 1. +pub fn sink_labels() -> StateVec { + let mut u = [E::ZERO; ROBP_WIDTH]; + u[1] = E::ONE; + u +} + +/// Initial forward vector: source state (carry=0, lt=0) = index 0. +pub fn source_vec() -> StateVec { + let mut v = [E::ZERO; ROBP_WIDTH]; + v[0] = E::ONE; + v +} + +#[inline] +pub fn dot4(a: &StateVec, b: &StateVec) -> E { + a[0] * b[0] + a[1] * b[1] + a[2] * b[2] + a[3] * b[3] +} + +/// Row-vector × matrix: out[j] = Σ_i v[i] * m[i][j]. +#[inline] +pub fn vec_mat_mul(v: &StateVec, m: &TransitionMatrix) -> StateVec { + std::array::from_fn(|j| v[0] * m[0][j] + v[1] * m[1][j] + v[2] * m[2][j] + v[3] * m[3][j]) +} + +/// Matrix × column-vector: out[i] = Σ_j m[i][j] * v[j]. +#[inline] +pub fn mat_vec_mul(m: &TransitionMatrix, v: &StateVec) -> StateVec { + std::array::from_fn(|i| m[i][0] * v[0] + m[i][1] * v[1] + m[i][2] * v[2] + m[i][3] * v[3]) +} + +/// Raw ROBP transition table for the indicator g(a,b,c,d) = [a+c=b ∧ b> 1; + let lt_in = from & 1; + let mut sym = 0u8; + while sym < 16 { + let a = sym >> 3; + let b = (sym >> 2) & 1; + let c = (sym >> 1) & 1; + let d = sym & 1; + let sum = a + c + carry_in; + if sum & 1 == b { + let carry_out = sum >> 1; + let lt_out = if b < d { + 1 + } else if b == d { + lt_in + } else { + 0 + }; + table[from as usize][sym as usize] = carry_out * 2 + lt_out; + } + sym += 1; + } + from += 1; + } + table +}; + +/// Per-symbol transition matrices M_i^{(c,d)} at one ROBP step, with +/// z1 (= z_row[i]) and z2 (= ρ[i]) eq-weights baked in. +/// +/// Returns 4 matrices indexed by `c * 2 + d` for `(c, d) ∈ {0,1}²`. +/// Matrix entry `m[from][to]` gives the transition weight from `from` to `to` +/// when reading symbol `(c, d)`. +/// +/// Derived from the raw ROBP transition table via: +/// M_i^{(c,d)}[from][to] = Σ_{a,b} eq₁(z1,a) · eq₁(z2,b) · [transition(from,(a,b,c,d)) = to] +/// +/// The full eq-weighted transition matrix is recovered as: +/// T_i = Σ_{c,d} eq₁(z3, c) · eq₁(z4, d) · M_i^{(c,d)} +pub fn symbol_transition_matrices(z1i: E, z2i: E) -> [TransitionMatrix; 4] { + let eq_ab: [E; 4] = { + let (nz1, nz2) = (E::ONE - z1i, E::ONE - z2i); + [nz1 * nz2, nz1 * z2i, z1i * nz2, z1i * z2i] + }; + + let mut mats = [[[E::ZERO; ROBP_WIDTH]; ROBP_WIDTH]; 4]; + for cd in 0..4u8 { + #[allow(clippy::needless_range_loop)] + for from in 0..ROBP_WIDTH { + for ab in 0..4u8 { + let sym = (ab << 2) | cd; + let to = ROBP_TRANSITION[from][sym as usize]; + if to != REJECT { + mats[cd as usize][from][to as usize] += eq_ab[ab as usize]; + } + } + } + } + mats +} + /// Compute the eq-weighted transition matrix at step `i`. /// /// Returns `(T_same, T_lt1, T_lt0)`: three 2×2 matrices (indexed by carry) @@ -191,7 +308,7 @@ mod tests { use p3::field::FieldAlgebra; use rand::thread_rng; - use super::{evaluate_g_backward, evaluate_g_forward}; + use super::*; type E = BabyBearExt4; @@ -263,4 +380,94 @@ mod tests { assert_eq!(fwd, bwd, "forward != backward at n={n}"); } } + + /// Verify that recombining per-symbol matrices reproduces transition_weights. + #[test] + fn test_symbol_matrices_match_transition_weights() { + let mut rng = thread_rng(); + for _ in 0..20 { + let z1i = E::random(&mut rng); + let z2i = E::random(&mut rng); + let z3i = E::random(&mut rng); + let z4i = E::random(&mut rng); + + let mats = symbol_transition_matrices(z1i, z2i); + let (nz3, nz4) = (E::ONE - z3i, E::ONE - z4i); + let cd_weights = [nz3 * nz4, nz3 * z4i, z3i * nz4, z3i * z4i]; + + // Reconstruct T = Σ_{c,d} eq₁(z3,c)·eq₁(z4,d)·M^{(c,d)} + let mut t_recon = [[E::ZERO; 4]; 4]; + for cd in 0..4 { + #[allow(clippy::needless_range_loop)] + for i in 0..4 { + for j in 0..4 { + t_recon[i][j] += cd_weights[cd] * mats[cd][i][j]; + } + } + } + + // Compare against transition_weights-based backward step. + let transitions = transition_weights(z1i, z2i, z3i, z4i); + let mut t_expected = [[E::ZERO; 4]; 4]; + for &(ci, co, w_same, w_lt1, w_lt0) in &transitions { + // M[from][to] convention, matching backward pass. + t_expected[ci * 2][co * 2] += w_same + w_lt0; + t_expected[ci * 2][co * 2 + 1] += w_lt1; + t_expected[ci * 2 + 1][co * 2] += w_lt0; + t_expected[ci * 2 + 1][co * 2 + 1] += w_same + w_lt1; + } + + assert_eq!( + t_recon, t_expected, + "symbol matrices don't match transition_weights" + ); + } + } + + /// Verify backward precomputation: using per-symbol matrices with Boolean + /// (c, d) reproduces evaluate_g. + #[test] + fn test_backward_via_symbol_matrices() { + let mut rng = thread_rng(); + for n in 1..=5 { + let z1: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + let z2: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + let z3: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + let z4: Vec = (0..n).map(|_| E::random(&mut rng)).collect(); + + let expected = evaluate_g_backward(&z1, &z2, &z3, &z4); + + // Compute using per-symbol matrices with Boolean (c, d) from z3/z4. + // Round each z3[i]/z4[i] to the nearest bit for this test — instead, + // use evaluate_g with actual field elements and compare using the + // matrix product approach. + let step_mats: Vec<_> = (0..n) + .map(|i| symbol_transition_matrices(z1[i], z2[i])) + .collect(); + + // Build T_i = Σ_{c,d} eq1(z3[i],c)*eq1(z4[i],d)*M^{(c,d)} + // and multiply backward. + let u = sink_labels::(); + let mut val = u; + for i in (0..n).rev() { + let (nz3, nz4) = (E::ONE - z3[i], E::ONE - z4[i]); + let cd_w = [nz3 * nz4, nz3 * z4[i], z3[i] * nz4, z3[i] * z4[i]]; + let mut t_i = [[E::ZERO; 4]; 4]; + #[allow(clippy::needless_range_loop)] + for cd in 0..4 { + for r in 0..4 { + for c in 0..4 { + t_i[r][c] += cd_w[cd] * step_mats[i][cd][r][c]; + } + } + } + val = mat_vec_mul(&t_i, &val); + } + let result = val[0]; // source state + assert_eq!( + result, expected, + "symbol-matrix backward != evaluate_g at n={n}" + ); + } + } } diff --git a/crates/mpcs/src/jagged/mod.rs b/crates/mpcs/src/jagged/mod.rs index c62118f..d29a214 100644 --- a/crates/mpcs/src/jagged/mod.rs +++ b/crates/mpcs/src/jagged/mod.rs @@ -139,10 +139,12 @@ //! `h(ρ)`. The ROBP makes `ĝ` efficiently evaluable, so the verifier computes //! `h(ρ) = Σ_i eq(z_c, i) · ĝ(z_r, ρ, t[i], t[i+1])` in `O(K·n)` time. +pub mod assist; pub mod evaluator; pub mod sumcheck; mod types; +pub use assist::{assist_sumcheck_prove, compute_q_at_assist_point}; pub use evaluator::{evaluate_g, evaluate_g_backward, evaluate_g_forward}; pub use sumcheck::{JaggedSumcheckInput, jagged_sumcheck_prove}; pub use types::{JaggedBatchOpenProof, JaggedCommitment, JaggedCommitmentWithWitness}; @@ -160,9 +162,7 @@ use multilinear_extensions::{ }; use p3::{ matrix::{Matrix, bitrev::BitReversableMatrix}, - maybe_rayon::prelude::{ - IndexedParallelIterator, IntoParallelIterator, ParallelIterator, ParallelSliceMut, - }, + maybe_rayon::prelude::*, }; use transcript::Transcript; use types::int_to_field_bits; @@ -396,26 +396,54 @@ pub fn jagged_batch_open, + }; + let assist_subclaim = + IOPVerifierState::::verify(proof.f_at_rho, &proof.assist_proof, &assist_aux, transcript); + let assist_point: Vec = assist_subclaim.point.iter().map(|c| c.elements).collect(); + + // De-interleave the assist point: (z3[0], z4[0], z3[1], z4[1], ...) + let rho_star_c: Vec = (0..n_robp).map(|i| assist_point[2 * i]).collect(); + let rho_star_d: Vec = (0..n_robp).map(|i| assist_point[2 * i + 1]).collect(); + + // h(ρ*) = ĝ(z_row_padded, ρ_padded, ρ*_c, ρ*_d) — one ROBP evaluation. let mut z_row_padded = z_row; z_row_padded.resize(n_robp, E::ZERO); let mut rho_padded = rho.clone(); rho_padded.resize(n_robp, E::ZERO); - let f_at_rho = compute_f_at_point( - &z_row_padded, - &rho_padded, - &eq_col, - &comm.cumulative_heights, - n_robp, - ); + let h_at_rho_star = evaluate_g(&z_row_padded, &rho_padded, &rho_star_c, &rho_star_d); - // Write q_eval to transcript (must match prover). - transcript.append_field_element_ext(&proof.q_eval); + // Q(ρ*) = Σ_y eq_col[y] · eq(ρ*, x_y). + let q_at_rho_star = + compute_q_at_assist_point(&assist_point, &eq_col, &comm.cumulative_heights, n_robp); - // Check subclaim: q'(ρ) · f̂(ρ) == expected_evaluation. - if proof.q_eval * f_at_rho != subclaim.expected_evaluation { + // Check assist subclaim: h(ρ*) · Q(ρ*) == expected_evaluation. + if h_at_rho_star * q_at_rho_star != assist_subclaim.expected_evaluation { return Err(Error::InvalidPcsOpen( - "jagged_batch_verify: q_eval * f(rho) != subclaim expected evaluation".into(), + "jagged_batch_verify: assist sumcheck final check failed".into(), )); } diff --git a/crates/mpcs/src/jagged/sumcheck.rs b/crates/mpcs/src/jagged/sumcheck.rs index 0ebe66f..47d21e1 100644 --- a/crates/mpcs/src/jagged/sumcheck.rs +++ b/crates/mpcs/src/jagged/sumcheck.rs @@ -7,10 +7,7 @@ use ff_ext::ExtensionField; use multilinear_extensions::{ mle::MultilinearExtension, util::max_usable_threads, virtual_poly::build_eq_x_r_vec, }; -use p3::maybe_rayon::prelude::{ - IndexedParallelIterator, IntoParallelIterator, ParallelIterator, ParallelSlice, - ParallelSliceMut, -}; +use p3::maybe_rayon::prelude::*; use sumcheck::{ macros::{entered_span, exit_span}, structs::{IOPProof, IOPProverMessage, IOPProverState}, diff --git a/crates/mpcs/src/jagged/types.rs b/crates/mpcs/src/jagged/types.rs index bc2381e..8dfa509 100644 --- a/crates/mpcs/src/jagged/types.rs +++ b/crates/mpcs/src/jagged/types.rs @@ -58,12 +58,15 @@ impl> /// Proof for the jagged batch opening protocol. /// /// Contains a sumcheck proof (reducing K column evaluation claims to a single -/// point on q'), the evaluation q'(ρ), and an inner PCS opening proof. +/// point on q'), the evaluation q'(ρ), an assist sumcheck proof (reducing K +/// ROBP evaluations to one), and an inner PCS opening proof. #[derive(Clone, Serialize, Deserialize)] #[serde(bound(serialize = "", deserialize = ""))] pub struct JaggedBatchOpenProof> { pub sumcheck_proof: IOPProof, pub q_eval: E, + pub f_at_rho: E, + pub assist_proof: IOPProof, pub inner_proof: InnerPcs::Proof, } diff --git a/crates/mpcs/src/lib.rs b/crates/mpcs/src/lib.rs index cfbb5fb..f37669a 100644 --- a/crates/mpcs/src/lib.rs +++ b/crates/mpcs/src/lib.rs @@ -270,8 +270,8 @@ pub use basefold::{ pub mod jagged; pub use jagged::{ JaggedBatchOpenProof, JaggedCommitment, JaggedCommitmentWithWitness, JaggedSumcheckInput, - evaluate_g, evaluate_g_backward, evaluate_g_forward, jagged_batch_open, jagged_batch_verify, - jagged_commit, jagged_sumcheck_prove, + assist_sumcheck_prove, evaluate_g, evaluate_g_backward, evaluate_g_forward, jagged_batch_open, + jagged_batch_verify, jagged_commit, jagged_sumcheck_prove, }; #[cfg(feature = "whir")] extern crate whir as whir_external; From 560f5a9b5b351edbe68b7e58c9a43d30f2b3172e Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Thu, 30 Apr 2026 12:59:03 +0000 Subject: [PATCH 16/21] docs: label batched opening claim and document correction factor inclusion in prover Agent-Logs-Url: https://github.com/scroll-tech/gkr-backend/sessions/22b3ba65-8e69-45eb-ae09-0a5267f8b518 Co-authored-by: kunxian-xia <1082586+kunxian-xia@users.noreply.github.com> --- crates/mpcs/src/jagged/mod.rs | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/crates/mpcs/src/jagged/mod.rs b/crates/mpcs/src/jagged/mod.rs index d29a214..b89aa8d 100644 --- a/crates/mpcs/src/jagged/mod.rs +++ b/crates/mpcs/src/jagged/mod.rs @@ -389,7 +389,11 @@ pub fn jagged_batch_open = transcript.sample_and_append_vec(b"jagged_z_col", num_col_vars); - // claimed_sum = Σ_i eq_col[i] · C_i · evals[i] + // Batched opening claim: v = Σ_i eq_col[i] · C_i · evals[i] // where C_i = eq(z_row[s_i..], 0) = Π_{j=s_i}^{max_s-1} (1 - z_row[j]) is the // correction factor from zero-padding p_i to max_s variables (see module docs). let eq_col = build_eq_x_r_vec(&z_col); From 41d40f7f53e34abf273bd9d8b541eac782c90af5 Mon Sep 17 00:00:00 2001 From: xkx Date: Fri, 1 May 2026 01:08:22 +0800 Subject: [PATCH 17/21] feat: add reshape step to jagged PCS (#48) --- crates/mpcs/benches/jagged_pcs.rs | 95 ++++++----- crates/mpcs/src/jagged/mod.rs | 261 +++++++++++++++++++++++++----- crates/mpcs/src/jagged/types.rs | 8 +- 3 files changed, 276 insertions(+), 88 deletions(-) diff --git a/crates/mpcs/benches/jagged_pcs.rs b/crates/mpcs/benches/jagged_pcs.rs index ba44ea4..4b9650f 100644 --- a/crates/mpcs/benches/jagged_pcs.rs +++ b/crates/mpcs/benches/jagged_pcs.rs @@ -59,38 +59,13 @@ fn bench_jagged_pcs(c: &mut Criterion) { let max_s = *log_heights.iter().max().unwrap(); let total_evals: usize = rmms.iter().map(|rmm| rmm.height() * rmm.width()).sum(); let num_giga_vars = ceil_log2(total_evals); - let num_polys = NUM_MATRICES * NUM_COLS; println!( - "num_matrices={}, num_cols={}, num_polys={}, total_evals={}, num_giga_vars={}, max_s={}", - NUM_MATRICES, NUM_COLS, num_polys, total_evals, num_giga_vars, max_s + "num_matrices={}, num_cols={}, total_evals={}, num_giga_vars={}, max_s={}", + NUM_MATRICES, NUM_COLS, total_evals, num_giga_vars, max_s ); - let poly_size = 1usize << num_giga_vars; - let param = Pcs::setup(poly_size, SecurityLevel::Conjecture100bits).unwrap(); - let (pp, vp) = Pcs::trim(param, poly_size).unwrap(); - - // --- Bench commit --- - let comm = jagged_commit::(&pp, rmms.clone()).expect("commit failed"); - - group.bench_function( - BenchmarkId::new("commit", format!("{}x{}", NUM_MATRICES, NUM_COLS)), - |b| { - b.iter_custom(|iters| { - let mut time = Duration::new(0, 0); - for _ in 0..iters { - let rmms_clone = rmms.clone(); - let instant = std::time::Instant::now(); - let _ = jagged_commit::(&pp, rmms_clone).unwrap(); - time += instant.elapsed(); - } - time - }) - }, - ); - - // --- Prepare evaluation data --- - // Extract column polynomials (before bit-reversal) for computing true evaluations. + // Prepare evaluation data (shared across all reshape_log_height values). let col_polys: Vec> = rmms .iter() .flat_map(|rmm| { @@ -112,10 +87,42 @@ fn bench_jagged_pcs(c: &mut Criterion) { .map(|(col, s_i)| eval_column_poly_at_point(col, &point[(max_s - s_i)..])) .collect(); - // --- Bench batch open --- - group.bench_function( - BenchmarkId::new("batch_open", format!("{}x{}", NUM_MATRICES, NUM_COLS)), - |b| { + // Sweep reshape_log_height values: baseline (no reshape) + a few smaller heights. + let reshape_log_heights: Vec = { + let mut vals = vec![num_giga_vars]; + for step in [2, 4, 6] { + if num_giga_vars > step { + vals.push(num_giga_vars - step); + } + } + vals + }; + + for &log_h in &reshape_log_heights { + let h = 1usize << log_h; + let w = total_evals.div_ceil(h); + let label = format!("log_h={log_h}_w={w}"); + + let poly_size = 1usize << log_h; + let param = Pcs::setup(poly_size, SecurityLevel::Conjecture100bits).unwrap(); + let (pp, vp) = Pcs::trim(param, poly_size).unwrap(); + + let comm = jagged_commit::(&pp, rmms.clone(), log_h).expect("commit failed"); + + group.bench_function(BenchmarkId::new("commit", &label), |b| { + b.iter_custom(|iters| { + let mut time = Duration::new(0, 0); + for _ in 0..iters { + let rmms_clone = rmms.clone(); + let instant = std::time::Instant::now(); + let _ = jagged_commit::(&pp, rmms_clone, log_h).unwrap(); + time += instant.elapsed(); + } + time + }) + }); + + group.bench_function(BenchmarkId::new("batch_open", &label), |b| { b.iter_batched( || { let mut transcript = BasicTranscript::::new(b"jagged_bench"); @@ -128,18 +135,18 @@ fn bench_jagged_pcs(c: &mut Criterion) { }, BatchSize::SmallInput, ); - }, - ); + }); + + let mut transcript_p = BasicTranscript::::new(b"jagged_bench"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); + let proof = + jagged_batch_open::(&pp, &comm, &point, &evals, &mut transcript_p).unwrap(); + let pure_comm = comm.to_commitment(); - // --- Bench batch verify --- - let mut transcript_p = BasicTranscript::::new(b"jagged_bench"); - Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); - let proof = jagged_batch_open::(&pp, &comm, &point, &evals, &mut transcript_p).unwrap(); - let pure_comm = comm.to_commitment(); + let proof_size = bincode::serialize(&proof).map(|v| v.len()).unwrap_or(0); + println!("{label}: proof_size={proof_size} bytes, col_evals.len={w}"); - group.bench_function( - BenchmarkId::new("batch_verify", format!("{}x{}", NUM_MATRICES, NUM_COLS)), - |b| { + group.bench_function(BenchmarkId::new("batch_verify", &label), |b| { b.iter_batched( || { let mut transcript = BasicTranscript::::new(b"jagged_bench"); @@ -159,8 +166,8 @@ fn bench_jagged_pcs(c: &mut Criterion) { }, BatchSize::SmallInput, ); - }, - ); + }); + } group.finish(); } diff --git a/crates/mpcs/src/jagged/mod.rs b/crates/mpcs/src/jagged/mod.rs index b89aa8d..1206a7a 100644 --- a/crates/mpcs/src/jagged/mod.rs +++ b/crates/mpcs/src/jagged/mod.rs @@ -161,6 +161,7 @@ use multilinear_extensions::{ virtual_poly::{VPAuxInfo, build_eq_x_r_vec}, }; use p3::{ + field::FieldAlgebra, matrix::{Matrix, bitrev::BitReversableMatrix}, maybe_rayon::prelude::*, }; @@ -188,6 +189,7 @@ use witness::{InstancePaddingStrategy, RowMajorMatrix}; pub fn jagged_commit>( pp: &InnerPcs::ProverParam, rmms: Vec>, + reshape_log_height: usize, ) -> Result, Error> { if rmms.is_empty() { return Err(Error::InvalidPcsParam( @@ -273,11 +275,47 @@ pub fn jagged_commit> poly_idx += n_cols; } - // --- Step 4: Commit via the inner PCS --- - // q' is committed as a single-column matrix with height = total_size. + // --- Step 4: Reshape and commit via the inner PCS --- + let log_h = reshape_log_height; + let h = 1usize << log_h; + let w = total_size.div_ceil(h); + + let giga_data = if w == 1 { + concatenated + } else { + // Transpose the flat q' evaluations to row-major for the inner PCS. + // In the flat array, column i occupies concatenated[i*h .. (i+1)*h]; + // the last column may be short and is zero-padded. + let padded_total = w * h; + let mut row_major: Vec = Vec::with_capacity(padded_total); + #[allow(clippy::uninit_vec)] + unsafe { + row_major.set_len(padded_total) + }; + + (0..w).into_par_iter().for_each(|i| { + let src_start = i * h; + let col_len = if i < w - 1 { h } else { total_size - src_start }; + let dst = unsafe { + &mut *std::ptr::slice_from_raw_parts_mut( + row_major.as_ptr() as *mut E::BaseField, + padded_total, + ) + }; + for b in 0..col_len { + dst[b * w + i] = concatenated[src_start + b]; + } + for b in col_len..h { + dst[b * w + i] = E::BaseField::ZERO; + } + }); + + row_major + }; + let giga_rmm = RowMajorMatrix::::new_by_values( - concatenated, - 1, // width = 1 (single polynomial q') + giga_data, + w, InstancePaddingStrategy::Default, ); @@ -287,6 +325,7 @@ pub fn jagged_commit> inner, cumulative_heights, poly_heights, + reshape_log_height, }) } @@ -360,7 +399,11 @@ pub fn jagged_batch_open = point.iter().rev().cloned().collect(); @@ -372,21 +415,34 @@ pub fn jagged_batch_open slice, - _ => { - return Err(Error::InvalidPcsParam( - "jagged_batch_open: expected base-field evaluations for q'".into(), - )); + assert_eq!(q_mles.len(), w); + let q_evals_base_owned: Vec; + let q_evals_base: &[E::BaseField] = if w == 1 { + match q_mles[0].evaluations() { + FieldType::Base(slice) => slice, + _ => { + return Err(Error::InvalidPcsParam( + "jagged_batch_open: expected base-field evaluations for q'".into(), + )); + } } + } else { + let target_len = 1usize << num_giga_vars; + q_evals_base_owned = q_mles + .iter() + .flat_map(|mle| match mle.evaluations() { + FieldType::Base(slice) => slice.iter().copied(), + _ => unreachable!("expected base-field evaluations"), + }) + .chain(std::iter::repeat_n( + E::BaseField::ZERO, + target_len - padded_total, + )) + .collect(); + &q_evals_base_owned }; // Batched opening claim: v = Σ_i eq_col[i] · C_i · p_i(z_row[..s_i]) @@ -405,11 +461,8 @@ pub fn jagged_batch_open = q_mles.par_iter().map(|mle| mle.evaluate(rho_row)).collect(); + + // Write col_evals and f_at_rho to transcript. + transcript.append_field_element_exts(&col_evals); transcript.append_field_element_ext(&f_at_rho); // Run the assist sumcheck to prove f̂(ρ) is correct. @@ -436,16 +496,16 @@ pub fn jagged_batch_open = subclaim.point.iter().map(|c| c.elements).collect(); // The ROBP needs enough bits to represent the max cumulative height (= total_evals). - // When total_evals is an exact power of 2, num_giga_vars bits can't represent it. - let n_robp = num_giga_vars + if total_evals.is_power_of_two() { 1 } else { 0 }; + // When padded_total is an exact power of 2, num_giga_vars bits can't represent it. + let n_robp = num_giga_vars + if padded_total.is_power_of_two() { 1 } else { 0 }; - // Write q_eval and f_at_rho to transcript (must match prover). - transcript.append_field_element_ext(&proof.q_eval); + // Split ρ: low bits = row point, high bits = column selector. + let rho_row = &rho[..log_h]; + let rho_col = &rho[log_h..]; + + // Reconstruct q'(ρ) from col_evals: q'(ρ) = Σ_{i=0}^{w-1} eq(ρ_col, i) · col_evals[i]. + let eq_rho_col = build_eq_x_r_vec(rho_col); + let q_eval: E = eq_rho_col[..w] + .iter() + .zip(&proof.col_evals) + .map(|(e, v)| *e * *v) + .sum(); + + // Write col_evals and f_at_rho to transcript (must match prover). + transcript.append_field_element_exts(&proof.col_evals); transcript.append_field_element_ext(&proof.f_at_rho); // Check multiplicative subclaim: q'(ρ) · f̂(ρ) == expected_evaluation. - if proof.q_eval * proof.f_at_rho != subclaim.expected_evaluation { + if q_eval * proof.f_at_rho != subclaim.expected_evaluation { return Err(Error::InvalidPcsOpen( "jagged_batch_verify: q_eval * f(rho) != subclaim expected evaluation".into(), )); @@ -560,12 +636,12 @@ pub fn jagged_batch_verify(4); + let num_giga_vars = 4; + let (pp, _vp) = setup_pcs::(num_giga_vars); let m1 = make_rmm(4, 1); let m2 = make_rmm(4, 2); - let comm = jagged_commit::(&pp, vec![m1, m2]).expect("commit should succeed"); + let comm = jagged_commit::(&pp, vec![m1, m2], num_giga_vars) + .expect("commit should succeed"); assert_eq!(comm.num_polys(), 3); assert_eq!(comm.poly_heights, vec![4, 4, 4]); @@ -654,10 +732,12 @@ mod tests { #[test] fn test_jagged_commit_single_poly() { // 8x1 matrix → 1 polynomial, 8 evals, no padding needed - let (pp, _vp) = setup_pcs::(3); + let num_giga_vars = 3; + let (pp, _vp) = setup_pcs::(num_giga_vars); let m = make_rmm(8, 1); - let comm = jagged_commit::(&pp, vec![m]).expect("commit should succeed"); + let comm = + jagged_commit::(&pp, vec![m], num_giga_vars).expect("commit should succeed"); assert_eq!(comm.num_polys(), 1); assert_eq!(comm.poly_heights, vec![8]); @@ -668,7 +748,7 @@ mod tests { #[test] fn test_jagged_commit_empty_error() { let (pp, _vp) = setup_pcs::(4); - let result = jagged_commit::(&pp, vec![]); + let result = jagged_commit::(&pp, vec![], 4); assert!(matches!(result, Err(Error::InvalidPcsParam(_)))); } @@ -709,7 +789,8 @@ mod tests { // Commit. let mut transcript_p = BasicTranscript::::new(b"jagged_batch_test"); - let comm = jagged_commit::(&pp, rmms).expect("commit should succeed"); + let comm = + jagged_commit::(&pp, rmms, num_giga_vars).expect("commit should succeed"); Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); // Random evaluation point of length max_s. @@ -756,7 +837,7 @@ mod tests { let col_poly: Vec = (0..num_rows).map(|r| rmm.values[r]).collect(); let mut transcript_p = BasicTranscript::::new(b"jagged_single"); - let comm = jagged_commit::(&pp, vec![rmm]).expect("commit"); + let comm = jagged_commit::(&pp, vec![rmm], num_giga_vars).expect("commit"); Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); let point: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); @@ -800,7 +881,7 @@ mod tests { .collect(); let mut transcript_p = BasicTranscript::::new(b"jagged_soundness"); - let comm = jagged_commit::(&pp, vec![rmm]).expect("commit"); + let comm = jagged_commit::(&pp, vec![rmm], num_giga_vars).expect("commit"); Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); let point: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); @@ -828,4 +909,98 @@ mod tests { ); assert!(result.is_err(), "verify should reject tampered evaluations"); } + + // --- Reshape tests (reshape_log_height < num_giga_vars) --- + + #[test] + fn test_jagged_reshape_single_poly() { + let mut rng = thread_rng(); + + let num_rows = 1024usize; // s=10 + let s = 10; + let total_evals = num_rows; + let reshape_log_height = 6; // h=64, w=ceil(1024/64)=16 + let h = 1usize << reshape_log_height; + let w = total_evals.div_ceil(h); + assert_eq!(w, 16); + + let (pp, vp) = setup_pcs::(reshape_log_height); + let rmm = make_rmm(num_rows, 1); + let col_poly: Vec = (0..num_rows).map(|r| rmm.values[r]).collect(); + + let mut transcript_p = BasicTranscript::::new(b"jagged_reshape_single"); + let comm = jagged_commit::(&pp, vec![rmm], reshape_log_height).expect("commit"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); + + let point: Vec = (0..s).map(|_| E::random(&mut rng)).collect(); + let evals = vec![eval_column_poly_at_point(&col_poly, &point)]; + + let proof = jagged_batch_open::(&pp, &comm, &point, &evals, &mut transcript_p) + .expect("batch open"); + assert_eq!(proof.col_evals.len(), w); + + let mut transcript_v = BasicTranscript::::new(b"jagged_reshape_single"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_v).unwrap(); + + jagged_batch_verify::( + &vp, + &comm.to_commitment(), + &point, + &evals, + &proof, + &mut transcript_v, + ) + .expect("batch verify"); + } + + #[test] + fn test_jagged_reshape_multiple_polys() { + let mut rng = thread_rng(); + + // 3 matrices with different heights: 2^10, 2^11, 2^9 (each 1 column). + let log_heights = [10usize, 11, 9]; + let heights: Vec = log_heights.iter().map(|&s| 1 << s).collect(); + let max_s = *log_heights.iter().max().unwrap(); + let total_evals: usize = heights.iter().sum(); + let reshape_log_height = 8; // h=256 + let h = 1usize << reshape_log_height; + let w = total_evals.div_ceil(h); + + let (pp, vp) = setup_pcs::(reshape_log_height); + let rmms: Vec<_> = heights.iter().map(|&h| make_rmm(h, 1)).collect(); + + let col_polys: Vec> = rmms + .iter() + .map(|rmm| (0..rmm.height()).map(|r| rmm.values[r]).collect()) + .collect(); + + let mut transcript_p = BasicTranscript::::new(b"jagged_reshape_multi"); + let comm = + jagged_commit::(&pp, rmms, reshape_log_height).expect("commit should succeed"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); + + let point: Vec = (0..max_s).map(|_| E::random(&mut rng)).collect(); + let evals: Vec = col_polys + .iter() + .zip(log_heights.iter()) + .map(|(col, &s_i)| eval_column_poly_at_point(col, &point[(max_s - s_i)..])) + .collect(); + + let proof = jagged_batch_open::(&pp, &comm, &point, &evals, &mut transcript_p) + .expect("batch open should succeed"); + assert_eq!(proof.col_evals.len(), w); + + let mut transcript_v = BasicTranscript::::new(b"jagged_reshape_multi"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_v).unwrap(); + + jagged_batch_verify::( + &vp, + &comm.to_commitment(), + &point, + &evals, + &proof, + &mut transcript_v, + ) + .expect("batch verify should succeed"); + } } diff --git a/crates/mpcs/src/jagged/types.rs b/crates/mpcs/src/jagged/types.rs index 8dfa509..4c9c38b 100644 --- a/crates/mpcs/src/jagged/types.rs +++ b/crates/mpcs/src/jagged/types.rs @@ -19,6 +19,9 @@ pub struct JaggedCommitmentWithWitness, + /// Log2 of the reshape column height. Each column MLE has this many variables. + /// When equal to `num_giga_vars`, no reshape is performed (single column). + pub reshape_log_height: usize, } /// The pure commitment (without witness data) for a jagged polynomial `q'`. @@ -30,6 +33,8 @@ pub struct JaggedCommitment, + /// Log2 of the reshape column height (verifier needs this to split ρ). + pub reshape_log_height: usize, } impl> @@ -40,6 +45,7 @@ impl> JaggedCommitment { inner: InnerPcs::get_pure_commitment(&self.inner), cumulative_heights: self.cumulative_heights.clone(), + reshape_log_height: self.reshape_log_height, } } @@ -64,7 +70,7 @@ impl> #[serde(bound(serialize = "", deserialize = ""))] pub struct JaggedBatchOpenProof> { pub sumcheck_proof: IOPProof, - pub q_eval: E, + pub col_evals: Vec, pub f_at_rho: E, pub assist_proof: IOPProof, pub inner_proof: InnerPcs::Proof, From c530e7f722990b993a263034acf583bf98078904 Mon Sep 17 00:00:00 2001 From: kunxian xia Date: Thu, 7 May 2026 22:29:17 +0800 Subject: [PATCH 18/21] feat: make JaggedSumcheckInput::compute_claimed_sum public --- crates/mpcs/src/jagged/sumcheck.rs | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/crates/mpcs/src/jagged/sumcheck.rs b/crates/mpcs/src/jagged/sumcheck.rs index 47d21e1..243dd99 100644 --- a/crates/mpcs/src/jagged/sumcheck.rs +++ b/crates/mpcs/src/jagged/sumcheck.rs @@ -88,8 +88,7 @@ impl<'a, E: ExtensionField> JaggedSumcheckInput<'a, E> { /// Brute-force computation of sum_b q'(b) * f(b). /// O(2^n) time — only for debugging and tests. - #[cfg(test)] - fn compute_claimed_sum(&self) -> E { + pub fn compute_claimed_sum(&self) -> E { self.q_evals[..self.total_evaluations()] .iter() .zip(self.col_row_iter(0)) From fba57ba001f43dff420d15e2479bc9e5840bc703 Mon Sep 17 00:00:00 2001 From: kunxian xia Date: Fri, 8 May 2026 23:29:02 +0800 Subject: [PATCH 19/21] feat: add jagged vs direct PCS comparison bench and use p3 RowMajorMatrix - Add comparison benchmark (comparison.rs) measuring commit, batch_open, batch_verify, and proof size for jagged PCS vs direct inner PCS - Refactor jagged_commit to accept p3::matrix::dense::RowMajorMatrix instead of witness::RowMajorMatrix, supporting non-power-of-two matrix heights (internally padded to next power of two before bit-reversal) - Update jagged_pcs bench to use BabyBearExt4, parallel make_rmm, jittered non-power-of-two heights, and eq-table-based column evaluation - Cap reshape_log_height to 25 to fit BabyBear two-adicity constraint Co-Authored-By: Claude Opus 4.6 --- crates/mpcs/Cargo.toml | 4 + crates/mpcs/benches/comparison.rs | 329 ++++++++++++++++++++++++++++++ crates/mpcs/benches/jagged_pcs.rs | 86 ++++---- crates/mpcs/src/jagged/mod.rs | 23 ++- 4 files changed, 392 insertions(+), 50 deletions(-) create mode 100644 crates/mpcs/benches/comparison.rs diff --git a/crates/mpcs/Cargo.toml b/crates/mpcs/Cargo.toml index fe997f9..d7598ea 100644 --- a/crates/mpcs/Cargo.toml +++ b/crates/mpcs/Cargo.toml @@ -62,3 +62,7 @@ name = "jagged_sumcheck" [[bench]] harness = false name = "jagged_pcs" + +[[bench]] +harness = false +name = "comparison" diff --git a/crates/mpcs/benches/comparison.rs b/crates/mpcs/benches/comparison.rs new file mode 100644 index 0000000..c5ee62d --- /dev/null +++ b/crates/mpcs/benches/comparison.rs @@ -0,0 +1,329 @@ +use std::time::Duration; + +use criterion::*; +use ff_ext::{BabyBearExt4, FromUniformBytes}; +use mpcs::{ + Basefold, BasefoldRSParams, PolynomialCommitmentScheme, SecurityLevel, jagged_batch_open, + jagged_batch_verify, jagged_commit, +}; +use multilinear_extensions::{util::ceil_log2, virtual_poly::build_eq_x_r_vec_sequential}; +use p3::{ + field::FieldAlgebra, + babybear::BabyBear, + matrix::{Matrix, dense::RowMajorMatrix}, + maybe_rayon::prelude::*, +}; +use rand::{Rng, thread_rng}; +use transcript::{BasicTranscript, Transcript}; +use witness::{InstancePaddingStrategy, RowMajorMatrix as WitnessRowMajorMatrix}; + +type E = BabyBearExt4; +type F = BabyBear; +type Pcs = Basefold; + +const NUM_SAMPLES: usize = 10; +const NUM_MATRICES: usize = 35; +const NUM_COLS: usize = 32; + +fn make_rmm(num_rows: usize, num_cols: usize) -> RowMajorMatrix { + let values: Vec = (0..num_rows * num_cols) + .into_par_iter() + .map(|i| F::from_canonical_u32(((i as u64 * 13 + 7) % (1 << 30)) as u32)) + .collect(); + RowMajorMatrix::new(values, num_cols) +} + +fn sample_heights(rng: &mut impl Rng, num_matrices: usize) -> Vec { + (0..num_matrices) + .map(|_| { + let log = rng.gen_range(16u32..=22); + let base = 1usize << log; + let lo = base - base / 4; + let hi = base + base / 4; + rng.gen_range(lo..=hi) + }) + .collect() +} + +fn eval_all_columns_at_point(rmm: &RowMajorMatrix, point: &[E]) -> Vec { + let w = rmm.width(); + let eq = build_eq_x_r_vec_sequential(point); + let mut col_evals = vec![E::ZERO; w]; + for (eq_r, row) in eq.iter().zip(rmm.rows()) { + for (col_eval, val) in col_evals.iter_mut().zip(row) { + *col_eval += *eq_r * val; + } + } + col_evals +} + +fn bench_comparison(c: &mut Criterion) { + let mut group = c.benchmark_group("jagged_vs_direct"); + group.sample_size(NUM_SAMPLES); + + let mut rng = thread_rng(); + let heights = sample_heights(&mut rng, NUM_MATRICES); + let log_heights: Vec = heights.iter().map(|h| ceil_log2(*h)).collect(); + let max_s = *log_heights.iter().max().unwrap(); + + println!("Matrix heights: {:?}", heights); + + let rmms: Vec<_> = heights.iter().map(|&h| make_rmm(h, NUM_COLS)).collect(); + let total_evals: usize = rmms.iter().map(|rmm| rmm.height() * rmm.width()).sum(); + let num_giga_vars = ceil_log2(total_evals); + + println!( + "num_matrices={NUM_MATRICES}, num_cols={NUM_COLS}, \ + total_evals={total_evals}, num_giga_vars={num_giga_vars}, max_s={max_s}" + ); + + let point: Vec = (0..max_s).map(|_| E::random(&mut rng)).collect(); + + let evals: Vec = rmms + .iter() + .zip(log_heights.iter()) + .flat_map(|(rmm, &s_i)| eval_all_columns_at_point(rmm, &point[(max_s - s_i)..])) + .collect(); + + // Per-matrix points and evals (used by direct batch_open). + let per_matrix_point_evals: Vec<(Vec, Vec)> = log_heights + .iter() + .enumerate() + .map(|(i, &s_i)| { + let matrix_point = point[(max_s - s_i)..].to_vec(); + let matrix_evals = evals[i * NUM_COLS..(i + 1) * NUM_COLS].to_vec(); + (matrix_point, matrix_evals) + }) + .collect(); + + // ======================== Jagged PCS ======================== + // BabyBear two-adicity is 27; RS rate_log=1 needs level+1 ≤ 27, so max poly_size is 2^25. + let reshape_log_height = num_giga_vars.saturating_sub(4).min(25); + let jagged_poly_size = 1usize << reshape_log_height; + let jagged_param = Pcs::setup(jagged_poly_size, SecurityLevel::Conjecture100bits).unwrap(); + let (jagged_pp, jagged_vp) = Pcs::trim(jagged_param, jagged_poly_size).unwrap(); + + group.bench_function("jagged/commit", |b| { + b.iter_custom(|iters| { + let mut time = Duration::new(0, 0); + for _ in 0..iters { + let rmms_clone = rmms.clone(); + let instant = std::time::Instant::now(); + let _ = + jagged_commit::(&jagged_pp, rmms_clone, reshape_log_height).unwrap(); + time += instant.elapsed(); + } + time + }) + }); + + let t0 = std::time::Instant::now(); + let jagged_comm = + jagged_commit::(&jagged_pp, rmms.clone(), reshape_log_height).unwrap(); + println!("jagged_commit: {:?}", t0.elapsed()); + let jagged_pure_comm = jagged_comm.to_commitment(); + + group.bench_function("jagged/batch_open", |b| { + b.iter_batched( + || { + let mut t = BasicTranscript::::new(b"bench"); + Pcs::write_commitment(&jagged_pure_comm.inner, &mut t).unwrap(); + t + }, + |mut t| { + jagged_batch_open::(&jagged_pp, &jagged_comm, &point, &evals, &mut t) + .unwrap(); + }, + BatchSize::SmallInput, + ); + }); + + let jagged_proof = { + let mut t = BasicTranscript::::new(b"bench"); + Pcs::write_commitment(&jagged_pure_comm.inner, &mut t).unwrap(); + let t0 = std::time::Instant::now(); + let proof = + jagged_batch_open::(&jagged_pp, &jagged_comm, &point, &evals, &mut t).unwrap(); + println!("jagged_batch_open: {:?}", t0.elapsed()); + proof + }; + let jagged_proof_size = bincode::serialize(&jagged_proof) + .map(|v| v.len()) + .unwrap_or(0); + + { + let mut t = BasicTranscript::::new(b"bench"); + Pcs::write_commitment(&jagged_pure_comm.inner, &mut t).unwrap(); + let t0 = std::time::Instant::now(); + jagged_batch_verify::( + &jagged_vp, + &jagged_pure_comm, + &point, + &evals, + &jagged_proof, + &mut t, + ) + .unwrap(); + println!("jagged_batch_verify: {:?}", t0.elapsed()); + } + + group.bench_function("jagged/batch_verify", |b| { + b.iter_batched( + || { + let mut t = BasicTranscript::::new(b"bench"); + Pcs::write_commitment(&jagged_pure_comm.inner, &mut t).unwrap(); + t + }, + |mut t| { + jagged_batch_verify::( + &jagged_vp, + &jagged_pure_comm, + &point, + &evals, + &jagged_proof, + &mut t, + ) + .unwrap(); + }, + BatchSize::SmallInput, + ); + }); + + // ======================== Direct Inner PCS ======================== + // Pcs::batch_commit expects witness::RowMajorMatrix, so convert. + let to_witness = |rmms: &[RowMajorMatrix]| -> Vec> { + rmms.iter() + .map(|rmm| { + WitnessRowMajorMatrix::new_by_values( + rmm.values.clone(), + rmm.width(), + InstancePaddingStrategy::Default, + ) + }) + .collect() + }; + + let direct_poly_size = 1usize << max_s; + let direct_param = Pcs::setup(direct_poly_size, SecurityLevel::Conjecture100bits).unwrap(); + let (direct_pp, direct_vp) = Pcs::trim(direct_param, direct_poly_size).unwrap(); + + group.bench_function("direct/commit", |b| { + b.iter_custom(|iters| { + let mut time = Duration::new(0, 0); + for _ in 0..iters { + let w_rmms = to_witness(&rmms); + let instant = std::time::Instant::now(); + let _ = Pcs::batch_commit(&direct_pp, w_rmms).unwrap(); + time += instant.elapsed(); + } + time + }) + }); + + let t0 = std::time::Instant::now(); + let direct_comm = Pcs::batch_commit(&direct_pp, to_witness(&rmms)).unwrap(); + println!("direct_commit: {:?}", t0.elapsed()); + let direct_pure_comm = Pcs::get_pure_commitment(&direct_comm); + + let make_direct_transcript = |comm: &>::Commitment| { + let mut t = BasicTranscript::::new(b"bench"); + Pcs::write_commitment(comm, &mut t).unwrap(); + for (_, matrix_evals) in &per_matrix_point_evals { + t.append_field_element_exts(matrix_evals); + } + t + }; + + group.bench_function("direct/batch_open", |b| { + b.iter_batched( + || make_direct_transcript(&direct_pure_comm), + |mut t| { + Pcs::batch_open( + &direct_pp, + vec![(&direct_comm, per_matrix_point_evals.clone())], + &mut t, + ) + .unwrap(); + }, + BatchSize::SmallInput, + ); + }); + + let direct_proof = { + let mut t = make_direct_transcript(&direct_pure_comm); + let t0 = std::time::Instant::now(); + let proof = Pcs::batch_open( + &direct_pp, + vec![(&direct_comm, per_matrix_point_evals.clone())], + &mut t, + ) + .unwrap(); + println!("direct_batch_open: {:?}", t0.elapsed()); + proof + }; + let direct_proof_size = bincode::serialize(&direct_proof) + .map(|v| v.len()) + .unwrap_or(0); + + let direct_verify_rounds: Vec<_> = log_heights + .iter() + .enumerate() + .map(|(i, &s_i)| { + let matrix_point = point[(max_s - s_i)..].to_vec(); + let matrix_evals = evals[i * NUM_COLS..(i + 1) * NUM_COLS].to_vec(); + (s_i, (matrix_point, matrix_evals)) + }) + .collect(); + + { + let mut t = make_direct_transcript(&direct_pure_comm); + let t0 = std::time::Instant::now(); + Pcs::batch_verify( + &direct_vp, + vec![(direct_pure_comm.clone(), direct_verify_rounds.clone())], + &direct_proof, + &mut t, + ) + .unwrap(); + println!("direct_batch_verify: {:?}", t0.elapsed()); + } + + group.bench_function("direct/batch_verify", |b| { + b.iter_batched( + || make_direct_transcript(&direct_pure_comm), + |mut t| { + Pcs::batch_verify( + &direct_vp, + vec![(direct_pure_comm.clone(), direct_verify_rounds.clone())], + &direct_proof, + &mut t, + ) + .unwrap(); + }, + BatchSize::SmallInput, + ); + }); + + group.finish(); + + println!("\n=== Proof Size Comparison ==="); + println!( + "Jagged PCS: {jagged_proof_size:>10} bytes ({:.1} KB)", + jagged_proof_size as f64 / 1024.0 + ); + println!( + "Direct PCS: {direct_proof_size:>10} bytes ({:.1} KB)", + direct_proof_size as f64 / 1024.0 + ); + println!( + "Ratio (direct / jagged): {:.2}x", + direct_proof_size as f64 / jagged_proof_size as f64 + ); +} + +criterion_group! { + name = benches; + config = Criterion::default().warm_up_time(Duration::from_millis(3000)); + targets = bench_comparison, +} +criterion_main!(benches); diff --git a/crates/mpcs/benches/jagged_pcs.rs b/crates/mpcs/benches/jagged_pcs.rs index 4b9650f..83e7838 100644 --- a/crates/mpcs/benches/jagged_pcs.rs +++ b/crates/mpcs/benches/jagged_pcs.rs @@ -1,19 +1,23 @@ use std::time::Duration; use criterion::*; -use ff_ext::{FromUniformBytes, GoldilocksExt2}; +use ff_ext::{BabyBearExt4, FromUniformBytes}; use mpcs::{ Basefold, BasefoldRSParams, PolynomialCommitmentScheme, SecurityLevel, jagged_batch_open, jagged_commit, }; -use multilinear_extensions::{mle::MultilinearExtension, util::ceil_log2}; -use p3::{field::FieldAlgebra, goldilocks::Goldilocks, matrix::Matrix}; +use multilinear_extensions::{util::ceil_log2, virtual_poly::build_eq_x_r_vec_sequential}; +use p3::{ + babybear::BabyBear, + field::FieldAlgebra, + matrix::{Matrix, dense::RowMajorMatrix}, + maybe_rayon::prelude::*, +}; use rand::{Rng, thread_rng}; use transcript::BasicTranscript; -use witness::{InstancePaddingStrategy, RowMajorMatrix}; -type E = GoldilocksExt2; -type F = Goldilocks; +type E = BabyBearExt4; +type F = BabyBear; type Pcs = Basefold; const NUM_SAMPLES: usize = 10; @@ -22,23 +26,35 @@ const NUM_COLS: usize = 32; fn make_rmm(num_rows: usize, num_cols: usize) -> RowMajorMatrix { let values: Vec = (0..num_rows * num_cols) + .into_par_iter() .map(|i| F::from_canonical_u32(((i as u64 * 13 + 7) % (1 << 30)) as u32)) .collect(); - RowMajorMatrix::::new_by_values(values, num_cols, InstancePaddingStrategy::Default) + RowMajorMatrix::new(values, num_cols) } fn sample_heights(rng: &mut impl Rng, num_matrices: usize) -> Vec { - let log_range: Vec = (16..=22).collect(); (0..num_matrices) - .map(|_| 1usize << log_range[rng.gen_range(0..log_range.len())]) + .map(|_| { + let log = rng.gen_range(16u32..=22); + let base = 1usize << log; + // Jitter within ±25% of the power-of-two base. + let lo = base - base / 4; + let hi = base + base / 4; + rng.gen_range(lo..=hi) + }) .collect() } -fn eval_column_poly_at_point(col_evals: &[F], point: &[E]) -> E { - let s = point.len(); - assert_eq!(col_evals.len(), 1 << s); - let mle = MultilinearExtension::from_evaluations_vec(s, col_evals.to_vec()); - mle.evaluate(point) +fn eval_all_columns_at_point(rmm: &RowMajorMatrix, point: &[E]) -> Vec { + let w = rmm.width(); + let eq = build_eq_x_r_vec_sequential(point); + let mut col_evals = vec![E::ZERO; w]; + for (eq_r, row) in eq.iter().zip(rmm.rows()) { + for (col_eval, val) in col_evals.iter_mut().zip(row) { + *col_eval += *eq_r * val; + } + } + col_evals } fn bench_jagged_pcs(c: &mut Criterion) { @@ -48,10 +64,7 @@ fn bench_jagged_pcs(c: &mut Criterion) { let mut rng = thread_rng(); let heights = sample_heights(&mut rng, NUM_MATRICES); - println!( - "Matrix heights (log2): {:?}", - heights.iter().map(|h| ceil_log2(*h)).collect::>() - ); + println!("Matrix heights: {:?}", heights); let rmms: Vec<_> = heights.iter().map(|&h| make_rmm(h, NUM_COLS)).collect(); @@ -65,34 +78,23 @@ fn bench_jagged_pcs(c: &mut Criterion) { NUM_MATRICES, NUM_COLS, total_evals, num_giga_vars, max_s ); - // Prepare evaluation data (shared across all reshape_log_height values). - let col_polys: Vec> = rmms - .iter() - .flat_map(|rmm| { - let h = rmm.height(); - let w = rmm.width(); - (0..w).map(move |c| (0..h).map(|r| rmm.values[r * w + c]).collect()) - }) - .collect(); - let point: Vec = (0..max_s).map(|_| E::random(&mut rng)).collect(); - let evals: Vec = col_polys + let evals: Vec = rmms .iter() - .zip( - log_heights - .iter() - .flat_map(|&s| std::iter::repeat_n(s, NUM_COLS)), - ) - .map(|(col, s_i)| eval_column_poly_at_point(col, &point[(max_s - s_i)..])) + .zip(log_heights.iter()) + .flat_map(|(rmm, &s_i)| eval_all_columns_at_point(rmm, &point[(max_s - s_i)..])) .collect(); - // Sweep reshape_log_height values: baseline (no reshape) + a few smaller heights. + // BabyBear two-adicity is 27; RS rate_log=1 needs level+1 ≤ 27, so max poly_size is 2^25. + let max_log_h = 25usize; let reshape_log_heights: Vec = { - let mut vals = vec![num_giga_vars]; + let start = num_giga_vars.min(max_log_h); + let mut vals = vec![start]; for step in [2, 4, 6] { - if num_giga_vars > step { - vals.push(num_giga_vars - step); + let v = start.saturating_sub(step); + if v > 0 && !vals.contains(&v) { + vals.push(v); } } vals @@ -141,8 +143,6 @@ fn bench_jagged_pcs(c: &mut Criterion) { Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); let proof = jagged_batch_open::(&pp, &comm, &point, &evals, &mut transcript_p).unwrap(); - let pure_comm = comm.to_commitment(); - let proof_size = bincode::serialize(&proof).map(|v| v.len()).unwrap_or(0); println!("{label}: proof_size={proof_size} bytes, col_evals.len={w}"); @@ -150,13 +150,13 @@ fn bench_jagged_pcs(c: &mut Criterion) { b.iter_batched( || { let mut transcript = BasicTranscript::::new(b"jagged_bench"); - Pcs::write_commitment(&pure_comm.inner, &mut transcript).unwrap(); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript).unwrap(); transcript }, |mut transcript| { mpcs::jagged_batch_verify::( &vp, - &pure_comm, + &comm.to_commitment(), &point, &evals, &proof, diff --git a/crates/mpcs/src/jagged/mod.rs b/crates/mpcs/src/jagged/mod.rs index 1206a7a..7eb19a2 100644 --- a/crates/mpcs/src/jagged/mod.rs +++ b/crates/mpcs/src/jagged/mod.rs @@ -162,12 +162,12 @@ use multilinear_extensions::{ }; use p3::{ field::FieldAlgebra, - matrix::{Matrix, bitrev::BitReversableMatrix}, + matrix::{Matrix, bitrev::BitReversableMatrix, dense::RowMajorMatrix}, maybe_rayon::prelude::*, }; use transcript::Transcript; use types::int_to_field_bits; -use witness::{InstancePaddingStrategy, RowMajorMatrix}; +use witness::{InstancePaddingStrategy, RowMajorMatrix as WitnessRowMajorMatrix}; /// Commit to a sequence of row-major matrices using the Jagged PCS scheme. /// @@ -198,6 +198,7 @@ pub fn jagged_commit> } // --- Step 1: Compute cumulative heights --- + // Each column polynomial is padded to the next power of two. let mut poly_heights: Vec = Vec::new(); for rmm in &rmms { let num_rows = rmm.height(); @@ -213,8 +214,9 @@ pub fn jagged_commit> "jagged_commit: matrix has zero columns".to_string(), )); } + let padded_height = num_rows.next_power_of_two(); for _ in 0..num_cols { - poly_heights.push(num_rows); + poly_heights.push(padded_height); } } if poly_heights.is_empty() { @@ -246,10 +248,17 @@ pub fn jagged_commit> // `poly_idx` tracks which poly (column index in cumulative_heights) is the // first polynomial of the current matrix. let mut poly_idx = 0; - for rmm in &rmms { + for rmm in rmms { + // Pad to next power-of-two height for bit-reversal. + let padded_height = rmm.height().next_power_of_two(); + let mut padded = rmm; + if padded.height() < padded_height { + padded.pad_to_height(padded_height, E::BaseField::ZERO); + } + // Step 2: Bit-reverse the rows (suffix-to-prefix transformation). // br.values[i * n_cols + j] = original[bitrev(i)][j] - let br = rmm.as_view().bit_reverse_rows().to_row_major_matrix(); + let br = padded.as_view().bit_reverse_rows().to_row_major_matrix(); let n_cols = br.width(); let n_rows = br.height(); @@ -313,7 +322,7 @@ pub fn jagged_commit> row_major }; - let giga_rmm = RowMajorMatrix::::new_by_values( + let giga_rmm = WitnessRowMajorMatrix::::new_by_values( giga_data, w, InstancePaddingStrategy::Default, @@ -669,7 +678,7 @@ mod tests { let values: Vec = (0..num_rows * num_cols) .map(|i| F::from_canonical_u64(i as u64 + 1)) .collect(); - RowMajorMatrix::::new_by_values(values, num_cols, InstancePaddingStrategy::Default) + RowMajorMatrix::new(values, num_cols) } #[test] From 85dbfef466d02f0c9574387c73a2ab52d5d67c50 Mon Sep 17 00:00:00 2001 From: Ming Date: Tue, 12 May 2026 10:59:45 +0800 Subject: [PATCH 20/21] jagged: direct concat for prefix-aligned points (#56) misc: prefix align without bit reverse --- crates/mpcs/benches/comparison.rs | 8 +- crates/mpcs/benches/jagged_pcs.rs | 2 +- crates/mpcs/src/jagged/mod.rs | 188 ++++++++++++++++------------- crates/mpcs/src/jagged/sumcheck.rs | 2 +- 4 files changed, 108 insertions(+), 92 deletions(-) diff --git a/crates/mpcs/benches/comparison.rs b/crates/mpcs/benches/comparison.rs index c5ee62d..4a81678 100644 --- a/crates/mpcs/benches/comparison.rs +++ b/crates/mpcs/benches/comparison.rs @@ -8,8 +8,8 @@ use mpcs::{ }; use multilinear_extensions::{util::ceil_log2, virtual_poly::build_eq_x_r_vec_sequential}; use p3::{ - field::FieldAlgebra, babybear::BabyBear, + field::FieldAlgebra, matrix::{Matrix, dense::RowMajorMatrix}, maybe_rayon::prelude::*, }; @@ -82,7 +82,7 @@ fn bench_comparison(c: &mut Criterion) { let evals: Vec = rmms .iter() .zip(log_heights.iter()) - .flat_map(|(rmm, &s_i)| eval_all_columns_at_point(rmm, &point[(max_s - s_i)..])) + .flat_map(|(rmm, &s_i)| eval_all_columns_at_point(rmm, &point[..s_i])) .collect(); // Per-matrix points and evals (used by direct batch_open). @@ -90,7 +90,7 @@ fn bench_comparison(c: &mut Criterion) { .iter() .enumerate() .map(|(i, &s_i)| { - let matrix_point = point[(max_s - s_i)..].to_vec(); + let matrix_point = point[..s_i].to_vec(); let matrix_evals = evals[i * NUM_COLS..(i + 1) * NUM_COLS].to_vec(); (matrix_point, matrix_evals) }) @@ -269,7 +269,7 @@ fn bench_comparison(c: &mut Criterion) { .iter() .enumerate() .map(|(i, &s_i)| { - let matrix_point = point[(max_s - s_i)..].to_vec(); + let matrix_point = point[..s_i].to_vec(); let matrix_evals = evals[i * NUM_COLS..(i + 1) * NUM_COLS].to_vec(); (s_i, (matrix_point, matrix_evals)) }) diff --git a/crates/mpcs/benches/jagged_pcs.rs b/crates/mpcs/benches/jagged_pcs.rs index 83e7838..99dd95b 100644 --- a/crates/mpcs/benches/jagged_pcs.rs +++ b/crates/mpcs/benches/jagged_pcs.rs @@ -83,7 +83,7 @@ fn bench_jagged_pcs(c: &mut Criterion) { let evals: Vec = rmms .iter() .zip(log_heights.iter()) - .flat_map(|(rmm, &s_i)| eval_all_columns_at_point(rmm, &point[(max_s - s_i)..])) + .flat_map(|(rmm, &s_i)| eval_all_columns_at_point(rmm, &point[..s_i])) .collect(); // BabyBear two-adicity is 27; RS rate_log=1 needs level+1 ≤ 27, so max poly_size is 2^25. diff --git a/crates/mpcs/src/jagged/mod.rs b/crates/mpcs/src/jagged/mod.rs index 7eb19a2..c8f66df 100644 --- a/crates/mpcs/src/jagged/mod.rs +++ b/crates/mpcs/src/jagged/mod.rs @@ -11,50 +11,23 @@ //! chips into a single "giga" multilinear polynomial `q'`: //! //! ```text -//! q' = bitrev(p_0) || bitrev(p_1) || ... || bitrev(p_{N-1}) +//! q' = p_0 || p_1 || ... || p_{N-1} //! ``` //! //! where each `p_i` is a column polynomial extracted from the input trace matrices. -//! We cannot use the naive `q = p_0 || ... || p_{N-1}` because the jagged sumcheck -//! requires prefix-aligned evaluation points, but Ceno's main sumcheck evaluates -//! polynomials at a **suffix** of the challenge point. The `bitrev` (suffix-to-prefix -//! bit-reversal) bridges the two, making the jagged PCS applicable (see below). -//! -//! ## Suffix-to-Prefix Transformation -//! -//! Each `p_i` has `s_i` variables and let `m` be the length of `z_r`. The main sumcheck prover -//! outputs evaluations `v_i = p_i(z_r[(m-s_i)..m])` — i.e., at the suffix of the -//! random challenge point. -//! To make these evaluations compatible with the jagged sumcheck (which operates on -//! prefix-aligned polynomials), we apply a bit-reversal permutation to each polynomial's -//! evaluations: -//! -//! ```text -//! p_i'[j] = p_i[bitrev_s_i(j)] (for j in 0..2^s_i) -//! ``` -//! -//! After bit-reversal, -//! ```text -//! v_i = p_i(z_r[(m-s_i)..m]) -//! = p_i'(z_r'[..s_i]) -//! ``` -//! where `z_r' = reverse(z_r)`. +//! The jagged sumcheck uses prefix-aligned evaluation points, so a polynomial with +//! `s_i = ceil_log2(h_i)` variables is opened at `z_r[..s_i]`. //! //! ## Cumulative Heights //! //! Each polynomial `p_i` has `s_i = ceil_log2(h_i)` variables, where `h_i` is the -//! original (unpadded) number of evaluations. Before bit-reversal, `p_i` is -//! zero-padded to `2^{s_i}` entries. **After bit-reversal, these zeros are scattered -//! throughout the block** (not contiguous at the end), so each polynomial occupies -//! a full `2^{s_i}`-entry block in `q'`. -//! -//! **Note:** Unlike SP1's jagged PCS (which uses raw `h_i`), our `q'` contains -//! `Σ (2^{s_i} - h_i)` extra zeros from this padding — the cost of the -//! suffix-to-prefix bit-reversal required by Ceno's main sumcheck. +//! real number of evaluations from the input matrix column. `q'` stores exactly +//! those `h_i` evaluations; any implicit zero padding to `2^{s_i}` is only an MLE +//! evaluation convention and is not materialized inside the concatenation. //! //! The cumulative height sequence `t` tracks the starting position of each polynomial in `q'`: //! - `t[0] = 0` -//! - `t[i+1] = t[i] + 2^{s_i}` (the padded height, not the original `h_i`) +//! - `t[i+1] = t[i] + h_i` //! //! Given a position `b` in `q'`, the inverse mapping `inv(b) = (i, r)` is defined by: //! - `t[i] <= b < t[i+1]` @@ -65,21 +38,16 @@ //! //! ## Commit Protocol //! -//! 1. For each input matrix `M_k` (with `h_k` rows and `w_k` columns): -//! a. Zero-pad `M_k` to `2^{s_k}` rows (where `s_k = ceil_log2(h_k)`). -//! b. Apply bit-reversal to the rows. Note: after bit-reversal the zero-padding -//! is scattered throughout the block, not contiguous at the end. -//! c. Extract each column as a polynomial with `2^{s_k}` evaluations. -//! 2. Concatenate all bit-reversed polynomials: `cat = bitrev(p_0) || bitrev(p_1) || ...` +//! 1. For each input matrix `M_k` (with `h_k` rows and `w_k` columns), extract each +//! column as a polynomial with `h_k` evaluations. +//! 2. Concatenate all column polynomials: `cat = p_0 || p_1 || ...` //! 3. Compute cumulative heights `t[i]`. //! 4. Pad `cat` to the next power of two (required for MLE representation). //! 5. Commit to the padded `cat` as a single-column matrix using the inner PCS. //! //! ## Batch Open Protocol //! -//! For notational simplicity, we write `p_i` and `z_r` instead of `p_i'` and `z_r'` -//! (the bit-reversed variants) throughout this section. Let `s_i` denote the number -//! of variables of `p_i` and let `m = max(s_i)`. +//! Let `s_i` denote the number of variables of `p_i` and let `m = max(s_i)`. //! //! Each column opening `v_i = p_i(z_r[..s_i])` requires its own sumcheck. We batch //! all K openings into one using `eq(z_c, ·)` weights (soundness loss: `log2(N) / |E|`). @@ -162,7 +130,7 @@ use multilinear_extensions::{ }; use p3::{ field::FieldAlgebra, - matrix::{Matrix, bitrev::BitReversableMatrix, dense::RowMajorMatrix}, + matrix::{Matrix, dense::RowMajorMatrix}, maybe_rayon::prelude::*, }; use transcript::Transcript; @@ -172,12 +140,11 @@ use witness::{InstancePaddingStrategy, RowMajorMatrix as WitnessRowMajorMatrix}; /// Commit to a sequence of row-major matrices using the Jagged PCS scheme. /// /// This function implements the commit phase described in Ceno issue #1288: -/// 1. For each matrix, bit-reverse its rows (suffix-to-prefix transformation). -/// 2. Transpose the bit-reversed matrix (row-major → column-major), so each +/// 1. For each matrix, transpose it (row-major → column-major), so each /// column polynomial occupies a contiguous region in memory. -/// 3. Concatenate all column polynomials: `q' = col_0 || col_1 || ...` -/// 4. Compute the cumulative height sequence `t`. -/// 5. Commit to `q'` as a single-column matrix using `InnerPcs::batch_commit`. +/// 2. Concatenate all column polynomials: `q' = col_0 || col_1 || ...` +/// 3. Compute the cumulative height sequence `t`. +/// 4. Commit to `q'` as a single-column matrix using `InnerPcs::batch_commit`. /// /// # Arguments /// * `pp` — Prover parameters for `InnerPcs`. @@ -197,8 +164,7 @@ pub fn jagged_commit> )); } - // --- Step 1: Compute cumulative heights --- - // Each column polynomial is padded to the next power of two. + // --- Step 1: Compute cumulative heights from real matrix heights --- let mut poly_heights: Vec = Vec::new(); for rmm in &rmms { let num_rows = rmm.height(); @@ -214,9 +180,8 @@ pub fn jagged_commit> "jagged_commit: matrix has zero columns".to_string(), )); } - let padded_height = num_rows.next_power_of_two(); for _ in 0..num_cols { - poly_heights.push(padded_height); + poly_heights.push(num_rows); } } if poly_heights.is_empty() { @@ -235,7 +200,7 @@ pub fn jagged_commit> }) .collect::>(); - // --- Steps 2 & 3: Bit-reverse rows, transpose, and write to concatenated --- + // --- Steps 2 & 3: Transpose and write to concatenated --- let total_size = cumulative_heights.last().copied().unwrap(); let mut concatenated: Vec = Vec::with_capacity(total_size); // Safety: every element in `concatenated[0..total_size]` is fully written @@ -249,31 +214,20 @@ pub fn jagged_commit> // first polynomial of the current matrix. let mut poly_idx = 0; for rmm in rmms { - // Pad to next power-of-two height for bit-reversal. - let padded_height = rmm.height().next_power_of_two(); - let mut padded = rmm; - if padded.height() < padded_height { - padded.pad_to_height(padded_height, E::BaseField::ZERO); - } - - // Step 2: Bit-reverse the rows (suffix-to-prefix transformation). - // br.values[i * n_cols + j] = original[bitrev(i)][j] - let br = padded.as_view().bit_reverse_rows().to_row_major_matrix(); - - let n_cols = br.width(); - let n_rows = br.height(); + let n_cols = rmm.width(); + let n_rows = rmm.height(); let n_cells = n_cols * n_rows; // The start position in `concatenated` for this matrix's block of polynomials. let start = cumulative_heights[poly_idx]; - // Step 3: Transpose — write each column j of `br` (= one polynomial) + // Step 3: Transpose — write each column j of `rmm` (= one polynomial) // into its corresponding contiguous slice in `concatenated`. (0..n_cols) .into_par_iter() .zip(concatenated[start..start + n_cells].par_chunks_mut(n_rows)) .for_each(|(j, chunk)| { - br.values + rmm.values .iter() .skip(j) .step_by(n_cols) @@ -371,7 +325,7 @@ fn compute_f_at_point( /// /// Polynomials may have different heights. `point` is the evaluation point of /// length `max_s = max(log2(h_i))`. Polynomial `i` with `s_i` variables is -/// evaluated at the suffix `point[(max_s - s_i)..]`. +/// evaluated at the prefix `point[..s_i]`. /// /// The protocol: /// 1. Batch the K column claims via a random column challenge `z_col`. @@ -414,8 +368,7 @@ pub fn jagged_batch_open = point.iter().rev().cloned().collect(); + let z_row: Vec = point.to_vec(); // Write evals to transcript, then sample z_col. transcript.append_field_element_exts(evals); @@ -525,7 +478,7 @@ pub fn jagged_batch_open>( vp: &InnerPcs::VerifierParam, comm: &JaggedCommitment, @@ -551,8 +504,7 @@ pub fn jagged_batch_verify = point.iter().rev().cloned().collect(); + let z_row: Vec = point.to_vec(); // Replay transcript: write evals, sample z_col. transcript.append_field_element_exts(evals); @@ -563,15 +515,15 @@ pub fn jagged_batch_verify(reshape_log_height); + let m1 = make_rmm(3, 1); + let m2 = make_rmm(5, 2); + + let comm = jagged_commit::(&pp, vec![m1, m2], reshape_log_height) + .expect("commit should succeed"); + + assert_eq!(comm.num_polys(), 3); + assert_eq!(comm.poly_heights, vec![3, 5, 5]); + assert_eq!(comm.cumulative_heights, vec![0, 3, 8, 13]); + assert_eq!(comm.total_evaluations(), 13); + } + #[test] fn test_jagged_commit_smoke() { // Two matrices: 4x1 and 4x2 → 3 polynomials, total 12 evals, padded to 16 @@ -767,12 +736,13 @@ mod tests { use rand::thread_rng; use transcript::basic::BasicTranscript; - /// Evaluate a single column polynomial at `point` (the original, non-bit-reversed poly). - /// `col_evals` are the raw evaluations of the column (before bit-reversal). + /// Evaluate a single column polynomial at `point`, zero-padding to `2^point.len()`. fn eval_column_poly_at_point(col_evals: &[F], point: &[E]) -> E { let s = point.len(); - assert_eq!(col_evals.len(), 1 << s); - let mle = MultilinearExtension::from_evaluations_vec(s, col_evals.to_vec()); + assert!(col_evals.len() <= 1 << s); + let mut padded = col_evals.to_vec(); + padded.resize(1 << s, F::ZERO); + let mle = MultilinearExtension::from_evaluations_vec(s, padded); mle.evaluate(point) } @@ -790,7 +760,7 @@ mod tests { let (pp, vp) = setup_pcs::(num_giga_vars); let rmms: Vec<_> = heights.iter().map(|&h| make_rmm(h, 1)).collect(); - // Extract each column polynomial (before bit-reversal). + // Extract each column polynomial. let col_polys: Vec> = rmms .iter() .map(|rmm| (0..rmm.height()).map(|r| rmm.values[r]).collect()) @@ -803,13 +773,13 @@ mod tests { Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); // Random evaluation point of length max_s. - // Poly i with s_i variables is evaluated at the suffix point[(max_s - s_i)..]. + // Poly i with s_i variables is evaluated at the prefix point[..s_i]. let point: Vec = (0..max_s).map(|_| E::random(&mut rng)).collect(); let evals: Vec = col_polys .iter() .zip(log_heights.iter()) - .map(|(col, &s_i)| eval_column_poly_at_point(col, &point[(max_s - s_i)..])) + .map(|(col, &s_i)| eval_column_poly_at_point(col, &point[..s_i])) .collect(); // Prover: batch open. @@ -869,6 +839,52 @@ mod tests { .expect("batch verify"); } + #[test] + fn test_jagged_batch_open_verify_non_power_of_two_heights() { + let mut rng = thread_rng(); + + let heights = [1023usize, 777, 513]; + let log_heights: Vec = heights.iter().map(|&h| ceil_log2(h)).collect(); + let max_s = *log_heights.iter().max().unwrap(); + let total_evals: usize = heights.iter().sum(); + let num_giga_vars = ceil_log2(total_evals); + + let (pp, vp) = setup_pcs::(num_giga_vars); + let rmms: Vec<_> = heights.iter().map(|&h| make_rmm(h, 1)).collect(); + + let col_polys: Vec> = rmms + .iter() + .map(|rmm| (0..rmm.height()).map(|r| rmm.values[r]).collect()) + .collect(); + + let mut transcript_p = BasicTranscript::::new(b"jagged_non_power_two"); + let comm = jagged_commit::(&pp, rmms, num_giga_vars).expect("commit"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); + + let point: Vec = (0..max_s).map(|_| E::random(&mut rng)).collect(); + let evals: Vec = col_polys + .iter() + .zip(log_heights.iter()) + .map(|(col, &s_i)| eval_column_poly_at_point(col, &point[..s_i])) + .collect(); + + let proof = jagged_batch_open::(&pp, &comm, &point, &evals, &mut transcript_p) + .expect("batch open"); + + let mut transcript_v = BasicTranscript::::new(b"jagged_non_power_two"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_v).unwrap(); + + jagged_batch_verify::( + &vp, + &comm.to_commitment(), + &point, + &evals, + &proof, + &mut transcript_v, + ) + .expect("batch verify"); + } + #[test] fn test_jagged_batch_open_verify_soundness() { let mut rng = thread_rng(); @@ -992,7 +1008,7 @@ mod tests { let evals: Vec = col_polys .iter() .zip(log_heights.iter()) - .map(|(col, &s_i)| eval_column_poly_at_point(col, &point[(max_s - s_i)..])) + .map(|(col, &s_i)| eval_column_poly_at_point(col, &point[..s_i])) .collect(); let proof = jagged_batch_open::(&pp, &comm, &point, &evals, &mut transcript_p) diff --git a/crates/mpcs/src/jagged/sumcheck.rs b/crates/mpcs/src/jagged/sumcheck.rs index 243dd99..8e0eddc 100644 --- a/crates/mpcs/src/jagged/sumcheck.rs +++ b/crates/mpcs/src/jagged/sumcheck.rs @@ -21,7 +21,7 @@ const LOG2_MAX_EPOCH: u32 = 3; /// All inputs needed for the jagged sumcheck. pub struct JaggedSumcheckInput<'a, E: ExtensionField> { - /// Giga polynomial evaluations (concatenated, bit-reversed). + /// Giga polynomial evaluations (directly concatenated). pub q_evals: &'a [E::BaseField], /// n = log2(padded_total_size). pub num_giga_vars: usize, From 399c29e56d3b77082390149b800695d370835bec Mon Sep 17 00:00:00 2001 From: Ming Date: Wed, 13 May 2026 13:01:11 +0800 Subject: [PATCH 21/21] Integrate Jagged PCS for Ceno CPU (#58) * misc: prefix align without bit reverse * WIP integrate jagged pcs cpu * Fix jagged PCS padded opening normalization * misc: clippy fix * avoid reconstruct q_mles * more docs * misc: fmt and clippy --- crates/mpcs/benches/comparison.rs | 31 +- crates/mpcs/benches/jagged_pcs.rs | 14 +- crates/mpcs/benches/jagged_sumcheck.rs | 6 +- crates/mpcs/src/basefold/structure.rs | 27 +- crates/mpcs/src/jagged/mod.rs | 424 +++++++++++++++++++++---- crates/mpcs/src/jagged/sumcheck.rs | 112 +++++-- crates/mpcs/src/jagged/types.rs | 9 + crates/mpcs/src/lib.rs | 11 +- 8 files changed, 505 insertions(+), 129 deletions(-) diff --git a/crates/mpcs/benches/comparison.rs b/crates/mpcs/benches/comparison.rs index 4a81678..ae5305b 100644 --- a/crates/mpcs/benches/comparison.rs +++ b/crates/mpcs/benches/comparison.rs @@ -7,12 +7,7 @@ use mpcs::{ jagged_batch_verify, jagged_commit, }; use multilinear_extensions::{util::ceil_log2, virtual_poly::build_eq_x_r_vec_sequential}; -use p3::{ - babybear::BabyBear, - field::FieldAlgebra, - matrix::{Matrix, dense::RowMajorMatrix}, - maybe_rayon::prelude::*, -}; +use p3::{babybear::BabyBear, field::FieldAlgebra, matrix::Matrix, maybe_rayon::prelude::*}; use rand::{Rng, thread_rng}; use transcript::{BasicTranscript, Transcript}; use witness::{InstancePaddingStrategy, RowMajorMatrix as WitnessRowMajorMatrix}; @@ -25,12 +20,12 @@ const NUM_SAMPLES: usize = 10; const NUM_MATRICES: usize = 35; const NUM_COLS: usize = 32; -fn make_rmm(num_rows: usize, num_cols: usize) -> RowMajorMatrix { +fn make_rmm(num_rows: usize, num_cols: usize) -> WitnessRowMajorMatrix { let values: Vec = (0..num_rows * num_cols) .into_par_iter() .map(|i| F::from_canonical_u32(((i as u64 * 13 + 7) % (1 << 30)) as u32)) .collect(); - RowMajorMatrix::new(values, num_cols) + WitnessRowMajorMatrix::new_by_values(values, num_cols, InstancePaddingStrategy::Default) } fn sample_heights(rng: &mut impl Rng, num_matrices: usize) -> Vec { @@ -45,7 +40,7 @@ fn sample_heights(rng: &mut impl Rng, num_matrices: usize) -> Vec { .collect() } -fn eval_all_columns_at_point(rmm: &RowMajorMatrix, point: &[E]) -> Vec { +fn eval_all_columns_at_point(rmm: &WitnessRowMajorMatrix, point: &[E]) -> Vec { let w = rmm.width(); let eq = build_eq_x_r_vec_sequential(point); let mut col_evals = vec![E::ZERO; w]; @@ -190,19 +185,6 @@ fn bench_comparison(c: &mut Criterion) { }); // ======================== Direct Inner PCS ======================== - // Pcs::batch_commit expects witness::RowMajorMatrix, so convert. - let to_witness = |rmms: &[RowMajorMatrix]| -> Vec> { - rmms.iter() - .map(|rmm| { - WitnessRowMajorMatrix::new_by_values( - rmm.values.clone(), - rmm.width(), - InstancePaddingStrategy::Default, - ) - }) - .collect() - }; - let direct_poly_size = 1usize << max_s; let direct_param = Pcs::setup(direct_poly_size, SecurityLevel::Conjecture100bits).unwrap(); let (direct_pp, direct_vp) = Pcs::trim(direct_param, direct_poly_size).unwrap(); @@ -211,9 +193,8 @@ fn bench_comparison(c: &mut Criterion) { b.iter_custom(|iters| { let mut time = Duration::new(0, 0); for _ in 0..iters { - let w_rmms = to_witness(&rmms); let instant = std::time::Instant::now(); - let _ = Pcs::batch_commit(&direct_pp, w_rmms).unwrap(); + let _ = Pcs::batch_commit(&direct_pp, rmms.clone()).unwrap(); time += instant.elapsed(); } time @@ -221,7 +202,7 @@ fn bench_comparison(c: &mut Criterion) { }); let t0 = std::time::Instant::now(); - let direct_comm = Pcs::batch_commit(&direct_pp, to_witness(&rmms)).unwrap(); + let direct_comm = Pcs::batch_commit(&direct_pp, rmms.clone()).unwrap(); println!("direct_commit: {:?}", t0.elapsed()); let direct_pure_comm = Pcs::get_pure_commitment(&direct_comm); diff --git a/crates/mpcs/benches/jagged_pcs.rs b/crates/mpcs/benches/jagged_pcs.rs index 99dd95b..5f4fa5c 100644 --- a/crates/mpcs/benches/jagged_pcs.rs +++ b/crates/mpcs/benches/jagged_pcs.rs @@ -7,14 +7,10 @@ use mpcs::{ jagged_commit, }; use multilinear_extensions::{util::ceil_log2, virtual_poly::build_eq_x_r_vec_sequential}; -use p3::{ - babybear::BabyBear, - field::FieldAlgebra, - matrix::{Matrix, dense::RowMajorMatrix}, - maybe_rayon::prelude::*, -}; +use p3::{babybear::BabyBear, field::FieldAlgebra, matrix::Matrix, maybe_rayon::prelude::*}; use rand::{Rng, thread_rng}; use transcript::BasicTranscript; +use witness::{InstancePaddingStrategy, RowMajorMatrix as WitnessRowMajorMatrix}; type E = BabyBearExt4; type F = BabyBear; @@ -24,12 +20,12 @@ const NUM_SAMPLES: usize = 10; const NUM_MATRICES: usize = 30; const NUM_COLS: usize = 32; -fn make_rmm(num_rows: usize, num_cols: usize) -> RowMajorMatrix { +fn make_rmm(num_rows: usize, num_cols: usize) -> WitnessRowMajorMatrix { let values: Vec = (0..num_rows * num_cols) .into_par_iter() .map(|i| F::from_canonical_u32(((i as u64 * 13 + 7) % (1 << 30)) as u32)) .collect(); - RowMajorMatrix::new(values, num_cols) + WitnessRowMajorMatrix::new_by_values(values, num_cols, InstancePaddingStrategy::Default) } fn sample_heights(rng: &mut impl Rng, num_matrices: usize) -> Vec { @@ -45,7 +41,7 @@ fn sample_heights(rng: &mut impl Rng, num_matrices: usize) -> Vec { .collect() } -fn eval_all_columns_at_point(rmm: &RowMajorMatrix, point: &[E]) -> Vec { +fn eval_all_columns_at_point(rmm: &WitnessRowMajorMatrix, point: &[E]) -> Vec { let w = rmm.width(); let eq = build_eq_x_r_vec_sequential(point); let mut col_evals = vec![E::ZERO; w]; diff --git a/crates/mpcs/benches/jagged_sumcheck.rs b/crates/mpcs/benches/jagged_sumcheck.rs index a6efe4e..a9f6081 100644 --- a/crates/mpcs/benches/jagged_sumcheck.rs +++ b/crates/mpcs/benches/jagged_sumcheck.rs @@ -1,6 +1,8 @@ use criterion::*; use ff_ext::{BabyBearExt4, ExtensionField, FieldFrom, FromUniformBytes}; -use mpcs::{JaggedSumcheckInput, assist_sumcheck_prove, jagged_sumcheck_prove}; +use mpcs::jagged::{ + JaggedSumcheckInput, QPrimeEvaluations, assist_sumcheck_prove, jagged_sumcheck_prove, +}; use multilinear_extensions::{util::ceil_log2, virtual_poly::build_eq_x_r_vec}; use rand::thread_rng; use transcript::BasicTranscript; @@ -40,7 +42,7 @@ fn bench_jagged_sumcheck(c: &mut Criterion) { let z_col: Vec = (0..z_col_vars).map(|_| E::random(&mut rng)).collect(); let input = JaggedSumcheckInput { - q_evals: &q_evals, + q_evals: QPrimeEvaluations::Flat(&q_evals), num_giga_vars, cumulative_heights: &cumulative_heights, eq_row: build_eq_x_r_vec(&z_row), diff --git a/crates/mpcs/src/basefold/structure.rs b/crates/mpcs/src/basefold/structure.rs index f97c70b..0f94893 100644 --- a/crates/mpcs/src/basefold/structure.rs +++ b/crates/mpcs/src/basefold/structure.rs @@ -81,9 +81,34 @@ pub struct BasefoldVerifierParams> { pub(super) security_level: SecurityLevel, } -impl_pcs_fri_param!(BasefoldProverParams); impl_pcs_fri_param!(BasefoldVerifierParams); +// Do not use `impl_pcs_fri_param!(BasefoldProverParams)` here. +// +// The macro only provides the FRI proof-of-work security parameter and would +// leave `PCSFriParam::get_max_message_size_log` at its default `usize::MAX`. +// Prover params must expose the real encoding message-size bound because +// Jagged PCS uses it to choose the q' reshape height before committing through +// Basefold. Returning the default would let Jagged construct a q' shape that +// exceeds the inner Basefold encoding limit. +impl> PCSFriParam for BasefoldProverParams { + fn get_pow_bits_by_level(&self, pow_strategy: crate::PowStrategy) -> usize { + match ( + &self.security_level, + pow_strategy, + >::get_rate_log(), + >::get_number_queries(), + ) { + (SecurityLevel::Conjecture100bits, crate::PowStrategy::FriPow, 1, 100) => 16, + _ => unimplemented!(), + } + } + + fn get_max_message_size_log(&self) -> usize { + self.encoding_params.get_max_message_size_log() + } +} + /// A polynomial commitment together with all the data (e.g., the codeword, and Merkle tree) /// used to generate this commitment and for assistant in opening pub struct BasefoldCommitmentWithWitness diff --git a/crates/mpcs/src/jagged/mod.rs b/crates/mpcs/src/jagged/mod.rs index c8f66df..c7cee04 100644 --- a/crates/mpcs/src/jagged/mod.rs +++ b/crates/mpcs/src/jagged/mod.rs @@ -114,29 +114,45 @@ mod types; pub use assist::{assist_sumcheck_prove, compute_q_at_assist_point}; pub use evaluator::{evaluate_g, evaluate_g_backward, evaluate_g_forward}; -pub use sumcheck::{JaggedSumcheckInput, jagged_sumcheck_prove}; -pub use types::{JaggedBatchOpenProof, JaggedCommitment, JaggedCommitmentWithWitness}; +pub use sumcheck::{JaggedSumcheckInput, QPrimeEvaluations, jagged_sumcheck_prove}; +pub use types::{JaggedBatchOpenProof, JaggedCommitment, JaggedCommitmentWithWitness, JaggedProof}; use std::{iter::once, marker::PhantomData}; -use crate::{Error, PolynomialCommitmentScheme}; +use crate::{Error, PCSFriParam, Point, PolynomialCommitmentScheme}; use ::sumcheck::structs::IOPVerifierState; use ff_ext::ExtensionField; use itertools::Itertools; use multilinear_extensions::{ - mle::FieldType, + mle::{ArcMultilinearExtension, FieldType}, util::ceil_log2, virtual_poly::{VPAuxInfo, build_eq_x_r_vec}, }; -use p3::{ - field::FieldAlgebra, - matrix::{Matrix, dense::RowMajorMatrix}, - maybe_rayon::prelude::*, -}; +use p3::{field::FieldAlgebra, matrix::Matrix, maybe_rayon::prelude::*}; +use serde::{Serialize, Serializer, de::DeserializeOwned}; +use std::sync::Arc; use transcript::Transcript; use types::int_to_field_bits; use witness::{InstancePaddingStrategy, RowMajorMatrix as WitnessRowMajorMatrix}; +#[derive(Debug)] +pub struct Jagged(PhantomData); + +impl Clone for Jagged { + fn clone(&self) -> Self { + Self(PhantomData) + } +} + +impl Serialize for Jagged { + fn serialize(&self, serializer: S) -> Result + where + S: Serializer, + { + serializer.serialize_str("jagged") + } +} + /// Commit to a sequence of row-major matrices using the Jagged PCS scheme. /// /// This function implements the commit phase described in Ceno issue #1288: @@ -155,7 +171,7 @@ use witness::{InstancePaddingStrategy, RowMajorMatrix as WitnessRowMajorMatrix}; /// Any error from the inner `InnerPcs::batch_commit` is propagated as-is. pub fn jagged_commit>( pp: &InnerPcs::ProverParam, - rmms: Vec>, + rmms: Vec>, reshape_log_height: usize, ) -> Result, Error> { if rmms.is_empty() { @@ -164,10 +180,15 @@ pub fn jagged_commit> )); } + let polys = rmms + .iter() + .flat_map(|rmm| rmm.to_mles().into_iter().map(Arc::new)) + .collect_vec(); + // --- Step 1: Compute cumulative heights from real matrix heights --- let mut poly_heights: Vec = Vec::new(); for rmm in &rmms { - let num_rows = rmm.height(); + let num_rows = rmm.occupied_physical_rows(); let num_cols = rmm.width(); if num_rows == 0 { @@ -215,7 +236,7 @@ pub fn jagged_commit> let mut poly_idx = 0; for rmm in rmms { let n_cols = rmm.width(); - let n_rows = rmm.height(); + let n_rows = rmm.occupied_physical_rows(); let n_cells = n_cols * n_rows; // The start position in `concatenated` for this matrix's block of polynomials. @@ -229,6 +250,7 @@ pub fn jagged_commit> .for_each(|(j, chunk)| { rmm.values .iter() + .take(n_cells) .skip(j) .step_by(n_cols) .zip_eq(chunk.iter_mut()) @@ -239,7 +261,7 @@ pub fn jagged_commit> } // --- Step 4: Reshape and commit via the inner PCS --- - let log_h = reshape_log_height; + let log_h = reshape_log_height.min(ceil_log2(total_size.max(1))); let h = 1usize << log_h; let w = total_size.div_ceil(h); @@ -288,10 +310,93 @@ pub fn jagged_commit> inner, cumulative_heights, poly_heights, - reshape_log_height, + reshape_log_height: log_h, + polys, }) } +fn default_reshape_log_height>( + pp: &InnerPcs::ProverParam, + rmms: &[WitnessRowMajorMatrix], +) -> usize { + let total_evals = rmms + .iter() + .map(|rmm| rmm.occupied_physical_rows() * rmm.width()) + .sum::(); + let total_log = ceil_log2(total_evals.max(1)); + pp.get_max_message_size_log().min(total_log) +} + +fn flatten_padded_openings_as_native( + poly_heights: &[usize], + openings: Vec<(Point, Vec)>, +) -> Result<(Point, Vec), Error> { + let max_native_point_len = poly_heights + .iter() + .map(|&h| ceil_log2(h)) + .max() + .unwrap_or(0); + let mut common_point = vec![None; max_native_point_len]; + let mut evals = Vec::with_capacity(poly_heights.len()); + let mut poly_idx = 0; + + for (point, point_evals) in openings { + for value in point_evals { + let height = poly_heights.get(poly_idx).ok_or_else(|| { + Error::InvalidPcsParam("jagged: too many opening evaluations".to_string()) + })?; + let native_num_vars = ceil_log2(*height); + if point.len() < native_num_vars { + return Err(Error::InvalidPcsParam(format!( + "jagged: opening point length {} is smaller than poly num_vars {}", + point.len(), + native_num_vars + ))); + } + for (dst, src) in common_point.iter_mut().zip(point.iter()) { + match dst { + Some(existing) if *existing != *src => { + return Err(Error::InvalidPcsParam( + "jagged: opening points are not prefix-compatible".to_string(), + )); + } + Some(_) => {} + None => *dst = Some(*src), + } + } + + let tail_zero_factor = point[native_num_vars..] + .iter() + .fold(E::ONE, |acc, r| acc * (E::ONE - *r)); + if tail_zero_factor == E::ZERO { + return Err(Error::InvalidPcsParam( + "jagged: padded opening tail factor is zero".to_string(), + )); + } + evals.push(value * tail_zero_factor.inverse()); + poly_idx += 1; + } + } + + if poly_idx != poly_heights.len() { + return Err(Error::InvalidPcsParam(format!( + "jagged: expected {} opening evaluations, got {}", + poly_heights.len(), + poly_idx + ))); + } + + let point = common_point + .into_iter() + .map(|value| { + value.ok_or_else(|| { + Error::InvalidPcsParam("jagged: missing common opening point".to_string()) + }) + }) + .collect::, _>>()?; + Ok((point, evals)) +} + // --------------------------------------------------------------------------- // Jagged Batch Open / Verify // --------------------------------------------------------------------------- @@ -377,35 +482,20 @@ pub fn jagged_batch_open> log_h, index & (h - 1)) gives q' without + // reconstructing a second flat buffer. let q_mles = InnerPcs::get_arc_mle_witness_from_commitment(&comm.inner); assert_eq!(q_mles.len(), w); - let q_evals_base_owned: Vec; - let q_evals_base: &[E::BaseField] = if w == 1 { - match q_mles[0].evaluations() { - FieldType::Base(slice) => slice, - _ => { - return Err(Error::InvalidPcsParam( - "jagged_batch_open: expected base-field evaluations for q'".into(), - )); - } - } - } else { - let target_len = 1usize << num_giga_vars; - q_evals_base_owned = q_mles - .iter() - .flat_map(|mle| match mle.evaluations() { - FieldType::Base(slice) => slice.iter().copied(), - _ => unreachable!("expected base-field evaluations"), - }) - .chain(std::iter::repeat_n( - E::BaseField::ZERO, - target_len - padded_total, - )) - .collect(); - &q_evals_base_owned - }; + let q_columns = q_mles + .iter() + .map(|mle| match mle.evaluations() { + FieldType::Base(slice) => Ok(slice.as_slice()), + _ => Err(Error::InvalidPcsParam( + "jagged_batch_open: expected base-field evaluations for q'".into(), + )), + }) + .collect::, _>>()?; // Batched opening claim: v = Σ_i eq_col[i] · C_i · p_i(z_row[..s_i]) // where C_i = Π_{j=s_i}^{max_s-1}(1 - z_row[j]) is the correction factor for poly i. @@ -414,7 +504,10 @@ pub fn jagged_batch_open = point.to_vec(); // Replay transcript: write evals, sample z_col. @@ -513,7 +605,7 @@ pub fn jagged_batch_verify PolynomialCommitmentScheme for Jagged +where + E: ExtensionField + Serialize + DeserializeOwned, + E::BaseField: Serialize + DeserializeOwned, + InnerPcs: PolynomialCommitmentScheme, +{ + type Param = InnerPcs::Param; + type ProverParam = InnerPcs::ProverParam; + type VerifierParam = InnerPcs::VerifierParam; + type CommitmentWithWitness = JaggedCommitmentWithWitness; + type Commitment = JaggedCommitment; + type CommitmentChunk = InnerPcs::CommitmentChunk; + type Proof = JaggedProof; + + fn setup(poly_size: usize, security_level: crate::SecurityLevel) -> Result { + InnerPcs::setup(poly_size, security_level) + } + + fn trim( + param: Self::Param, + poly_size: usize, + ) -> Result<(Self::ProverParam, Self::VerifierParam), Error> { + InnerPcs::trim(param, poly_size) + } + + fn commit( + pp: &Self::ProverParam, + rmm: WitnessRowMajorMatrix, + ) -> Result { + Self::batch_commit(pp, vec![rmm]) + } + + fn write_commitment( + comm: &Self::Commitment, + transcript: &mut impl Transcript, + ) -> Result<(), Error> { + InnerPcs::write_commitment(&comm.inner, transcript)?; + transcript + .append_field_element(&E::BaseField::from_canonical_usize(comm.reshape_log_height)); + transcript.append_field_element(&E::BaseField::from_canonical_usize( + comm.cumulative_heights.len(), + )); + for height in &comm.cumulative_heights { + transcript.append_field_element(&E::BaseField::from_canonical_usize(*height)); + } + Ok(()) + } + + fn get_pure_commitment(comm: &Self::CommitmentWithWitness) -> Self::Commitment { + comm.to_commitment() + } + + fn batch_commit( + pp: &Self::ProverParam, + rmms: Vec>, + ) -> Result { + let reshape_log_height = default_reshape_log_height::(pp, &rmms); + jagged_commit::(pp, rmms, reshape_log_height) + } + + fn open( + _pp: &Self::ProverParam, + _poly: &ArcMultilinearExtension, + _comm: &Self::CommitmentWithWitness, + _point: &[E], + _eval: &E, + _transcript: &mut impl Transcript, + ) -> Result { + unimplemented!() + } + + fn batch_open( + pp: &Self::ProverParam, + rounds: Vec<(&Self::CommitmentWithWitness, Vec<(Point, Vec)>)>, + transcript: &mut impl Transcript, + ) -> Result { + let mut proofs = Vec::with_capacity(rounds.len()); + for (comm, openings) in rounds { + let (point, evals) = flatten_padded_openings_as_native(&comm.poly_heights, openings)?; + proofs.push(jagged_batch_open::( + pp, comm, &point, &evals, transcript, + )?); + } + Ok(JaggedProof { rounds: proofs }) + } + + fn simple_batch_open( + _pp: &Self::ProverParam, + _polys: &[ArcMultilinearExtension], + _comm: &Self::CommitmentWithWitness, + _point: &[E], + _evals: &[E], + _transcript: &mut impl Transcript, + ) -> Result { + unimplemented!() + } + + fn verify( + _vp: &Self::VerifierParam, + _comm: &Self::Commitment, + _point: &[E], + _eval: &E, + _proof: &Self::Proof, + _transcript: &mut impl Transcript, + ) -> Result<(), Error> { + unimplemented!() + } + + fn batch_verify( + vp: &Self::VerifierParam, + rounds: Vec<(Self::Commitment, Vec<(usize, (Point, Vec))>)>, + proof: &Self::Proof, + transcript: &mut impl Transcript, + ) -> Result<(), Error> { + if rounds.len() != proof.rounds.len() { + return Err(Error::InvalidPcsOpeningProof(format!( + "jagged: expected {} proof rounds, got {}", + rounds.len(), + proof.rounds.len() + ))); + } + for ((comm, openings), round_proof) in rounds.into_iter().zip_eq(proof.rounds.iter()) { + let poly_heights = comm + .cumulative_heights + .windows(2) + .map(|window| window[1] - window[0]) + .collect_vec(); + let openings = openings + .into_iter() + .map(|(_, (point, evals))| (point, evals)) + .collect_vec(); + let (point, evals) = flatten_padded_openings_as_native(&poly_heights, openings)?; + jagged_batch_verify::(vp, &comm, &point, &evals, round_proof, transcript)?; + } + Ok(()) + } + + fn simple_batch_verify( + _vp: &Self::VerifierParam, + _comm: &Self::Commitment, + _point: &[E], + _evals: &[E], + _proof: &Self::Proof, + _transcript: &mut impl Transcript, + ) -> Result<(), Error> { + unimplemented!() + } + + fn get_arc_mle_witness_from_commitment( + commitment: &Self::CommitmentWithWitness, + ) -> Vec> { + commitment.polys.clone() + } +} + #[cfg(test)] mod tests { use super::*; @@ -620,17 +867,20 @@ mod tests { }; use ff_ext::GoldilocksExt2; use multilinear_extensions::mle::MultilinearExtension; - use p3::{field::FieldAlgebra, goldilocks::Goldilocks}; + use p3::{field::FieldAlgebra, goldilocks::Goldilocks, matrix::dense::RowMajorMatrix}; type F = Goldilocks; type E = GoldilocksExt2; type Pcs = Basefold; - fn make_rmm(num_rows: usize, num_cols: usize) -> RowMajorMatrix { + fn make_rmm(num_rows: usize, num_cols: usize) -> WitnessRowMajorMatrix { let values: Vec = (0..num_rows * num_cols) .map(|i| F::from_canonical_u64(i as u64 + 1)) .collect(); - RowMajorMatrix::new(values, num_cols) + WitnessRowMajorMatrix::new_by_inner_matrix( + RowMajorMatrix::new(values, num_cols), + InstancePaddingStrategy::Default, + ) } #[test] @@ -687,6 +937,22 @@ mod tests { assert_eq!(comm.total_evaluations(), 13); } + #[test] + fn test_jagged_commit_uses_unpadded_physical_rotation_height() { + // 3 logical rows with 4-way rotation occupy 12 real physical rows, padded to 16. + let reshape_log_height = 4; + let (pp, _vp) = setup_pcs::(reshape_log_height); + let rmm = WitnessRowMajorMatrix::new_by_rotation(3, 2, 2, InstancePaddingStrategy::Default); + + let comm = jagged_commit::(&pp, vec![rmm], reshape_log_height) + .expect("commit should succeed"); + + assert_eq!(comm.num_polys(), 2); + assert_eq!(comm.poly_heights, vec![12, 12]); + assert_eq!(comm.cumulative_heights, vec![0, 12, 24]); + assert_eq!(comm.total_evaluations(), 24); + } + #[test] fn test_jagged_commit_smoke() { // Two matrices: 4x1 and 4x2 → 3 polynomials, total 12 evals, padded to 16 @@ -736,14 +1002,15 @@ mod tests { use rand::thread_rng; use transcript::basic::BasicTranscript; - /// Evaluate a single column polynomial at `point`, zero-padding to `2^point.len()`. + /// Evaluate a single column polynomial at its native prefix of `point`. fn eval_column_poly_at_point(col_evals: &[F], point: &[E]) -> E { - let s = point.len(); + let s = ceil_log2(col_evals.len()); + assert!(point.len() >= s); assert!(col_evals.len() <= 1 << s); let mut padded = col_evals.to_vec(); padded.resize(1 << s, F::ZERO); let mle = MultilinearExtension::from_evaluations_vec(s, padded); - mle.evaluate(point) + mle.evaluate(&point[..s]) } #[test] @@ -773,13 +1040,13 @@ mod tests { Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); // Random evaluation point of length max_s. - // Poly i with s_i variables is evaluated at the prefix point[..s_i]. + // Poly i is evaluated at its native prefix of the common point. let point: Vec = (0..max_s).map(|_| E::random(&mut rng)).collect(); let evals: Vec = col_polys .iter() .zip(log_heights.iter()) - .map(|(col, &s_i)| eval_column_poly_at_point(col, &point[..s_i])) + .map(|(col, _)| eval_column_poly_at_point(col, &point)) .collect(); // Prover: batch open. @@ -865,7 +1132,7 @@ mod tests { let evals: Vec = col_polys .iter() .zip(log_heights.iter()) - .map(|(col, &s_i)| eval_column_poly_at_point(col, &point[..s_i])) + .map(|(col, _)| eval_column_poly_at_point(col, &point)) .collect(); let proof = jagged_batch_open::(&pp, &comm, &point, &evals, &mut transcript_p) @@ -1008,7 +1275,7 @@ mod tests { let evals: Vec = col_polys .iter() .zip(log_heights.iter()) - .map(|(col, &s_i)| eval_column_poly_at_point(col, &point[..s_i])) + .map(|(col, _)| eval_column_poly_at_point(col, &point)) .collect(); let proof = jagged_batch_open::(&pp, &comm, &point, &evals, &mut transcript_p) @@ -1028,4 +1295,53 @@ mod tests { ) .expect("batch verify should succeed"); } + + #[test] + fn test_jagged_reshape_non_power_of_two_heights() { + let mut rng = thread_rng(); + + let heights = [1023usize, 777, 513]; + let log_heights: Vec = heights.iter().map(|&h| ceil_log2(h)).collect(); + let max_s = *log_heights.iter().max().unwrap(); + let total_evals: usize = heights.iter().sum(); + let reshape_log_height = 8; + let h = 1usize << reshape_log_height; + let w = total_evals.div_ceil(h); + + let (pp, vp) = setup_pcs::(reshape_log_height); + let rmms: Vec<_> = heights.iter().map(|&h| make_rmm(h, 1)).collect(); + + let col_polys: Vec> = rmms + .iter() + .map(|rmm| (0..rmm.height()).map(|r| rmm.values[r]).collect()) + .collect(); + + let mut transcript_p = BasicTranscript::::new(b"jagged_reshape_non_power"); + let comm = jagged_commit::(&pp, rmms, reshape_log_height).expect("commit"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_p).unwrap(); + + let point: Vec = (0..max_s).map(|_| E::random(&mut rng)).collect(); + let evals: Vec = col_polys + .iter() + .zip(log_heights.iter()) + .map(|(col, _)| eval_column_poly_at_point(col, &point)) + .collect(); + + let proof = jagged_batch_open::(&pp, &comm, &point, &evals, &mut transcript_p) + .expect("batch open"); + assert_eq!(proof.col_evals.len(), w); + + let mut transcript_v = BasicTranscript::::new(b"jagged_reshape_non_power"); + Pcs::write_commitment(&comm.to_commitment().inner, &mut transcript_v).unwrap(); + + jagged_batch_verify::( + &vp, + &comm.to_commitment(), + &point, + &evals, + &proof, + &mut transcript_v, + ) + .expect("batch verify"); + } } diff --git a/crates/mpcs/src/jagged/sumcheck.rs b/crates/mpcs/src/jagged/sumcheck.rs index 8e0eddc..6aa434e 100644 --- a/crates/mpcs/src/jagged/sumcheck.rs +++ b/crates/mpcs/src/jagged/sumcheck.rs @@ -21,8 +21,9 @@ const LOG2_MAX_EPOCH: u32 = 3; /// All inputs needed for the jagged sumcheck. pub struct JaggedSumcheckInput<'a, E: ExtensionField> { - /// Giga polynomial evaluations (directly concatenated). - pub q_evals: &'a [E::BaseField], + /// Giga polynomial evaluations q'. This may be a flat concatenation or a + /// structured view over inner PCS column MLEs. + pub q_evals: QPrimeEvaluations<'a, E::BaseField>, /// n = log2(padded_total_size). pub num_giga_vars: usize, /// Cumulative height sequence t[j], length num_polys + 1. @@ -33,6 +34,47 @@ pub struct JaggedSumcheckInput<'a, E: ExtensionField> { pub eq_col: Vec, } +/// Non-owning view of q' evaluations. +/// +/// Jagged commit first concatenates all original column polynomials into a +/// single giga polynomial q'. For opening, the sumcheck only needs random +/// indexed reads from q'; it does not require q' to be materialized as one flat +/// buffer. +/// +/// `Flat` is the compact layout used by standalone sumcheck tests and by +/// callers that already own `p_0 || p_1 || ...`. +/// +/// `Columns` is the layout produced by the inner PCS commitment: q' is reshaped +/// into `w` column MLEs of height `2^log_h`. The logical flat index maps to +/// `columns[index >> log_h][index & ((1 << log_h) - 1)]`. This lets +/// `jagged_batch_open` reuse q' from `comm.inner` directly and avoids rebuilding +/// a second full-size flat q' buffer during opening. +pub enum QPrimeEvaluations<'a, F> { + /// Direct `p_0 || p_1 || ...` flat concatenation. + Flat(&'a [F]), + /// Inner PCS column layout for q'. + Columns { columns: Vec<&'a [F]>, log_h: usize }, +} + +impl QPrimeEvaluations<'_, F> { + #[inline] + fn get(&self, index: usize) -> F { + match self { + QPrimeEvaluations::Flat(evals) => evals.get(index).copied().unwrap_or_default(), + QPrimeEvaluations::Columns { columns, log_h } => { + let row_mask = (1usize << log_h) - 1; + let col = index >> log_h; + let row = index & row_mask; + columns + .get(col) + .and_then(|column| column.get(row)) + .copied() + .unwrap_or_default() + } + } + } +} + /// Iterator that yields `(col, row)` pairs for consecutive giga indices. /// Uses one binary search at construction, then O(1) per step. struct ColRowIter<'a> { @@ -89,10 +131,10 @@ impl<'a, E: ExtensionField> JaggedSumcheckInput<'a, E> { /// Brute-force computation of sum_b q'(b) * f(b). /// O(2^n) time — only for debugging and tests. pub fn compute_claimed_sum(&self) -> E { - self.q_evals[..self.total_evaluations()] - .iter() - .zip(self.col_row_iter(0)) - .fold(E::ZERO, |sum, (&q, (col, row))| { + self.col_row_iter(0) + .enumerate() + .fold(E::ZERO, |sum, (index, (col, row))| { + let q = self.q_evals.get(index); sum + (self.eq_row[row] * self.eq_col[col]) * q }) } @@ -106,7 +148,9 @@ impl<'a, E: ExtensionField> JaggedSumcheckInput<'a, E> { let total_evals = self.total_evaluations(); // Build q' MLE (padded with zeros) and evaluate at point. - let mut q_padded: Vec = self.q_evals.to_vec(); + let mut q_padded: Vec = (0..total_evals) + .map(|index| self.q_evals.get(index)) + .collect(); q_padded.resize(1 << n, Default::default()); let q_mle = MultilinearExtension::from_evaluations_vec(n, q_padded); let q_at_point = q_mle.evaluate(point); @@ -280,19 +324,18 @@ fn build_m_table( .enumerate() .for_each(|(beta, (q_b, f_b))| { let base = chunk_start + beta * a_count; - let (q_acc, f_acc) = eq_r - .iter() - .zip(input.q_evals.get(base..).unwrap_or(&[])) - .zip(input.col_row_iter(base)) - .fold( - (E::ZERO, E::ZERO), - |(q_acc, f_acc), ((&eq_r_a, &q), (col, row))| { - ( - q_acc + eq_r_a * q, - f_acc + eq_r_a * (input.eq_row[row] * input.eq_col[col]), - ) - }, - ); + let (q_acc, f_acc) = input + .col_row_iter(base) + .take(a_count) + .enumerate() + .fold((E::ZERO, E::ZERO), |(q_acc, f_acc), (a, (col, row))| { + let eq_r_a = eq_r[a]; + let q = input.q_evals.get(base + a); + ( + q_acc + eq_r_a * q, + f_acc + eq_r_a * (input.eq_row[row] * input.eq_col[col]), + ) + }); *q_b = q_acc; *f_b = f_acc; }); @@ -415,18 +458,17 @@ fn bind_and_materialize( .into_par_iter() .map(|idx| { let base = idx * a_count; - eq_r.iter() - .zip(input.q_evals.get(base..).unwrap_or(&[])) - .zip(input.col_row_iter(base)) - .fold( - (E::ZERO, E::ZERO), - |(q_acc, f_acc), ((&eq_r_a, &q), (col, row))| { - ( - q_acc + eq_r_a * q, - f_acc + eq_r_a * (input.eq_row[row] * input.eq_col[col]), - ) - }, - ) + input.col_row_iter(base).take(a_count).enumerate().fold( + (E::ZERO, E::ZERO), + |(q_acc, f_acc), (a, (col, row))| { + let eq_r_a = eq_r[a]; + let q = input.q_evals.get(base + a); + ( + q_acc + eq_r_a * q, + f_acc + eq_r_a * (input.eq_row[row] * input.eq_col[col]), + ) + }, + ) }) .collect(); @@ -469,7 +511,7 @@ mod tests { let z_col: Vec = (0..2).map(|_| E::random(&mut rng)).collect(); let input = JaggedSumcheckInput { - q_evals: &q_evals, + q_evals: QPrimeEvaluations::Flat(&q_evals), num_giga_vars, cumulative_heights: &cumulative_heights, eq_row: build_eq_x_r_vec(&z_row), @@ -529,7 +571,7 @@ mod tests { let z_col: Vec = (0..3).map(|_| E::random(&mut rng)).collect(); // ceil(log2(8))=3 let input = JaggedSumcheckInput { - q_evals: &q_evals, + q_evals: QPrimeEvaluations::Flat(&q_evals), num_giga_vars, cumulative_heights: &cumulative_heights, eq_row: build_eq_x_r_vec(&z_row), @@ -581,7 +623,7 @@ mod tests { let z_col: Vec = (0..10).map(|_| E::random(&mut rng)).collect(); // ceil(log2(1024))=10 let input = JaggedSumcheckInput { - q_evals: &q_evals, + q_evals: QPrimeEvaluations::Flat(&q_evals), num_giga_vars, cumulative_heights: &cumulative_heights, eq_row: build_eq_x_r_vec(&z_row), diff --git a/crates/mpcs/src/jagged/types.rs b/crates/mpcs/src/jagged/types.rs index 4c9c38b..0ba154a 100644 --- a/crates/mpcs/src/jagged/types.rs +++ b/crates/mpcs/src/jagged/types.rs @@ -1,6 +1,7 @@ use crate::PolynomialCommitmentScheme; use ::sumcheck::structs::IOPProof; use ff_ext::ExtensionField; +use multilinear_extensions::mle::ArcMultilinearExtension; use serde::{Deserialize, Serialize}; /// Commitment to a jagged polynomial `q'`, together with all witness data needed @@ -22,6 +23,8 @@ pub struct JaggedCommitmentWithWitness>, } /// The pure commitment (without witness data) for a jagged polynomial `q'`. @@ -76,6 +79,12 @@ pub struct JaggedBatchOpenProof> { + pub rounds: Vec>, +} + /// Convert a `usize` to its little-endian binary representation as field elements. pub(super) fn int_to_field_bits(val: usize, num_bits: usize) -> Vec { (0..num_bits) diff --git a/crates/mpcs/src/lib.rs b/crates/mpcs/src/lib.rs index f37669a..37835f6 100644 --- a/crates/mpcs/src/lib.rs +++ b/crates/mpcs/src/lib.rs @@ -217,6 +217,10 @@ pub enum PowStrategy { pub trait PCSFriParam { fn get_pow_bits_by_level(&self, pow_strategy: PowStrategy) -> usize; + + fn get_max_message_size_log(&self) -> usize { + usize::MAX + } } #[derive(Clone, Debug)] @@ -269,9 +273,10 @@ pub use basefold::{ }; pub mod jagged; pub use jagged::{ - JaggedBatchOpenProof, JaggedCommitment, JaggedCommitmentWithWitness, JaggedSumcheckInput, - assist_sumcheck_prove, evaluate_g, evaluate_g_backward, evaluate_g_forward, jagged_batch_open, - jagged_batch_verify, jagged_commit, jagged_sumcheck_prove, + Jagged, JaggedBatchOpenProof, JaggedCommitment, JaggedCommitmentWithWitness, JaggedProof, + JaggedSumcheckInput, assist_sumcheck_prove, evaluate_g, evaluate_g_backward, + evaluate_g_forward, jagged_batch_open, jagged_batch_verify, jagged_commit, + jagged_sumcheck_prove, }; #[cfg(feature = "whir")] extern crate whir as whir_external;