diff --git a/Cargo.lock b/Cargo.lock index 303a17c..a4a5bea 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -834,10 +834,10 @@ dependencies = [ "p3", "rand", "rand_chacha", - "rayon", "serde", "sumcheck", "tracing", + "tracing-forest", "tracing-subscriber", "transcript", "whir", @@ -1919,6 +1919,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..d7598ea 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 @@ -30,10 +29,11 @@ witness.workspace = true [dev-dependencies] criterion.workspace = true +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 = [] @@ -54,3 +54,15 @@ name = "interpolate" harness = false name = "whir" required-features = ["whir"] + +[[bench]] +harness = false +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..ae5305b --- /dev/null +++ b/crates/mpcs/benches/comparison.rs @@ -0,0 +1,310 @@ +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::{babybear::BabyBear, field::FieldAlgebra, matrix::Matrix, 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) -> 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(); + WitnessRowMajorMatrix::new_by_values(values, num_cols, InstancePaddingStrategy::Default) +} + +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: &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]; + 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[..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[..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 ======================== + 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 instant = std::time::Instant::now(); + let _ = Pcs::batch_commit(&direct_pp, rmms.clone()).unwrap(); + time += instant.elapsed(); + } + time + }) + }); + + let t0 = std::time::Instant::now(); + 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); + + 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[..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 new file mode 100644 index 0000000..5f4fa5c --- /dev/null +++ b/crates/mpcs/benches/jagged_pcs.rs @@ -0,0 +1,172 @@ +use std::time::Duration; + +use criterion::*; +use ff_ext::{BabyBearExt4, FromUniformBytes}; +use mpcs::{ + Basefold, BasefoldRSParams, PolynomialCommitmentScheme, SecurityLevel, jagged_batch_open, + jagged_commit, +}; +use multilinear_extensions::{util::ceil_log2, virtual_poly::build_eq_x_r_vec_sequential}; +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; +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) -> 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(); + WitnessRowMajorMatrix::new_by_values(values, num_cols, InstancePaddingStrategy::Default) +} + +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; + // 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_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]; + 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) { + 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: {:?}", heights); + + 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); + + println!( + "num_matrices={}, num_cols={}, total_evals={}, num_giga_vars={}, max_s={}", + NUM_MATRICES, NUM_COLS, total_evals, num_giga_vars, 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[..s_i])) + .collect(); + + // 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 start = num_giga_vars.min(max_log_h); + let mut vals = vec![start]; + for step in [2, 4, 6] { + let v = start.saturating_sub(step); + if v > 0 && !vals.contains(&v) { + vals.push(v); + } + } + 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"); + 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, + ); + }); + + 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 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", &label), |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| { + mpcs::jagged_batch_verify::( + &vp, + &comm.to_commitment(), + &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 new file mode 100644 index 0000000..a9f6081 --- /dev/null +++ b/crates/mpcs/benches/jagged_sumcheck.rs @@ -0,0 +1,107 @@ +use criterion::*; +use ff_ext::{BabyBearExt4, ExtensionField, FieldFrom, FromUniformBytes}; +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; + +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: QPrimeEvaluations::Flat(&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, None) + }) + }, + ); + } + + group.finish(); +} + +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/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/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 new file mode 100644 index 0000000..5ff1246 --- /dev/null +++ b/crates/mpcs/src/jagged/evaluator.rs @@ -0,0 +1,473 @@ +use ff_ext::ExtensionField; + +// --------------------------------------------------------------------------- +// Succinct evaluation of ĝ(z₁, z₂, z₃, z₄) where g(a,b,c,d) = [a+c=b ∧ b = [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) +/// that describe how the 4 ROBP states transition, grouped by effect on `lt`: +/// - `T_same[ci][co]`: weight of symbols that preserve `lt` +/// - `T_lt1[ci][co]`: weight of symbols that force `lt = 1` +/// - `T_lt0[ci][co]`: weight of symbols that force `lt = 0` +/// +/// For each (carry_in, carry_out) pair, we enumerate all consistent symbols +/// (where `a + c + carry_in` has LSB = `b` and MSB = `carry_out`) and sum +/// their eq-weights. Inconsistent symbols are absent — they correspond to +/// transitions to a rejecting sink with label 0. +#[inline] +fn transition_weights( + 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::*; + + 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}"); + } + } + + /// 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 new file mode 100644 index 0000000..c7cee04 --- /dev/null +++ b/crates/mpcs/src/jagged/mod.rs @@ -0,0 +1,1347 @@ +//! Jagged PCS +//! +//! 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 +//! +//! The "Jagged PCS" reduces proof size by packing all trace polynomials from multiple +//! chips into a single "giga" multilinear polynomial `q'`: +//! +//! ```text +//! q' = p_0 || p_1 || ... || p_{N-1} +//! ``` +//! +//! where each `p_i` is a column polynomial extracted from the input trace matrices. +//! 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 +//! 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] + h_i` +//! +//! 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]` +//! +//! 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 +//! +//! 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 +//! +//! 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|`). +//! +//! ### 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 +//! 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 +//! 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. + +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, QPrimeEvaluations, jagged_sumcheck_prove}; +pub use types::{JaggedBatchOpenProof, JaggedCommitment, JaggedCommitmentWithWitness, JaggedProof}; + +use std::{iter::once, marker::PhantomData}; + +use crate::{Error, PCSFriParam, Point, PolynomialCommitmentScheme}; +use ::sumcheck::structs::IOPVerifierState; +use ff_ext::ExtensionField; +use itertools::Itertools; +use multilinear_extensions::{ + mle::{ArcMultilinearExtension, FieldType}, + util::ceil_log2, + virtual_poly::{VPAuxInfo, build_eq_x_r_vec}, +}; +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: +/// 1. For each matrix, transpose it (row-major → column-major), so each +/// column polynomial occupies a contiguous region in memory. +/// 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`. +/// * `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 `InnerPcs::batch_commit` is propagated as-is. +pub fn jagged_commit>( + pp: &InnerPcs::ProverParam, + rmms: Vec>, + reshape_log_height: usize, +) -> Result, Error> { + if rmms.is_empty() { + return Err(Error::InvalidPcsParam( + "jagged_commit: cannot commit to empty sequence of matrices".to_string(), + )); + } + + 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.occupied_physical_rows(); + 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: 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. + #[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. + let mut poly_idx = 0; + for rmm in rmms { + let n_cols = rmm.width(); + 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. + let start = cumulative_heights[poly_idx]; + + // 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)| { + rmm.values + .iter() + .take(n_cells) + .skip(j) + .step_by(n_cols) + .zip_eq(chunk.iter_mut()) + .for_each(|(v, out)| *out = *v); + }); + + poly_idx += n_cols; + } + + // --- Step 4: Reshape and commit via the inner PCS --- + 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); + + 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 = WitnessRowMajorMatrix::::new_by_values( + giga_data, + w, + InstancePaddingStrategy::Default, + ); + + let inner = InnerPcs::batch_commit(pp, vec![giga_rmm])?; + + Ok(JaggedCommitmentWithWitness { + inner, + cumulative_heights, + poly_heights, + 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 +// --------------------------------------------------------------------------- + +/// 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_i)` are consistent with a +/// jagged commitment. +/// +/// 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 prefix `point[..s_i]`. +/// +/// 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: &InnerPcs::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 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 log_h = comm.reshape_log_height; + let h = 1usize << log_h; + let w = total_evals.div_ceil(h); + let padded_total = w * h; + let num_giga_vars = ceil_log2(padded_total); + + let z_row: Vec = point.to_vec(); + + // Write evals to transcript, then sample z_col. + transcript.append_field_element_exts(evals); + 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); + + // Reuse q' from the inner PCS MLEs. The inner PCS stores w column MLEs; + // indexing them as (index >> 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_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. + // eq_row[r] = eq(z_row, r) naturally incorporates C_i for row r within poly i + // (the high bits of r are 0, so eq_row[r] = eq(z_row[..s_i], r) * C_i). This means + // the jagged sumcheck sum Σ_{i,r} eq_col[i]*eq_row[r]*q'(t_i+r) equals the batched claim. + let eq_row = build_eq_x_r_vec(&z_row); + let input = JaggedSumcheckInput { + q_evals: QPrimeEvaluations::Columns { + columns: q_columns, + log_h, + }, + num_giga_vars, + cumulative_heights: &comm.cumulative_heights, + eq_row, + eq_col: eq_col.to_vec(), + }; + let (sumcheck_proof, challenges) = jagged_sumcheck_prove(&input, transcript, None); + let rho = challenges; + + // Compute f̂(ρ) = Σ_y eq_col[y] · ĝ(z_row, ρ, bits(t_y), bits(t_{y+1})). + let n_robp = num_giga_vars + if padded_total.is_power_of_two() { 1 } else { 0 }; + 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, + ); + + // Split ρ: low bits = row point, high bits = column selector. + let rho_row = &rho[..log_h]; + + // Compute column evaluations: evaluate each column MLE at ρ_row. + assert_eq!(q_mles.len(), w); + let col_evals: Vec = 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. + let (assist_proof, _assist_challenges) = assist_sumcheck_prove( + &z_row_padded, + &rho_padded, + &eq_col, + &comm.cumulative_heights, + n_robp, + transcript, + ); + + // Open column MLEs at ρ_row via inner PCS. + let inner_proof = InnerPcs::batch_open( + pp, + vec![(&comm.inner, vec![(rho_row.to_vec(), col_evals.clone())])], + transcript, + )?; + + Ok(JaggedBatchOpenProof { + sumcheck_proof, + col_evals, + f_at_rho, + assist_proof, + inner_proof, + }) +} + +/// 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 prefix `point[..s_i]`. +pub fn jagged_batch_verify>( + vp: &InnerPcs::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 log_h = comm.reshape_log_height; + let h = 1usize << log_h; + let w = total_evals.div_ceil(h); + let padded_total = w * h; + let num_giga_vars = ceil_log2(padded_total); + let max_s = point.len(); + let z_row: Vec = point.to_vec(); + + // Replay transcript: write evals, sample z_col. + transcript.append_field_element_exts(evals); + 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); + + // 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. + let eq_col = build_eq_x_r_vec(&z_col); + let mut tail_zero_prod = vec![E::ONE; max_s + 1]; + for j in (0..max_s).rev() { + tail_zero_prod[j] = tail_zero_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] * tail_zero_prod[s_i] * evals[i] + }) + .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 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 }; + + // 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 q_eval * proof.f_at_rho != subclaim.expected_evaluation { + return Err(Error::InvalidPcsOpen( + "jagged_batch_verify: q_eval * f(rho) != subclaim expected evaluation".into(), + )); + } + + // Verify the assist sumcheck: proves that f_at_rho is correct. + let n_assist = 2 * n_robp; + let assist_aux = VPAuxInfo { + max_degree: 2, + max_num_variables: n_assist, + phantom: PhantomData::, + }; + 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 h_at_rho_star = evaluate_g(&z_row_padded, &rho_padded, &rho_star_c, &rho_star_d); + + // 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 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: assist sumcheck final check failed".into(), + )); + } + + // Verify the inner PCS opening at ρ_row with col_evals. + InnerPcs::batch_verify( + vp, + vec![( + comm.inner.clone(), + vec![(log_h, (rho_row.to_vec(), proof.col_evals.clone()))], + )], + &proof.inner_proof, + transcript, + )?; + + Ok(()) +} + +impl 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::*; + use crate::{ + basefold::{Basefold, BasefoldRSParams}, + test_util::setup_pcs, + }; + use ff_ext::GoldilocksExt2; + use multilinear_extensions::mle::MultilinearExtension; + 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) -> WitnessRowMajorMatrix { + let values: Vec = (0..num_rows * num_cols) + .map(|i| F::from_canonical_u64(i as u64 + 1)) + .collect(); + WitnessRowMajorMatrix::new_by_inner_matrix( + RowMajorMatrix::new(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_uses_real_heights() { + // 3x1 + 5x2 → heights [3, 5, 5], not [4, 8, 8]. + let reshape_log_height = 4; + let (pp, _vp) = setup_pcs::(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_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 + 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], num_giga_vars) + .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 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], num_giga_vars).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![], 4); + assert!(matches!(result, Err(Error::InvalidPcsParam(_)))); + } + + // --- Batch open/verify tests --- + + use ff_ext::FromUniformBytes; + use rand::thread_rng; + use transcript::basic::BasicTranscript; + + /// 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 = 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[..s]) + } + + #[test] + fn test_jagged_batch_open_verify_small() { + 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 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(); + + // Extract each column polynomial. + 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, 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. + // 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, _)| 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 = ceil_log2(num_rows * num_cols); + + 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], 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(); + 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_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, _)| eval_column_poly_at_point(col, &point)) + .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(); + + let num_rows = 1024usize; + let num_cols = 3usize; + let s = 10; + 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); + + 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], 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(); + 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"); + } + + // --- 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, _)| eval_column_poly_at_point(col, &point)) + .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"); + } + + #[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 new file mode 100644 index 0000000..6aa434e --- /dev/null +++ b/crates/mpcs/src/jagged/sumcheck.rs @@ -0,0 +1,662 @@ +//! 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::*; +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 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. + 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, +} + +/// 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> { + 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. + pub fn compute_claimed_sum(&self) -> E { + 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 + }) + } + + /// 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 = (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); + + // 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) = 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; + }); + + // 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; + 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(); + + 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: QPrimeEvaluations::Flat(&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: QPrimeEvaluations::Flat(&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: QPrimeEvaluations::Flat(&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/jagged/types.rs b/crates/mpcs/src/jagged/types.rs new file mode 100644 index 0000000..0ba154a --- /dev/null +++ b/crates/mpcs/src/jagged/types.rs @@ -0,0 +1,93 @@ +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 +/// for opening proofs. +/// +/// 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 `InnerPcs`. + pub inner: InnerPcs::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, + /// 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, + /// Original per-column MLEs, in the same order as the packed jagged columns. + pub polys: 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: InnerPcs::Commitment, + /// Cumulative height sequence `t` (verifier needs this to evaluate `f(b)`). + pub cumulative_heights: Vec, + /// Log2 of the reshape column height (verifier needs this to split ρ). + pub reshape_log_height: usize, +} + +impl> + JaggedCommitmentWithWitness +{ + /// Extract the pure commitment (without witness data). + pub fn to_commitment(&self) -> JaggedCommitment { + JaggedCommitment { + inner: InnerPcs::get_pure_commitment(&self.inner), + cumulative_heights: self.cumulative_heights.clone(), + reshape_log_height: self.reshape_log_height, + } + } + + /// 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'(ρ), 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 col_evals: Vec, + pub f_at_rho: E, + pub assist_proof: IOPProof, + pub inner_proof: InnerPcs::Proof, +} + +#[derive(Clone, Serialize, Deserialize)] +#[serde(bound(serialize = "", deserialize = ""))] +pub struct JaggedProof> { + 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) + .map(|i| if (val >> i) & 1 == 1 { E::ONE } else { E::ZERO }) + .collect() +} diff --git a/crates/mpcs/src/lib.rs b/crates/mpcs/src/lib.rs index a8defea..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)] @@ -267,6 +271,13 @@ pub use basefold::{ Basefold, BasefoldCommitment, BasefoldCommitmentWithWitness, BasefoldDefault, BasefoldParams, BasefoldRSParams, BasefoldSpec, EncodingScheme, RSCode, RSCodeDefaultSpec, }; +pub mod jagged; +pub use jagged::{ + 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; #[cfg(feature = "whir")]