diff --git a/benches/expand_from_coeff.rs b/benches/expand_from_coeff.rs index 4e5905d2..cd2793ca 100644 --- a/benches/expand_from_coeff.rs +++ b/benches/expand_from_coeff.rs @@ -1,5 +1,8 @@ use divan::{black_box, AllocProfiler, Bencher}; -use whir::algebra::{fields::Field64, ntt, random_vector}; +use whir::{ + algebra::{fields::Field64, ntt, random_vector}, + buffer::{Buffer, BufferMath, BufferOps}, +}; #[global_allocator] static ALLOC: AllocProfiler = AllocProfiler::system(); @@ -28,15 +31,18 @@ fn interleaved_rs_encode(bencher: Bencher, case: &(usize, usize, usize)) { let message_length = 1 << (exp - coset_sz); let num_messages = 1 << coset_sz; let mut rng = ark_std::rand::thread_rng(); - let coeffs: Vec> = (0..num_messages) - .map(|_| random_vector(&mut rng, message_length)) + let coeffs: Vec> = (0..num_messages) + .map(|_| Buffer::from(random_vector(&mut rng, message_length))) .collect(); - (coeffs, expansion, coset_sz) + let masks = Buffer::zeros(0); + (coeffs, masks, expansion, coset_sz) }) - .bench_values(|(coeffs, expansion, _coset_sz)| { - let coeffs_refs = coeffs.iter().map(|v| v.as_slice()).collect::>(); + .bench_values(|(coeffs, masks, expansion, _coset_sz)| { + let coeffs_refs = coeffs.iter().collect::>(); + let messages = ntt::Messages::new(&coeffs_refs, coeffs[0].len(), 1); black_box(ntt::interleaved_rs_encode( - &coeffs_refs, + messages, + &masks, coeffs[0].len() * expansion, )) }); diff --git a/src/algebra/linear_form/univariate_evaluation.rs b/src/algebra/linear_form/univariate_evaluation.rs index 4d565198..4a18162d 100644 --- a/src/algebra/linear_form/univariate_evaluation.rs +++ b/src/algebra/linear_form/univariate_evaluation.rs @@ -1,9 +1,7 @@ use ark_ff::Field; use super::LinearForm; -use crate::algebra::{ - embedding::Embedding, geometric_accumulate, linear_form::Evaluate, mixed_univariate_evaluate, -}; +use crate::algebra::{embedding::Embedding, linear_form::Evaluate, mixed_univariate_evaluate}; /// Linear form to represent univariate polynomial evaluation. /// @@ -21,21 +19,6 @@ impl UnivariateEvaluation { pub const fn new(point: F, size: usize) -> Self { Self { size, point } } - - /// Batched version of [`LinearForm::accumulate`] for many [`UnivariateEvaluation`]s. - pub fn accumulate_many(evaluators: &[Self], accumulator: &mut [F], scalars: &[F]) { - assert_eq!(evaluators.len(), scalars.len()); - let Some(size) = evaluators.first().map(|e| e.size) else { - return; - }; - assert_eq!(accumulator.len(), size); - for evaluator in evaluators { - assert_eq!(evaluator.size, size); - } - let points = evaluators.iter().map(|e| e.point).collect::>(); - let scalars = scalars.to_vec(); - geometric_accumulate(accumulator, scalars, &points); - } } impl LinearForm for UnivariateEvaluation { @@ -55,7 +38,6 @@ impl LinearForm for UnivariateEvaluation { result } - /// See also [`Self::accumulate_many`] for a more efficient batched version. fn accumulate(&self, accumulator: &mut [F], scalar: F) { assert_eq!(accumulator.len(), self.size); let mut power = scalar; diff --git a/src/algebra/ntt/cooley_tukey.rs b/src/algebra/ntt/cooley_tukey.rs index 3c0377b3..729a42d1 100644 --- a/src/algebra/ntt/cooley_tukey.rs +++ b/src/algebra/ntt/cooley_tukey.rs @@ -14,11 +14,15 @@ use {crate::utils::workload_size, rayon::prelude::*, std::cmp::max}; use super::{ transpose, utils::{lcm, sqrt_factor}, - ReedSolomon, + Messages, ReedSolomon, }; #[cfg(not(feature = "rs_in_order"))] use crate::algebra::ntt::transpose::transpose_permute; -use crate::{algebra::ntt::utils::divisors, buffer::Buffer}; +use crate::{ + algebra::ntt::utils::divisors, + buffer::{Buffer, BufferOps}, + utils::{chunks_exact_or_empty, zip_strict}, +}; // Supported primes const PRIMES: [usize; 2] = [2, 3]; @@ -398,19 +402,41 @@ impl ReedSolomon for NttEngine { result } - #[cfg_attr(feature = "tracing", instrument(skip(self, polys), fields( - num_polys = polys.len(), - poly_length = polys.first().map(|p| p.len()), + #[cfg_attr(feature = "tracing", instrument(skip(self, messages, masks), fields( + num_polys = messages.vectors.len() * messages.interleaving_depth, + message_length = messages.message_length, codeword_length = codeword_length, )))] - fn interleaved_encode(&self, polys: &[&[F]], codeword_length: usize) -> Buffer { + fn interleaved_encode( + &self, + messages: Messages<'_, F>, + masks: &Buffer, + codeword_length: usize, + ) -> Buffer { + let vectors = messages + .vectors + .iter() + .map(|vector| vector.to_slice()) + .collect::>(); + let messages = vectors + .iter() + .flat_map(|vector| { + chunks_exact_or_empty(vector, messages.message_length, messages.interleaving_depth) + }) + .collect::>(); assert!(self.order.is_multiple_of(codeword_length)); - if polys.is_empty() { + if messages.is_empty() { + assert!(masks.is_empty()); return Buffer::from(Vec::new()); } - let num_polys = polys.len(); - let poly_length = polys[0].len(); - assert!(polys.iter().all(|p| p.len() == poly_length)); + let num_polys = messages.len(); + let message_length = messages[0].len(); + assert!(messages + .iter() + .all(|message| message.len() == message_length)); + assert!(masks.len().is_multiple_of(num_polys)); + let mask_length = masks.len() / num_polys; + let poly_length = message_length + mask_length; assert!(poly_length <= codeword_length); // Coset-NTT: instead of doing one codeword-length NTT on mostly zeros, @@ -434,10 +460,14 @@ impl ReedSolomon for NttEngine { // Lay out twisted coefficients in contiguous coset blocks of length // `coset_size`, zero-padding each block as needed. let mut result = Vec::with_capacity(num_polys * codeword_length); - for poly in polys { + for (message, mask) in zip_strict( + messages, + chunks_exact_or_empty(masks.to_slice(), mask_length, num_polys), + ) { // FFT[a 0 0 0] = [a a a a], so just replicate input in coset dimension. for _ in 0..num_cosets { - result.extend_from_slice(poly); + result.extend_from_slice(message); + result.extend_from_slice(mask); result.resize(result.len() + coset_padding, F::ZERO); } } diff --git a/src/algebra/ntt/mod.rs b/src/algebra/ntt/mod.rs index 3f86b043..60fc9422 100644 --- a/src/algebra/ntt/mod.rs +++ b/src/algebra/ntt/mod.rs @@ -22,7 +22,7 @@ pub use self::{ }; use crate::{ algebra::fields, - buffer::{Buffer, DefaultRs}, + buffer::{Buffer, BufferOps, DefaultRs}, type_map::{self, TypeMap}, }; @@ -62,12 +62,42 @@ impl type_map::Family for NttFamily { type Dyn = dyn ReedSolomon; } +/// Buffer-native description of interleaved Reed-Solomon messages. +/// +/// Each vector contains `interleaving_depth` consecutive messages of +/// `message_length` elements. The encoder decides how to gather those chunks, +/// append masks, and pad them for the NTT. +// TODO: Generalize `Messages` plus the separate mask buffer into an +// IRS-agnostic segmented-polynomial batch so the encoder sees only resident +// prefix/suffix segments, not message/mask protocol semantics. +#[derive(Clone)] +pub struct Messages<'a, F> { + pub vectors: &'a [&'a Buffer], + pub message_length: usize, + pub interleaving_depth: usize, +} + +impl<'a, F: Copy> Messages<'a, F> { + pub fn new( + vectors: &'a [&'a Buffer], + message_length: usize, + interleaving_depth: usize, + ) -> Self { + assert!(vectors + .iter() + .all(|vector| vector.len() == message_length * interleaving_depth)); + Self { + vectors, + message_length, + interleaving_depth, + } + } +} + /// Reed-Solomon encoder for a given field `F`. /// -/// Pure-NTT abstraction: encodes polynomials, knows nothing about how callers -/// structure those polynomials (whir's IRS, for example, concatenates a -/// message and a mask into a single polynomial before calling this trait — -/// that split lives entirely on the caller side). +/// The input remains buffer-native so each backend can gather messages, append +/// masks, and pad without forcing protocol code through host slices. pub trait ReedSolomon: Debug + Send + Sync { /// Smallest supported codeword length `>= size`, or `None` if `size` /// exceeds the engine's maximum order. The returned length is always @@ -96,16 +126,19 @@ pub trait ReedSolomon: Debug + Send + Sync { indices: &[usize], ) -> Vec; - /// Batch-encode polynomials in parallel. + /// Batch-encode masked polynomials in parallel. /// - /// All `polys[i]` must have the same length. Output is a flat buffer of - /// `polys.len() * codeword_length` elements in row-major - /// `(eval_index, poly)` layout: `result[i * polys.len() + j]` is poly - /// `j`'s value at the `i`-th evaluation point. + /// Each logical polynomial is a message chunk followed by its mask row. + /// Output is a flat buffer in row-major `(eval_index, polynomial)` layout. /// /// `codeword_length` must be NTT-smooth for this engine and at least the /// polynomial length. - fn interleaved_encode(&self, polys: &[&[F]], codeword_length: usize) -> Buffer; + fn interleaved_encode( + &self, + messages: Messages<'_, F>, + masks: &Buffer, + codeword_length: usize, + ) -> Buffer; } assert_obj_safe!(ReedSolomon); @@ -126,10 +159,14 @@ pub fn evaluation_points( .evaluation_points(poly_length, codeword_length, indices) } -pub fn interleaved_rs_encode(polys: &[&[F]], codeword_length: usize) -> Buffer { +pub fn interleaved_rs_encode( + messages: Messages<'_, F>, + masks: &Buffer, + codeword_length: usize, +) -> Buffer { NTT.get::() .expect("Unsupported NTT field.") - .interleaved_encode(polys, codeword_length) + .interleaved_encode(messages, masks, codeword_length) } pub fn generator(codeword_length: usize) -> F { @@ -149,8 +186,8 @@ mod tests { use super::*; use crate::{ - algebra::{random_vector, univariate_evaluate}, - buffer::BufferOps, + algebra::univariate_evaluate, + buffer::{BufferMath, BufferOps}, utils::zip_strict, }; @@ -191,19 +228,12 @@ mod tests { )| { let mut rng = StdRng::seed_from_u64(seed); let messages = (0..num_messages) - .map(|_| random_vector(&mut rng, message_length)) + .map(|_| Buffer::random(&mut rng, message_length)) .collect::>(); - let masks: Vec> = (0..num_messages) - .map(|_| random_vector(&mut rng, mask_length)) - .collect(); - // Build each polynomial as `message || mask`. The engine takes - // unified polynomial slices; the message/mask split is purely a - // caller-side concept. - let polys: Vec> = (0..num_messages) - .map(|i| messages[i].iter().chain(masks[i].iter()).copied().collect()) - .collect(); - let poly_refs: Vec<&[F]> = polys.iter().map(Vec::as_slice).collect(); - let codeword = ntt.interleaved_encode(&poly_refs, codeword_length); + let masks = Buffer::random(&mut rng, mask_length * num_messages); + let message_refs = messages.iter().collect::>(); + let input = Messages::new(&message_refs, message_length, 1); + let codeword = ntt.interleaved_encode(input, &masks, codeword_length); // Output must be the right size. assert_eq!(codeword.len(), codeword_length * num_messages); @@ -213,9 +243,10 @@ mod tests { let codeword = codeword.to_slice(); for (&index, &evaluation_point) in zip_strict(&sampled_indices, &evaluation_points) { let evaluations = &codeword[index * num_messages.. (index + 1) * num_messages]; - for ((message, mask), value) in zip_strict(zip_strict(&messages, &masks), evaluations) { + for (poly_index, (message, value)) in zip_strict(&messages, evaluations).enumerate() { + let mask = &masks.to_slice()[poly_index * mask_length..(poly_index + 1) * mask_length]; assert_eq!(*value, - univariate_evaluate(message, evaluation_point) + univariate_evaluate(message.to_slice(), evaluation_point) + evaluation_point.pow([message_length as u64]) * univariate_evaluate(mask, evaluation_point)); } diff --git a/src/buffer/cpu.rs b/src/buffer/cpu.rs index 47e9e917..35310cea 100644 --- a/src/buffer/cpu.rs +++ b/src/buffer/cpu.rs @@ -9,7 +9,7 @@ use zeroize::Zeroize; use crate::{ algebra::{ embedding::Embedding, - linear_form::{Covector, LinearForm, UnivariateEvaluation}, + linear_form::{Covector, LinearForm}, }, buffer::{BufferMath, BufferOps}, engines::EngineId, @@ -79,6 +79,13 @@ impl BufferOps for CpuBuffer { self.data.get(index) } + fn concat(&self, other: &Self) -> Self { + let mut data = Vec::with_capacity(self.data.len() + other.data.len()); + data.extend_from_slice(&self.data); + data.extend_from_slice(&other.data); + Self { data } + } + fn wipe(&mut self) where T: Zeroize, @@ -136,6 +143,10 @@ impl BufferMath for CpuBuffer { } } + fn resize_zeroed(&mut self, new_len: usize) { + self.data.resize(new_len, F::ZERO); + } + fn dot(&self, other: &Self) -> F { crate::algebra::dot(&self.data, &other.data) } @@ -180,13 +191,6 @@ impl BufferMath for CpuBuffer { } } - fn concat(&self, other: &Self) -> Self { - let mut data = Vec::with_capacity(self.data.len() + other.data.len()); - data.extend_from_slice(&self.data); - data.extend_from_slice(&other.data); - Self { data } - } - fn eq_weights(point: &[F]) -> Self { Self { data: crate::algebra::eq_weights(point), @@ -213,18 +217,15 @@ impl BufferMath for CpuBuffer { crate::algebra::scalar_mul(&mut self.data, weight); } - fn accumulate_univariate_evaluations( - &mut self, - evaluators: &[UnivariateEvaluation], - scalars: &Self, - ) { - let Some(size) = evaluators.first().map(|e| e.size) else { + fn accumulate_geometric(&mut self, points: &[F], scalars: &Self, prefix_len: usize) { + assert_eq!(points.len(), scalars.len()); + if points.is_empty() { return; - }; - UnivariateEvaluation::accumulate_many( - evaluators, - &mut self.data[..size], - scalars.to_slice(), + } + crate::algebra::geometric_accumulate( + &mut self.data[..prefix_len], + scalars.to_slice().to_vec(), + points, ); } @@ -236,6 +237,12 @@ impl BufferMath for CpuBuffer { crate::algebra::mixed_univariate_evaluate(embedding, &self.data, point) } + fn mixed_lift>(&self, embedding: &M) -> CpuBuffer { + CpuBuffer { + data: crate::algebra::lift(embedding, &self.data), + } + } + fn mixed_dot>( &self, embedding: &M, @@ -244,6 +251,29 @@ impl BufferMath for CpuBuffer { crate::algebra::mixed_dot(embedding, &other.data, &self.data) } + fn mixed_mat_vec>( + &self, + embedding: &M, + vector: &CpuBuffer, + ) -> CpuBuffer { + assert!( + !vector.data.is_empty(), + "matrix-vector product requires a non-empty vector" + ); + assert_eq!( + self.data.len() % vector.data.len(), + 0, + "matrix-vector dimensions mismatch" + ); + CpuBuffer { + data: self + .data + .chunks_exact(vector.data.len()) + .map(|row| crate::algebra::mixed_dot(embedding, &vector.data, row)) + .collect(), + } + } + fn mixed_sumcheck_polynomial>( &self, embedding: &M, @@ -329,7 +359,11 @@ mod tests { use ark_ff::AdditiveGroup; use super::*; - use crate::algebra::{fields::Field64, geometric_accumulate}; + use crate::algebra::{ + embedding::Basefield, + fields::{Field64, Field64_2}, + geometric_accumulate, + }; type F = Field64; @@ -343,6 +377,18 @@ mod tests { assert_eq!(buffer.to_slice(), expected.as_slice()); } + #[test] + fn resize_zeroed_preserves_prefix_and_zero_fills() { + let mut buffer = CpuBuffer::from(vec![F::from(1u64), F::from(2u64)]); + buffer.resize_zeroed(4); + assert_eq!( + buffer.to_slice(), + &[F::from(1u64), F::from(2u64), F::ZERO, F::ZERO] + ); + buffer.resize_zeroed(1); + assert_eq!(buffer.to_slice(), &[F::from(1u64)]); + } + #[test] fn accumulate_matches_geometric_accumulate_over_prefix() { let len = 8usize; @@ -352,12 +398,7 @@ mod tests { // Full-length and prefix accumulation. for size in [len, 5] { let mut buffer = CpuBuffer::from(vec![F::ZERO; len]); - let evaluators: Vec<_> = points - .iter() - .map(|&point| UnivariateEvaluation::new(point, size)) - .collect(); - buffer - .accumulate_univariate_evaluations(&evaluators, &CpuBuffer::from(scalars.clone())); + buffer.accumulate_geometric(&points, &CpuBuffer::from(scalars.clone()), size); // Reference: accumulate Σ_j scalars[j]·points[j]^i into the prefix // of a plain vector. @@ -371,4 +412,29 @@ mod tests { ); } } + + #[test] + fn mixed_mat_vec_matches_row_wise_mixed_dot() { + let matrix = CpuBuffer::from(vec![ + F::from(1u64), + F::from(2u64), + F::from(3u64), + F::from(4u64), + F::from(5u64), + F::from(6u64), + ]); + let vector = CpuBuffer::from(vec![ + Field64_2::new(F::from(7u64), F::from(1u64)), + Field64_2::new(F::from(8u64), F::from(2u64)), + Field64_2::new(F::from(9u64), F::from(3u64)), + ]); + let embedding = Basefield::::new(); + let result = matrix.mixed_mat_vec(&embedding, &vector); + let expected = matrix + .to_slice() + .chunks_exact(vector.len()) + .map(|row| crate::algebra::mixed_dot(&embedding, vector.to_slice(), row)) + .collect::>(); + assert_eq!(result.to_slice(), expected); + } } diff --git a/src/buffer/mod.rs b/src/buffer/mod.rs index 9f9488bd..14739d6e 100644 --- a/src/buffer/mod.rs +++ b/src/buffer/mod.rs @@ -21,13 +21,19 @@ use ark_std::rand::{ }; pub use cpu::CpuBuffer; -use crate::algebra::{ - embedding::Embedding, - linear_form::{LinearForm, UnivariateEvaluation}, -}; +use crate::algebra::{embedding::Embedding, linear_form::LinearForm}; + +/// Compile-time-selected prover backend. +/// +/// Keep the storage and Reed-Solomon families in one module so they cannot be +/// selected independently. A GPU backend replaces this module as a unit. +mod active_backend { + pub type Buffer = super::CpuBuffer; + pub type ReedSolomon = crate::algebra::ntt::NttEngine; +} -pub type Buffer = CpuBuffer; -pub type DefaultRs = crate::algebra::ntt::NttEngine; +pub type Buffer = active_backend::Buffer; +pub type DefaultRs = active_backend::ReedSolomon; /// Host communication for owned buffers over any copyable element type. /// @@ -53,6 +59,10 @@ pub trait BufferOps { /// Gather elements at arbitrary indices. fn gather_at_indices(&self, indices: &[usize]) -> Vec; fn get(&self, index: usize) -> Option<&T>; + /// Concatenation `[self, other]` into a single buffer of length + /// `self.len() + other.len()`. + #[must_use] + fn concat(&self, other: &Self) -> Self; /// Best-effort in-place zeroization of the buffer's contents. /// @@ -85,6 +95,10 @@ pub trait BufferMath: Clone { R: RngCore + CryptoRng, Standard: Distribution; + /// Change the logical length, filling newly exposed entries with zeroes. + /// Backends may reuse capacity or allocate/copy entirely on-device. + fn resize_zeroed(&mut self, new_len: usize); + /// Inner product with another buffer of the same length. fn dot(&self, other: &Self) -> F; @@ -108,11 +122,6 @@ pub trait BufferMath: Clone { #[must_use] fn mat_vec(&self, vector: &Self) -> Self; - /// Concatenation `[self, other]` into a single buffer of length - /// `self.len() + other.len()`. - #[must_use] - fn concat(&self, other: &Self) -> Self; - /// Equality-polynomial weights `eq(point, ·)` over the Boolean hypercube /// `{0,1}^{point.len()}`, as a buffer of length `1 << point.len()`. /// @@ -125,11 +134,13 @@ pub trait BufferMath: Clone { fn fold(&mut self, weight: F); + /// backends should override this to avoid two fold calls fn fold_pair(&mut self, other: &mut Self, weight: F) { self.fold(weight); other.fold(weight); } + /// backends should override this and use a single kernel fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { self.fold_pair(other, weight); self.sumcheck_polynomial(other) @@ -138,17 +149,9 @@ pub trait BufferMath: Clone { /// In-place scalar multiplication: `self[i] *= weight`. fn scalar_mul(&mut self, weight: F); - /// Accumulate `Σ_j scalars[j] · evaluators[j].point^i` into entry `i`. - /// - /// The evaluators must share a common size `s ≤ self.len()`; only the - /// first `s` entries are updated. This allows accumulating constraints - /// that cover a prefix of the buffer (e.g. the unmasked message part of - /// a covector). - fn accumulate_univariate_evaluations( - &mut self, - evaluators: &[UnivariateEvaluation], - scalars: &Self, - ); + /// Accumulate `Σ_j scalars[j] · points[j]^i` into entry `i` for the + /// first `prefix_len` entries of the buffer. + fn accumulate_geometric(&mut self, points: &[F], scalars: &Self, prefix_len: usize); /// Random linear combination of linear forms into a covector buffer. fn linear_forms_rlc( @@ -166,6 +169,10 @@ pub trait BufferMath: Clone { point: M::Target, ) -> M::Target; + /// Lift every element into the target field without host materialization. + #[must_use] + fn mixed_lift>(&self, embedding: &M) -> Self::TargetBuffer; + /// Inner product with a target-field buffer. fn mixed_dot>( &self, @@ -173,6 +180,16 @@ pub trait BufferMath: Clone { other: &Self::TargetBuffer, ) -> M::Target; + /// Mixed-field matrix-vector product. `self` is a row-major source-field + /// matrix with `vector.len()` columns; the result contains one target-field + /// inner product per row. + #[must_use] + fn mixed_mat_vec>( + &self, + embedding: &M, + vector: &Self::TargetBuffer, + ) -> Self::TargetBuffer; + /// Sumcheck round coefficients `(c0, c2)` for the mixed inner product of /// source-field `self` against a target-field covector. /// diff --git a/src/protocols/code_switch.rs b/src/protocols/code_switch.rs index 43e602ef..8c965a66 100644 --- a/src/protocols/code_switch.rs +++ b/src/protocols/code_switch.rs @@ -5,7 +5,7 @@ use std::{fmt, num::NonZeroUsize}; -use ark_ff::Field; +use ark_ff::{AdditiveGroup, Field}; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; use serde::{Deserialize, Serialize}; #[cfg(feature = "tracing")] @@ -15,12 +15,12 @@ use crate::{ algebra::{ dot, embedding::{Embedding, Identity}, - eq_weights, geometric_accumulate, lift, mixed_dot, scalar_mul, univariate_evaluate, + eq_weights, lift, mixed_dot, univariate_evaluate, }, - buffer::{Buffer, BufferOps}, + buffer::{Buffer, BufferMath, BufferOps}, hash::Hash, protocols::{ - geometric_challenge::geometric_challenge, + geometric_challenge::{geometric_challenge, geometric_challenge_groups_with_offset}, irs_commit::{Commitment as IrsCommitment, Config as IrsConfig, Witness as IrsWitness}, proof_of_work, }, @@ -54,7 +54,7 @@ pub struct Config { #[must_use] #[derive(Clone, Debug)] pub struct Witness { - pub message: Vec, + pub message: Buffer, pub target_witness: IrsWitness, } @@ -62,7 +62,7 @@ pub struct Witness { /// with `ℓ_zk` slack in ZK mode) paired with the running sum `μ` such that /// `μ = ⟨vector, covector⟩` after each protocol step. pub struct Claim<'a, F: Field> { - pub covector: &'a mut [F], + pub covector: &'a mut Buffer, pub sum: &'a mut F, } @@ -237,7 +237,7 @@ impl Config { pub fn prove( &self, prover_state: &mut ProverState, - message: Vec, + message: Buffer, witness: IrsWitness, claim: Claim<'_, M::Target>, folding_randomness: &[M::Target], @@ -266,8 +266,7 @@ impl Config { ); // Step 1: g := Enc_{C'}(f, r') — Construction 9.7 Step 1, p.55 - let message_buffer = Buffer::from(message.as_slice()); - let target_witness = self.target.commit(prover_state, &[&message_buffer]); + let target_witness = self.target.commit(prover_state, &[&message]); // Grind Lemma 9.9 OOD gap before α is sampled. self.pow.prove(prover_state); @@ -281,35 +280,42 @@ impl Config { // Source IRS matrix is no longer needed; release it before the trailing // arithmetic and the caller's mask-discharge phase. drop(witness); - let collapse_weights = eq_weights(folding_randomness); - let collapsed_values: Vec = source_evaluations - .matrix - .to_slice() - .chunks_exact(self.source.interleaving_depth()) - .map(|row| mixed_dot(self.source.embedding(), &collapse_weights, row)) - .collect(); + let collapse_weights = Buffer::::eq_weights(folding_randomness); + let collapsed_values = + source_evaluations.values_buffer(self.source.embedding(), &collapse_weights); // Step 4.1: batching — Construction 9.7 Step 4, p.55 let num_ood = self.out_domain_samples; let num_in_domain = source_evaluations.points.len(); - let batching_coeffs = - geometric_challenge::<_, M::Target>(prover_state, 1 + num_ood + num_in_domain); - let (&original_sl_coeff, constraint_rlc_coeffs) = batching_coeffs.split_first().unwrap(); - let (ood_rlc_coeffs, in_domain_rlc_coeffs) = constraint_rlc_coeffs.split_at(num_ood); + let (batching_base, batching_coeffs) = geometric_challenge_groups_with_offset::<_, M::Target>( + prover_state, + 1, + &[num_ood, num_in_domain], + ); + let mut batching_coeffs = batching_coeffs.into_iter(); + + // x⁰ belongs to the original claim. OOD answers are already + // transcript-sized host values, while their covector weights must stay + // resident, so derive both runs from the same Fiat–Shamir base. + let ood_rlc_coeffs = batching_coeffs.next().unwrap(); + let in_domain_rlc_coeffs = batching_coeffs.next().unwrap(); + let mut next_coeff = batching_base; + let mut ood_sum = M::Target::ZERO; + for answer in ood_answers { + ood_sum += next_coeff * answer; + next_coeff *= batching_base; + } // Mirror verifier's sum update — Construction 9.7 Decision phase, p.55. - *sum = original_sl_coeff * *sum - + dot(ood_rlc_coeffs, &ood_answers) - + dot(in_domain_rlc_coeffs, &collapsed_values); + *sum += ood_sum + in_domain_rlc_coeffs.dot(&collapsed_values); // Covector update — sl' from Completeness proof (p.55-56) let eval_points = lift(self.source.embedding(), &source_evaluations.points); - scalar_mul(covector, original_sl_coeff); self.update_covector( covector, - ood_rlc_coeffs, + &ood_rlc_coeffs, &ood_points, - in_domain_rlc_coeffs, + &in_domain_rlc_coeffs, &eval_points, ); @@ -325,7 +331,7 @@ impl Config { fn maybe_send_ood_answers( &self, prover_state: &mut ProverState, - message: &[M::Target], + message: &Buffer, mask: &[M::Target], ood_points: &[M::Target], ) -> Vec @@ -337,7 +343,7 @@ impl Config { let msg_len = message.len(); let mut answers = Vec::with_capacity(ood_points.len()); for &point in ood_points { - let f_eval = univariate_evaluate(message, point); + let f_eval = message.mixed_univariate_evaluate(&Identity::::new(), point); let answer = match &self.mode { CodeSwitchMode::Standard => f_eval, CodeSwitchMode::ZeroKnowledge { .. } => { @@ -352,34 +358,30 @@ impl Config { answers } - /// Accumulate OOD and in-domain weights into the covector. - /// Standard mode treats all points uniformly; ZK mode applies OOD over - /// the full `[f; r; s]` and in-domain over the `[f; r]` prefix only. + /// Update the resident covector. Constraint metadata is transcript-sized; + /// only the witness-sized covector remains on the selected backend. fn update_covector( &self, - covector: &mut [M::Target], - ood_rlc_coeffs: &[M::Target], + covector: &mut Buffer, + ood_rlc_coeffs: &Buffer, ood_points: &[M::Target], - in_domain_rlc_coeffs: &[M::Target], + in_domain_rlc_coeffs: &Buffer, in_domain_points: &[M::Target], ) { match &self.mode { CodeSwitchMode::Standard => { - let all_points: Vec<_> = - ood_points.iter().chain(in_domain_points).copied().collect(); - let pows: Vec<_> = ood_rlc_coeffs - .iter() - .chain(in_domain_rlc_coeffs) - .copied() - .collect(); - geometric_accumulate(covector, pows, &all_points); + let mut points = Vec::with_capacity(ood_points.len() + in_domain_points.len()); + points.extend_from_slice(ood_points); + points.extend_from_slice(in_domain_points); + let scalars = ood_rlc_coeffs.concat(in_domain_rlc_coeffs); + covector.accumulate_geometric(&points, &scalars, covector.len()); } CodeSwitchMode::ZeroKnowledge { .. } => { - geometric_accumulate(covector, ood_rlc_coeffs.to_vec(), ood_points); - geometric_accumulate( - &mut covector[..self.source.masked_message_length()], - in_domain_rlc_coeffs.to_vec(), + covector.accumulate_geometric(ood_points, ood_rlc_coeffs, covector.len()); + covector.accumulate_geometric( in_domain_points, + in_domain_rlc_coeffs, + self.source.masked_message_length(), ); } } @@ -477,7 +479,7 @@ impl Config { &self, verifier_state: &mut VerifierState, sum: &mut M::Target, - covector: &mut [M::Target], + covector: &mut Buffer, folding_randomness: &[M::Target], commitment: &IrsCommitment, ) -> VerificationResult @@ -496,12 +498,14 @@ impl Config { let (target_commitment, params) = self.verify_inner(verifier_state, sum, folding_randomness, commitment)?; - scalar_mul(covector, params.original_sl_coeff); + covector.scalar_mul(params.original_sl_coeff); + let ood_rlc_coeffs = Buffer::from(params.ood_rlc_coeffs); + let in_domain_rlc_coeffs = Buffer::from(params.in_domain_rlc_coeffs); self.update_covector( covector, - ¶ms.ood_rlc_coeffs, + &ood_rlc_coeffs, ¶ms.ood_eval_points, - ¶ms.in_domain_rlc_coeffs, + &in_domain_rlc_coeffs, ¶ms.in_domain_eval_points, ); @@ -742,7 +746,7 @@ mod tests { let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); covector.resize(config.covector_length(), F::ZERO); - let mut verifier_covector = covector.clone(); + let mut verifier_covector = Buffer::from(covector.clone()); let mut prover_sum = initial_sum; let instance = U64(seed); @@ -759,10 +763,12 @@ mod tests { let folded_message = fold_chunks(&f_full, config.source.message_length(), &folding_randomness); let mask_msg = build_mask_msg(config, &source_witness, &folding_randomness, &mut rng); + let folded_message_buffer = Buffer::from(folded_message.as_slice()); + let mut covector = Buffer::from(covector); let witness = config.prove( &mut prover_state, - folded_message.clone(), + folded_message_buffer, source_witness, Claim { covector: &mut covector, @@ -789,7 +795,7 @@ mod tests { ) .unwrap(); verifier_state.check_eof().unwrap(); - assert_eq!(witness.message, folded_message); + assert_eq!(witness.message.to_slice(), folded_message); assert_eq!(covector, verifier_covector); } @@ -803,7 +809,7 @@ mod tests { let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); covector.resize(config.covector_length(), F::ZERO); - let mut verifier_covector = covector.clone(); + let mut verifier_covector = Buffer::from(covector.clone()); let instance = U64(seed); let ds = DomainSeparator::protocol(config) @@ -833,6 +839,8 @@ mod tests { }; let initial_mu = dot(&h, &covector); let mut prover_sum = initial_mu; + let folded_message = Buffer::from(folded_message); + let mut covector = Buffer::from(covector); let _witness = config.prove( &mut prover_state, @@ -865,7 +873,7 @@ mod tests { verifier_state.check_eof().unwrap(); assert_eq!(covector, verifier_covector); - assert_eq!(dot(&h, &verifier_covector), verifier_sum); + assert_eq!(dot(&h, verifier_covector.to_slice()), verifier_sum); } fn test_tampered_ood_config>(seed: u64, config: &Config>) @@ -882,7 +890,7 @@ mod tests { let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); covector.resize(config.covector_length(), F::ZERO); - let mut verifier_covector = covector.clone(); + let mut verifier_covector = Buffer::from(covector.clone()); // Commit honest f_full, fold to get the honest post-fold message. let mut prover_state = ProverState::new_std(&ds); @@ -899,6 +907,8 @@ mod tests { // Tamper the post-fold message before proving. let mut tampered = folded_message.clone(); tampered[0] += F::ONE; + let tampered = Buffer::from(tampered); + let mut covector = Buffer::from(covector); let _witness = config.prove( &mut prover_state, tampered, @@ -930,7 +940,10 @@ mod tests { verifier_state.check_eof().unwrap(); // Sum diverges — downstream sumcheck would reject - assert_ne!(dot(&folded_message, &verifier_covector), verifier_sum); + assert_ne!( + dot(&folded_message, verifier_covector.to_slice()), + verifier_sum + ); } fn test + 'static>() diff --git a/src/protocols/geometric_challenge.rs b/src/protocols/geometric_challenge.rs index 09cacc3a..8146afff 100644 --- a/src/protocols/geometric_challenge.rs +++ b/src/protocols/geometric_challenge.rs @@ -16,14 +16,23 @@ where T: VerifierMessage, F: Field + Decoding<[T::U]>, { - match count { - 0 => Vec::new(), - 1 => vec![F::ONE], - _ => { - // Only source entropy when required - let x = transcript.verifier_message(); - geometric_sequence(F::ONE, x, count) - } + let base = geometric_challenge_base(transcript, count); + geometric_sequence(F::ONE, base, count) +} + +/// Draw the base `x` for a geometric challenge of the given total length. +/// +/// No entropy is needed for an empty challenge or the singleton `[1]`, so +/// those cases return one without touching the transcript. +pub fn geometric_challenge_base(transcript: &mut T, count: usize) -> F +where + T: VerifierMessage, + F: Field + Decoding<[T::U]>, +{ + if count > 1 { + transcript.verifier_message() + } else { + F::ONE } } @@ -40,8 +49,6 @@ where .unwrap() } -/// Sample a single geometric challenge base `x`. -/// /// Split the sequence `[1, x, x², …]` into consecutive groups of the given /// `lengths`, each returned as its own on-device buffer. Group `k` starts at /// `x^(lengths[0] + … + lengths[k-1])`. @@ -56,19 +63,35 @@ where T: VerifierMessage, F: Field + Decoding<[T::U]>, { - let total: usize = lengths.iter().sum(); - let base = if total > 1 { - transcript.verifier_message() - } else { - F::ONE - }; - let mut current = F::ONE; - lengths + geometric_challenge_groups_with_offset(transcript, 0, lengths).1 +} + +/// Draw one base and return consecutive resident groups beginning at +/// `x^offset`. +/// +/// The entropy condition includes the omitted prefix, so the result matches a +/// host [`geometric_challenge`] of length `offset + lengths.iter().sum()`. +/// Returning the base lets callers combine resident groups with +/// transcript-sized host values without transferring either one. +pub fn geometric_challenge_groups_with_offset( + transcript: &mut T, + offset: usize, + lengths: &[usize], +) -> (F, Vec>) +where + T: VerifierMessage, + F: Field + Decoding<[T::U]>, +{ + let total = offset + lengths.iter().sum::(); + let base: F = geometric_challenge_base(transcript, total); + let mut current = base.pow([offset as u64]); + let groups = lengths .iter() .map(|&len| { let group = Buffer::::geometric_challenge(current, base, len); current *= base.pow([len as u64]); group }) - .collect() + .collect(); + (base, groups) } diff --git a/src/protocols/irs_commit.rs b/src/protocols/irs_commit.rs index 859a0cf6..ec1ad4f2 100644 --- a/src/protocols/irs_commit.rs +++ b/src/protocols/irs_commit.rs @@ -38,7 +38,7 @@ use crate::{ VerifierMessage, VerifierState, }, type_info::Typed, - utils::{chunks_exact_or_empty, zip_strict}, + utils::zip_strict, }; #[derive(Clone, PartialEq, Eq, Debug, Hash, Serialize, Deserialize)] @@ -384,27 +384,8 @@ impl Config { let num_polys = self.num_messages(); let masks = Buffer::::random(prover_state.rng(), mask_length * num_polys); - // Engine takes unified polynomial slices (message || mask); the - // message/mask split is an IRS-side concept, not part of the NTT API. - let message_length = self.message_length(); - let poly_length = message_length + mask_length; - let masks_slice = masks.to_slice(); - let mut poly_buf = Vec::with_capacity(num_polys * poly_length); - let mut poly_idx = 0; - for vector in vectors { - for message in - chunks_exact_or_empty(vector.to_slice(), message_length, self.interleaving_depth) - { - poly_buf.extend_from_slice(message); - poly_buf.extend_from_slice( - &masks_slice[poly_idx * mask_length..(poly_idx + 1) * mask_length], - ); - poly_idx += 1; - } - } - debug_assert_eq!(poly_idx, num_polys); - let polys: Vec<&[M::Source]> = poly_buf.chunks_exact(poly_length).collect(); - let matrix = ntt::interleaved_rs_encode(&polys, self.codeword_length); + let messages = ntt::Messages::new(vectors, self.message_length(), self.interleaving_depth); + let matrix = ntt::interleaved_rs_encode(messages, &masks, self.codeword_length); // Commit to the matrix let matrix_witness = self.matrix_commit.commit(prover_state, &matrix); @@ -653,11 +634,14 @@ impl Evaluations { self.rows().map(|row| dot(weights, row)) } - /// Buffer-native [`values`](Self::values): a matrix-vector product on the - /// backend, returning one value per point. Both the matrix and the weights - /// stay on-device, so no readback of the (potentially large) weights is - /// forced. Used by the prover. - pub fn values_buffer(&self, weights: &Buffer) -> Buffer { + /// Reduce each row against resident `weights` through `embedding`, yielding + /// one resident target-field value per point. An [`Identity`](crate::algebra::embedding::Identity) + /// embedding covers the same-field case, so prover code has one path for + /// both same- and mixed-field reductions. + pub fn values_buffer(&self, embedding: &M, weights: &Buffer) -> Buffer + where + M: Embedding, + { let num_points = self.num_points(); if num_points == 0 { assert!(self.matrix.is_empty(), "evaluation matrix has no points"); @@ -677,7 +661,7 @@ impl Evaluations { if num_columns == 0 { return Buffer::zeros(num_points); } - self.matrix.mat_vec(weights) + self.matrix.mixed_mat_vec(embedding, weights) } } diff --git a/src/protocols/sumcheck.rs b/src/protocols/sumcheck.rs index 7e2fb0b9..f095879f 100644 --- a/src/protocols/sumcheck.rs +++ b/src/protocols/sumcheck.rs @@ -9,7 +9,7 @@ use serde::{Deserialize, Serialize}; use tracing::instrument; use crate::{ - algebra::{embedding::Embedding, lift, univariate_evaluate}, + algebra::{embedding::Embedding, univariate_evaluate}, buffer::{Buffer, BufferMath, BufferOps}, protocols::proof_of_work, transcript::{ @@ -262,10 +262,10 @@ impl Config { folded } // No rounds: nothing folds, but the caller still expects a - // target-field buffer. Cold path; a plain lift is fine. + // target-field buffer. Keep the lift on the selected backend. (None, None) => { let a = a.take().expect("source buffer consumed once"); - Buffer::from(lift(embedding, a.to_slice())) + a.mixed_lift(embedding) } (Some(_), None) => unreachable!("folded buffer implies a prior challenge"), }; diff --git a/src/protocols/whir/prover.rs b/src/protocols/whir/prover.rs index 84ae72bc..814ab19d 100644 --- a/src/protocols/whir/prover.rs +++ b/src/protocols/whir/prover.rs @@ -98,36 +98,36 @@ impl Config { // Complete evaluations of EVERY vector at EVERY linear form. let (oods_evals, oods_matrix) = { - let mut oods_evals = Vec::new(); + let mut oods_points = Vec::new(); let mut oods_matrix = Vec::new(); // Out of domain samples. Compute missing cross-terms and send to verifier. let mut vector_offset = 0; for witness in &witnesses { for (oods_eval, oods_row) in zip_strict( - witness.out_of_domain.evaluators(self.initial_size()), + witness.out_of_domain.points.iter().copied(), witness.out_of_domain.rows(), ) { for (j, vector) in vectors.iter().enumerate() { if j >= vector_offset && j < oods_row.len() + vector_offset { debug_assert_eq!( oods_row[j - vector_offset], - vector.mixed_univariate_evaluate(self.embedding(), oods_eval.point) + vector.mixed_univariate_evaluate(self.embedding(), oods_eval) ); oods_matrix.push(oods_row[j - vector_offset]); } else { let eval = - vector.mixed_univariate_evaluate(self.embedding(), oods_eval.point); + vector.mixed_univariate_evaluate(self.embedding(), oods_eval); prover_state.prover_message(&eval); oods_matrix.push(eval); } } - oods_evals.push(oods_eval); + oods_points.push(oods_eval); } vector_offset += witness.num_vectors(); } - (oods_evals, oods_matrix) + (oods_points, oods_matrix) }; // Random linear combination of the vectors. @@ -172,7 +172,7 @@ impl Config { debug_assert!(!has_constraints || vector.dot(&covector) == the_sum); // Add OODS constraints - covector.accumulate_univariate_evaluations(&oods_evals, &oods_rlc_coeffs); + covector.accumulate_geometric(&oods_evals, &oods_rlc_coeffs, self.initial_size()); let oods_matrix = Buffer::from(oods_matrix); the_sum += oods_matrix.bilinear_form(&oods_rlc_coeffs, &vector_rlc_coeffs); drop(oods_evals); @@ -241,19 +241,25 @@ impl Config { }; // Collect constraints for this round and RLC them in - let stir_challenges = out_of_domain - .evaluators(round_config.initial_size()) - .chain(in_domain.evaluators(round_config.initial_size())) + let stir_points = out_of_domain + .points + .iter() + .chain(&in_domain.points) + .copied() .collect::>(); // Weights for the in-domain rows: vector_rlc_coeffs ⊗ eq(folding_randomness), // built directly on the backend so no readback is needed. let stir_weights = vector_rlc_coeffs.tensor_product(&Buffer::eq_weights(&folding_randomness)); let stir_evaluations = out_of_domain - .values_buffer(&Buffer::ones(1)) - .concat(&in_domain.values_buffer(&stir_weights)); - let stir_rlc_coeffs = geometric_challenge_buffer(prover_state, stir_challenges.len()); - covector.accumulate_univariate_evaluations(&stir_challenges, &stir_rlc_coeffs); + .values_buffer(&Identity::new(), &Buffer::ones(1)) + .concat(&in_domain.values_buffer(&Identity::new(), &stir_weights)); + let stir_rlc_coeffs = geometric_challenge_buffer(prover_state, stir_points.len()); + covector.accumulate_geometric( + &stir_points, + &stir_rlc_coeffs, + round_config.initial_size(), + ); the_sum += stir_rlc_coeffs.dot(&stir_evaluations); debug_assert_eq!(vector.dot(&covector), the_sum); diff --git a/src/protocols/zook/commit.rs b/src/protocols/zook/commit.rs index 75adb2ed..9e9b690f 100644 --- a/src/protocols/zook/commit.rs +++ b/src/protocols/zook/commit.rs @@ -8,8 +8,8 @@ use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, R use tracing::instrument; use crate::{ - algebra::{embedding::Embedding, lift}, - buffer::Buffer, + algebra::embedding::Embedding, + buffer::{Buffer, BufferMath, BufferOps}, hash::Hash, protocols::{ irs_commit::{Commitment as IrsCommitment, Witness as IrsWitness}, @@ -34,12 +34,12 @@ pub(crate) enum CommittedState { /// The message stays in `M::Source`: the first round's sumcheck lifts it /// into `M::Target` at its first fold. Round { - message: Vec, + message: Buffer, irs_witness: IrsWitness, }, /// Basecase-only plan; witness was lifted into `M::Target` first. Basecase { - message: Vec, + message: Buffer, irs_witness: IrsWitness, }, } @@ -52,12 +52,12 @@ pub struct Commitment { } impl ProtocolConfig { - /// Commit the initial witness to the protocol's first IRS codeword. + /// Commit an already-resident witness to the protocol's first IRS codeword. #[cfg_attr(feature = "tracing", instrument(skip_all, name = "zook::commit", fields(vector_size = self.tuning().vector_size, num_rounds = self.num_rounds())))] pub fn commit( &self, ps: &mut ProverState, - witness: &[M::Source], + witness: Buffer, ) -> CommittedWitness where Standard: Distribution + Distribution, @@ -73,18 +73,16 @@ impl ProtocolConfig { ); let state = if let Some(round) = self.first_round() { - let witness_buffer = Buffer::from(witness); - let irs_witness = round.code_switch().source().commit(ps, &[&witness_buffer]); + let irs_witness = round.code_switch().source().commit(ps, &[&witness]); CommittedState::Round { - message: witness.to_vec(), + message: witness, irs_witness, } } else { // Basecase IRS is over `M::Target`; lift before committing. let embedding = M::default(); - let message = lift(&embedding, witness); - let message_buffer = Buffer::from(message.as_slice()); - let irs_witness = self.basecase().commit().commit(ps, &[&message_buffer]); + let message = witness.mixed_lift(&embedding); + let irs_witness = self.basecase().commit().commit(ps, &[&message]); CommittedState::Basecase { message, irs_witness, @@ -115,7 +113,6 @@ mod tests { use super::*; use crate::{ - algebra::random_vector, hash, protocols::params::{ spec::{ @@ -167,14 +164,14 @@ mod tests { seed: u64, ) -> CommittedWitness { let mut rng = StdRng::seed_from_u64(seed); - let witness = random_vector::(&mut rng, config.tuning().vector_size); + let witness = Buffer::::random(&mut rng, config.tuning().vector_size); let ds = DomainSeparator::protocol(&"zook-commit-test") .session(&format!("commit roundtrip {}:{}", file!(), line!())) .instance(&Empty); let mut prover_state = ProverState::new_std(&ds); - let committed = config.commit(&mut prover_state, &witness); + let committed = config.commit(&mut prover_state, witness); let proof = prover_state.proof(); let mut verifier_state = VerifierState::new_std(&ds, &proof); @@ -232,12 +229,12 @@ mod tests { ) .unwrap(); let mut rng = StdRng::seed_from_u64(3); - let too_short = random_vector::(&mut rng, config.tuning().vector_size - 1); + let too_short = Buffer::::random(&mut rng, config.tuning().vector_size - 1); let ds = DomainSeparator::protocol(&"zook-commit-test") .session(&format!("wrong size {}:{}", file!(), line!())) .instance(&Empty); let mut prover_state = ProverState::new_std(&ds); - let _ = config.commit(&mut prover_state, &too_short); + let _ = config.commit(&mut prover_state, too_short); } } diff --git a/src/protocols/zook/mod.rs b/src/protocols/zook/mod.rs index b0212033..48eae971 100644 --- a/src/protocols/zook/mod.rs +++ b/src/protocols/zook/mod.rs @@ -71,6 +71,7 @@ mod tests { linear_form::{Evaluate, LinearForm, MultilinearExtension}, random_vector, }, + buffer::{Buffer, BufferMath, BufferOps}, hash, protocols::params::spec::{ DecodingRegime, FoldingFactor, Mode, PowBudget, RateSchedule, SecuritySpec, TuningSpec, @@ -141,7 +142,7 @@ mod tests { ) { let embedding = ::default(); let mut rng = StdRng::seed_from_u64(seed); - let witness: Vec = random_vector(&mut rng, config.tuning().vector_size); + let witness = Buffer::::random(&mut rng, config.tuning().vector_size); let mu = config.tuning().vector_size.trailing_zeros() as usize; let forms: Vec> = (0..num_claims) @@ -151,14 +152,14 @@ mod tests { .collect(); let values: Vec = forms .iter() - .map(|f| f.evaluate(&embedding, &witness)) + .map(|f| f.evaluate(&embedding, witness.to_slice())) .collect(); let form_refs: Vec<&dyn LinearForm> = forms.iter().map(|f| f as &dyn LinearForm).collect(); let ds = make_ds(label); let mut ps = ProverState::new_std(&ds); - let committed = config.commit(&mut ps, &witness); + let committed = config.commit(&mut ps, witness); config.prove(&mut ps, committed, &form_refs, &values); let proof = ps.proof(); @@ -204,7 +205,7 @@ mod tests { ) { let embedding = ::default(); let mut rng = StdRng::seed_from_u64(seed); - let witness: Vec = random_vector(&mut rng, config.tuning().vector_size); + let witness = Buffer::::random(&mut rng, config.tuning().vector_size); let mu = config.tuning().vector_size.trailing_zeros() as usize; let forms: Vec> = (0..num_claims) @@ -214,7 +215,7 @@ mod tests { .collect(); let values: Vec = forms .iter() - .map(|f| f.evaluate(&embedding, &witness)) + .map(|f| f.evaluate(&embedding, witness.to_slice())) .collect(); let form_refs: Vec<&dyn LinearForm> = forms .iter() @@ -223,7 +224,7 @@ mod tests { let ds = make_ds(label); let mut ps = ProverState::new_std(&ds); - let committed = config.commit(&mut ps, &witness); + let committed = config.commit(&mut ps, witness); config.prove(&mut ps, committed, &form_refs, &values); let proof = ps.proof(); @@ -334,7 +335,7 @@ mod tests { ProtocolConfig::::derive(small_spec(Mode::Standard), multi_round_tuning()) .unwrap(); let mut rng = StdRng::seed_from_u64(0); - let witness: Vec = random_vector(&mut rng, config.tuning().vector_size); + let witness = Buffer::::random(&mut rng, config.tuning().vector_size); let mu = config.tuning().vector_size.trailing_zeros() as usize; let form: MultilinearExtension = MultilinearExtension { point: random_vector(&mut rng, mu), @@ -343,7 +344,7 @@ mod tests { .session(&"count-mismatch".to_string()) .instance(&Empty); let mut ps = ProverState::new_std(&ds); - let committed = config.commit(&mut ps, &witness); + let committed = config.commit(&mut ps, witness); // 1 form but 2 values — should panic config.prove( &mut ps, @@ -360,12 +361,12 @@ mod tests { ProtocolConfig::::derive(small_spec(Mode::Standard), multi_round_tuning()) .unwrap(); let mut rng = StdRng::seed_from_u64(0); - let witness: Vec = random_vector(&mut rng, config.tuning().vector_size); + let witness = Buffer::::random(&mut rng, config.tuning().vector_size); let ds = DomainSeparator::protocol(&"zook-test") .session(&"empty-forms".to_string()) .instance(&Empty); let mut ps = ProverState::new_std(&ds); - let committed = config.commit(&mut ps, &witness); + let committed = config.commit(&mut ps, witness); // No forms at all — should panic config.prove(&mut ps, committed, &[], &[]); } @@ -539,7 +540,7 @@ mod tests { ProtocolConfig::::derive(large_spec(mode), large_tuning()).unwrap(); let mut rng = StdRng::seed_from_u64(seed); - let witness: Vec = random_vector(&mut rng, config.tuning().vector_size); + let witness = Buffer::::random(&mut rng, config.tuning().vector_size); let mu = config.tuning().vector_size.trailing_zeros() as usize; let embedding = ::default(); @@ -550,7 +551,7 @@ mod tests { .collect(); let values: Vec = forms .iter() - .map(|f| f.evaluate(&embedding, &witness)) + .map(|f| f.evaluate(&embedding, witness.to_slice())) .collect(); let form_refs: Vec<&dyn LinearForm> = forms.iter().map(|f| f as &dyn LinearForm).collect(); @@ -560,7 +561,7 @@ mod tests { .instance(&Empty); let mut ps = ProverState::new_std(&ds); - let committed = config.commit(&mut ps, &witness); + let committed = config.commit(&mut ps, witness); config.prove(&mut ps, committed, &form_refs, &values); let proof = ps.proof(); diff --git a/src/protocols/zook/prover.rs b/src/protocols/zook/prover.rs index 5e9710c8..3770a045 100644 --- a/src/protocols/zook/prover.rs +++ b/src/protocols/zook/prover.rs @@ -52,9 +52,9 @@ use crate::{ embedding::{Embedding, Identity}, geometric_sequence, linear_form::LinearForm, - mixed_dot, random_vector, univariate_evaluate, + random_vector, univariate_evaluate, }, - buffer::{Buffer, BufferOps}, + buffer::{Buffer, BufferMath, BufferOps}, hash::Hash, protocols::{ code_switch::{self, mixed_fold_chunks}, @@ -120,6 +120,7 @@ impl ProtocolConfig { .zip(&claim_weights) .map(|(v, weight)| *v * weight) .sum(); + let covector = Buffer::from(covector); // Reduce to basecase inputs `(message, witness, covector, sum)`. The // two arms differ only in how those are obtained. @@ -156,13 +157,9 @@ impl ProtocolConfig { // Standard mode (BasecaseMode::Standard) sends the full witness vector // and IRS randomness cleartext. Only call with Mode::ZeroKnowledge if // end-to-end hiding is required. - let _ = self.basecase().prove( - ps, - Buffer::from(message), - &basecase_witness, - Buffer::from(covector), - sum, - ); + let _ = self + .basecase() + .prove(ps, message, &basecase_witness, covector, sum); } } @@ -172,9 +169,9 @@ impl ProtocolConfig { /// into an `M::Target` one, so [`prove_round`] maps `ProverRoundState` to /// `ProverRoundState>`. struct ProverRoundState { - message: Vec, + message: Buffer, irs_witness: IrsWitness, - covector: Vec, + covector: Buffer, sum: M::Target, } @@ -208,7 +205,7 @@ where let embedding = round.code_switch().source().embedding(); debug_assert_eq!( - mixed_dot(embedding, &covector, &message), + message.mixed_dot(embedding, &covector), sum, "prove_round entry: dot(message, covector) must equal sum" ); @@ -218,22 +215,16 @@ where let mut masker = RoundMaskOracle::begin(round, ps); // Sumcheck lifts the source-field message into `M::Target` at its first - // fold and returns the folded buffer (the covector folds in place). Move - // the host-side round state into buffers, fold, and move the folded result - // back into the `Vec` state (which downstream steps resize/truncate/index - // directly). The hops are zero-copy on the CPU backend. - let message_buf = Buffer::from(message); - let mut covector_buf = Buffer::from(covector); - let (message_buf, opening) = round.sumcheck().prove( + // fold. Both outputs remain resident for code-switch and the next round. + let mut covector = covector; + let (message, opening) = round.sumcheck().prove( ps, embedding, - message_buf, - &mut covector_buf, + message, + &mut covector, &mut sum, masker.sumcheck_blinding(), ); - let message = message_buf.into_vec(); - let mut covector = covector_buf.into_vec(); // Build cs_mask = (folded_irs_masks ‖ cs_fresh_padding), commit its tree, // send mask_eval_sum cleartext, reconcile sum to the unmasked dot. @@ -246,13 +237,13 @@ where ); debug_assert_eq!( - dot(&message, &covector), + message.dot(&covector), sum, "post-reconcile: dot(message, covector) must equal sum" ); // Extend covector for ZK mask region; +0 in Standard mode. - covector.resize(msg_len + masker.covector_extension(), M::Target::ZERO); + covector.resize_zeroed(msg_len + masker.covector_extension()); let cs_witness = round.code_switch().prove( ps, message, @@ -266,19 +257,16 @@ where ); // Prove both mask trees; subtract cs_mask contribution to project sum to f-only. - masker.finish( - &opening.round_challenges, - &covector[msg_len..], - &mut sum, - ps, - ); + let cs_mask_indices = (msg_len..covector.len()).collect::>(); + let cs_mask_covector = covector.gather_at_indices(&cs_mask_indices); + masker.finish(&opening.round_challenges, &cs_mask_covector, &mut sum, ps); drop(opening); let message = cs_witness.message; - covector.truncate(message.len()); + covector.resize_zeroed(message.len()); debug_assert_eq!( - dot(&message, &covector), + message.dot(&covector), sum, "prove_round exit: dot(message, covector) must equal sum" ); diff --git a/tests/resident_transfer_budget.rs b/tests/resident_transfer_budget.rs new file mode 100644 index 00000000..4f0d53d0 --- /dev/null +++ b/tests/resident_transfer_budget.rs @@ -0,0 +1,47 @@ +//! Guard the two witness-sized transfer seams fixed by the resident prover path. +//! +//! This deliberately checks the protocol boundary rather than CPU backend +//! internals: CPU encoders may use slices privately, while resident protocol +//! code must not materialize large host vectors. + +fn section<'a>(source: &'a str, start: &str, end: &str) -> &'a str { + let start = source.find(start).expect("start marker"); + let tail = &source[start..]; + let end = tail.find(end).expect("end marker"); + &tail[..end] +} + +#[test] +fn resident_hot_path_has_no_full_host_bounces() { + let irs = include_str!("../src/protocols/irs_commit.rs"); + let irs_commit = section( + irs, + " pub fn commit(", + " /// Receive a commitment", + ); + assert!(!irs_commit.contains(".to_slice()")); + assert!(!irs_commit.contains(".into_vec()")); + assert!(!irs_commit.contains("poly_buf")); + + let zook = include_str!("../src/protocols/zook/prover.rs"); + let prove_round = section( + zook, + "fn prove_round(", + "/// The committed mask tree", + ); + assert!(!prove_round.contains(".into_vec()")); + assert!(!prove_round.contains("Buffer::from(message)")); + assert!(!prove_round.contains("Buffer::from(covector)")); + + let code_switch = include_str!("../src/protocols/code_switch.rs"); + let prove = section( + code_switch, + " pub fn prove(", + " /// Send OOD answers", + ); + assert!(!prove.contains("Buffer::from(message")); + assert!(!prove.contains("message.into_vec()")); + assert!(!prove.contains("covector.into_vec()")); + assert!(!prove.contains("Buffer::from(ood_answers)")); + assert!(!prove.contains(".to_slice()")); +}