diff --git a/benches/expand_from_coeff.rs b/benches/expand_from_coeff.rs index c244514e..bdcfde9a 100644 --- a/benches/expand_from_coeff.rs +++ b/benches/expand_from_coeff.rs @@ -1,5 +1,11 @@ use divan::{black_box, AllocProfiler, Bencher}; -use whir::algebra::{fields::Field64, ntt, random_vector}; +use whir::{ + algebra::{ + fields::Field64, + ntt::{Messages, NttEngine, ReedSolomon}, + }, + buffer::{ActiveBuffer, Buffer, BufferOps}, +}; #[global_allocator] static ALLOC: AllocProfiler = AllocProfiler::system(); @@ -28,18 +34,17 @@ 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(|_| ActiveBuffer::random(&mut rng, message_length)) .collect(); - (coeffs, expansion, coset_sz) + let engine = NttEngine::::new_from_fftfield(); + (engine, coeffs, expansion) }) - .bench_values(|(coeffs, expansion, _coset_sz)| { - let coeffs_refs = coeffs.iter().map(|v| v.as_slice()).collect::>(); - black_box(ntt::interleaved_rs_encode( - &coeffs_refs, - &[], - coeffs[0].len() * expansion, - )) + .bench_values(|(engine, coeffs, expansion)| { + let coeffs_refs = coeffs.iter().collect::>(); + let messages = Messages::new(&coeffs_refs, coeffs[0].len(), 1); + let masks = ActiveBuffer::from_slice(&[]); + black_box(engine.interleaved_encode(messages, &masks, coeffs[0].len() * expansion)) }); } diff --git a/src/algebra/ntt/cooley_tukey.rs b/src/algebra/ntt/cooley_tukey.rs index fde9dc5d..67e17f17 100644 --- a/src/algebra/ntt/cooley_tukey.rs +++ b/src/algebra/ntt/cooley_tukey.rs @@ -14,12 +14,13 @@ 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::{ActiveBuffer, BufferOps}, utils::{chunks_exact_or_empty, zip_strict}, }; @@ -402,17 +403,33 @@ impl ReedSolomon for NttEngine { } #[cfg_attr(feature = "tracing", instrument(skip(self, messages, masks), fields( - num_messages = messages.len(), - message_len = messages.first().map(|c| c.len()), + num_messages = messages.vectors.len() * messages.interleaving_depth, + message_len = messages.message_length, codeword_length = codeword_length, - mask_len = masks.len().checked_div(messages.len()) + mask_len = masks.len().checked_div(messages.vectors.len() * messages.interleaving_depth) )))] - fn interleaved_encode(&self, messages: &[&[F]], masks: &[F], codeword_length: usize) -> Vec { + fn interleaved_encode( + &self, + messages: Messages<'_, F>, + masks: &ActiveBuffer, + codeword_length: usize, + ) -> ActiveBuffer { + let vectors = messages + .vectors + .iter() + .map(|v| v.to_slice()) + .collect::>(); + let messages = vectors + .iter() + .flat_map(|v| { + chunks_exact_or_empty(v, messages.message_length, messages.interleaving_depth) + }) + .collect::>(); assert!(self.order.is_multiple_of(codeword_length)); if messages.is_empty() { assert!(masks.is_empty()); - return Vec::new(); + return ActiveBuffer::from_vec(Vec::new()); } let num_messages = messages.len(); let message_len = messages[0].len(); @@ -445,7 +462,7 @@ impl ReedSolomon for NttEngine { let mut result = Vec::with_capacity(num_messages * codeword_length); for (message, mask) in zip_strict( messages, - chunks_exact_or_empty(masks, mask_length, num_messages), + chunks_exact_or_empty(masks.to_slice(), mask_length, num_messages), ) { // FFT[a 0 0 0] = [a a a a], so just replicate input in coset dimension. for _ in 0..num_cosets { @@ -473,7 +490,7 @@ impl ReedSolomon for NttEngine { // Transpose to row-major order with vectors stacked horizontally. transpose(&mut result, num_messages, codeword_length); - result + ActiveBuffer::from_vec(result) } } diff --git a/src/algebra/ntt/mod.rs b/src/algebra/ntt/mod.rs index 9e3fc74d..9deeca44 100644 --- a/src/algebra/ntt/mod.rs +++ b/src/algebra/ntt/mod.rs @@ -22,34 +22,35 @@ pub use self::{ }; use crate::{ algebra::fields, + buffer::{ActiveBuffer, BufferOps, DefaultRs}, type_map::{self, TypeMap}, }; pub static NTT: LazyLock> = LazyLock::new(|| { let map = TypeMap::new(); map.insert( - Arc::new(NttEngine::::new_from_fftfield()) as Arc> + Arc::new(DefaultRs::::new_from_fftfield()) as Arc> ); map.insert( - Arc::new(NttEngine::::new_from_fftfield()) as Arc> + Arc::new(DefaultRs::::new_from_fftfield()) as Arc> ); map.insert( - Arc::new(NttEngine::::new_from_fftfield()) as Arc> + Arc::new(DefaultRs::::new_from_fftfield()) as Arc> ); map.insert( - Arc::new(NttEngine::::new_from_fftfield()) as Arc> + Arc::new(DefaultRs::::new_from_fftfield()) as Arc> ); map.insert( - Arc::new(NttEngine::::new_from_fftfield()) as Arc>, + Arc::new(DefaultRs::::new_from_fftfield()) as Arc>, ); map.insert( - Arc::new(NttEngine::::new_from_fftfield()) as Arc>, + Arc::new(DefaultRs::::new_from_fftfield()) as Arc>, ); map.insert(Arc::new( - NttEngine::<::BasePrimeField>::new_from_fftfield(), + DefaultRs::<::BasePrimeField>::new_from_fftfield(), ) as Arc>); map.insert(Arc::new( - NttEngine::<::BasePrimeField>::new_from_fftfield(), + DefaultRs::<::BasePrimeField>::new_from_fftfield(), ) as Arc>); map }); @@ -61,6 +62,37 @@ impl type_map::Family for NttFamily { type Dyn = dyn ReedSolomon; } +/// Interleaved Reed-Solomon input buffers. +/// +/// Chunking buffer-backed vectors into per-message views can require backend work. This descriptor +/// keeps the encoding input compact and lets each backend derive the chunked messages internally. +/// +/// Each vector stores `interleaving_depth` consecutive messages of `message_length` field +/// elements. +#[derive(Clone, Copy)] +pub struct Messages<'a, F> { + pub(crate) vectors: &'a [&'a ActiveBuffer], + pub(crate) message_length: usize, + pub(crate) interleaving_depth: usize, +} + +impl<'a, F: Field> Messages<'a, F> { + pub fn new( + vectors: &'a [&'a ActiveBuffer], + message_length: usize, + interleaving_depth: usize, + ) -> Self { + assert!(vectors + .iter() + .all(|vector| vector.len() == message_length * interleaving_depth)); + Self { + vectors, + message_length, + interleaving_depth, + } + } +} + /// Trait for a Reed-Solomon encoder implementation for a given field `F`. pub trait ReedSolomon: Debug + Send + Sync { /// Returns the next supported order equal or larger than `size`. @@ -88,15 +120,20 @@ pub trait ReedSolomon: Debug + Send + Sync { /// Compute a masked interleaved Reed-Solomon encoding. /// - /// `messages` are `num_messages` slices of `message_length` elements. - /// `masks` is a `num_messages` × `mask_length` matrix of blinding coefficients. + /// `messages` contains `num_vectors` buffers, each holding `interleaving_depth` messages of + /// `message_length` elements laid out contiguously. `masks` is a + /// `(num_vectors * interleaving_depth) × mask_length` matrix of blinding coefficients. /// `codeword_length` must be an NTT-smooth number >= `message_length + mask_length`. - /// returns an `codeword_length × num_messages` matrix. + /// Returns a `codeword_length × (num_vectors * interleaving_depth)` matrix. /// /// Each output value is the univariate polynomial evaluation in the evaluation point /// corresponding with the index of a coefficient list formed by concatenating message and mask. - /// - fn interleaved_encode(&self, messages: &[&[F]], masks: &[F], codeword_length: usize) -> Vec; + fn interleaved_encode( + &self, + messages: Messages<'_, F>, + masks: &ActiveBuffer, + codeword_length: usize, + ) -> ActiveBuffer; } assert_obj_safe!(ReedSolomon); @@ -118,10 +155,10 @@ pub fn evaluation_points( } pub fn interleaved_rs_encode( - messages: &[&[F]], - masks: &[F], + messages: Messages<'_, F>, + masks: &ActiveBuffer, codeword_length: usize, -) -> Vec { +) -> ActiveBuffer { NTT.get::() .expect("Unsupported NTT field.") .interleaved_encode(messages, masks, codeword_length) @@ -145,6 +182,7 @@ mod tests { use super::*; use crate::{ algebra::{random_vector, univariate_evaluate}, + buffer::BufferOps, utils::{chunks_exact_or_empty, zip_strict}, }; @@ -188,10 +226,16 @@ mod tests { .map(|_| random_vector(&mut rng, message_length)) .collect::>(); let masks = random_vector(&mut rng, mask_length * num_messages); - let message_refs = messages.iter().map(|v| v.as_slice()).collect::>(); + let vectors = messages + .iter() + .map(|message| ActiveBuffer::from_slice(message)) + .collect::>(); + let vector_refs = vectors.iter().collect::>(); + let mask_buffer = ActiveBuffer::from_slice(&masks); + let rs_messages = Messages::new(&vector_refs, message_length, 1); let codeword = ntt.interleaved_encode( - &message_refs, - &masks, + rs_messages, + &mask_buffer, codeword_length, ); @@ -200,6 +244,7 @@ mod tests { // Output values are polynomial evaluations in the evaluation points. let mut evaluation_points = ntt.evaluation_points(message_length + mask_length, codeword_length, &sampled_indices); + 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]; let masks = chunks_exact_or_empty(&masks, mask_length, num_messages); diff --git a/src/bin/benchmark.rs b/src/bin/benchmark.rs index 77574c42..b3014086 100644 --- a/src/bin/benchmark.rs +++ b/src/bin/benchmark.rs @@ -15,6 +15,7 @@ use whir::{ linear_form::{Evaluate, LinearForm, MultilinearExtension}, }, bits::Bits, + buffer::{ActiveBuffer, BufferOps}, cmdline_utils::{AvailableFields, AvailableHash}, hash::HASH_COUNTER, parameters::ProtocolParameters, @@ -162,12 +163,13 @@ where HASH_COUNTER.reset(); - let witness = params.commit(&mut prover_state, &[&vector]); + let vector_buffer = ActiveBuffer::from_slice(&vector); + let witness = params.commit(&mut prover_state, &[&vector_buffer]); let _ = params.prove( &mut prover_state, - vec![Cow::Borrowed(vector.as_slice())], - vec![Cow::Owned(witness)], + &[&vector_buffer], + vec![&witness], vec![], Cow::Owned(vec![]), ); @@ -239,7 +241,8 @@ where HASH_COUNTER.reset(); let whir_prover_time = Instant::now(); - let witness = params.commit(&mut prover_state, &[&vector]); + let vector_buffer = ActiveBuffer::from_slice(&vector); + let witness = params.commit(&mut prover_state, &[&vector_buffer]); let prove_linear_forms: Vec>> = points .iter() @@ -250,8 +253,8 @@ where let _ = params.prove( &mut prover_state, - vec![Cow::Borrowed(vector.as_slice())], - vec![Cow::Owned(witness)], + &[&vector_buffer], + vec![&witness], prove_linear_forms, Cow::Borrowed(evaluations.as_slice()), ); diff --git a/src/bin/main.rs b/src/bin/main.rs index 0faf4b36..56f97487 100644 --- a/src/bin/main.rs +++ b/src/bin/main.rs @@ -11,6 +11,7 @@ use whir::{ linear_form::{Covector, Evaluate, LinearForm, MultilinearExtension}, }, bits::Bits, + buffer::{ActiveBuffer, BufferOps}, cmdline_utils::{AvailableFields, AvailableHash}, hash::HASH_COUNTER, parameters::ProtocolParameters, @@ -149,9 +150,10 @@ where } let vector = (0..num_coeffs).map(M::Source::from).collect::>(); + let vector_buffer = ActiveBuffer::from_slice(&vector); let whir_commit_time = Instant::now(); - let witness = params.commit(&mut prover_state, &[&vector]); + let witness = params.commit(&mut prover_state, &[&vector_buffer]); let whir_commit_time = whir_commit_time.elapsed(); // Allocate constraints @@ -184,8 +186,8 @@ where let whir_prove_time = Instant::now(); let _ = params.prove( &mut prover_state, - vec![Cow::Borrowed(vector.as_slice())], - vec![Cow::Owned(witness)], + &[&vector_buffer], + vec![&witness], prove_linear_forms, Cow::Borrowed(evaluations.as_slice()), ); @@ -315,13 +317,14 @@ where } let whir_commit_time = Instant::now(); - let witness = params.commit(&mut prover_state, &[vector.as_slice()]); + let vector_buffer = ActiveBuffer::from_slice(&vector); + let witness = params.commit(&mut prover_state, &[&vector_buffer]); let whir_commit_time = whir_commit_time.elapsed(); let whir_prove_time = Instant::now(); let _ = params.prove( &mut prover_state, - vec![Cow::Borrowed(&vector)], + &[&vector_buffer], witness, prove_linear_forms, Cow::Borrowed(&evaluations), diff --git a/src/buffer/cpu.rs b/src/buffer/cpu.rs new file mode 100644 index 00000000..85bff3f7 --- /dev/null +++ b/src/buffer/cpu.rs @@ -0,0 +1,261 @@ +//! In-memory CPU backend for the buffer abstraction. + +use std::{any::Any, mem}; + +use ark_ff::Field; +use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, Rng, RngCore}; + +use crate::{ + algebra::{ + embedding::Embedding, + linear_form::{Covector, LinearForm, UnivariateEvaluation}, + }, + buffer::{Buffer, BufferOps}, + engines::EngineId, + hash::{self, Hash}, + protocols::{ + matrix_commit::{hash_rows, Encodable}, + merkle_tree, + }, +}; + +#[derive( + Clone, + PartialEq, + Eq, + PartialOrd, + Ord, + Hash, + Debug, + Default, + serde::Serialize, + serde::Deserialize, +)] +pub struct CpuBuffer { + data: Vec, +} + +impl BufferOps for CpuBuffer { + type Nodes = CpuBuffer; + + fn from_vec(source: Vec) -> Self { + Self { data: source } + } + + fn from_slice(source: &[T]) -> Self { + Self { + data: source.to_vec(), + } + } + + fn to_slice(&self) -> &[T] { + &self.data + } + + fn len(&self) -> usize { + self.data.len() + } + + fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec { + let mut result = Vec::with_capacity(indices.len() * num_cols); + for i in indices { + result.extend_from_slice(&self.data[i * num_cols..(i + 1) * num_cols]); + } + result + } + + fn gather_at_indices(&self, indices: &[usize]) -> Vec { + indices.iter().map(|&i| self.data[i]).collect() + } + + fn merklize( + &self, + num_cols: usize, + leaf_hash: EngineId, + merkle: &merkle_tree::Config, + ) -> (Self::Nodes, Hash) + where + T: Encodable + Send + Sync, + { + assert_eq!(self.len(), num_cols * merkle.num_leaves); + let engine = hash::ENGINES + .retrieve(leaf_hash) + .expect("Failed to retrieve hash engine"); + #[cfg(feature = "tracing")] + tracing::Span::current().record("engine", engine.name().as_ref()); + + let mut leaves = vec![Hash::default(); merkle.num_leaves]; + hash_rows(&*engine, &self.data, &mut leaves); + let nodes = merkle.build_nodes(leaves); + let root = nodes[merkle.num_nodes() - 1]; + (CpuBuffer { data: nodes }, root) + } +} + +impl Buffer for CpuBuffer { + type TargetBuffer = CpuBuffer; + + fn zeros(length: usize) -> Self { + Self { + data: vec![F::ZERO; length], + } + } + + fn random(rng: &mut R, length: usize) -> Self + where + R: RngCore + CryptoRng, + Standard: Distribution, + { + Self { + data: (0..length).map(|_| rng.gen()).collect(), + } + } + + fn dot(&self, other: &Self) -> F { + crate::algebra::dot(&self.data, &other.data) + } + + fn sumcheck_polynomial(&self, other: &Self) -> (F, F) { + crate::algebra::sumcheck::compute_sumcheck_polynomial(&self.data, &other.data) + } + + fn fold(&mut self, weight: F) { + crate::algebra::sumcheck::fold(&mut self.data, weight); + } + + fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { + crate::algebra::sumcheck::fold_and_compute_polynomial( + &mut self.data, + &mut other.data, + weight, + ) + } + + fn scalar_mul(&mut self, weight: F) { + crate::algebra::scalar_mul(&mut self.data, weight); + } + + fn accumulate_univariate_evaluations( + &mut self, + evaluators: &[UnivariateEvaluation], + scalars: &[F], + ) { + let Some(size) = evaluators.first().map(|e| e.size) else { + return; + }; + UnivariateEvaluation::accumulate_many(evaluators, &mut self.data[..size], scalars); + } + + fn mixed_univariate_evaluate>( + &self, + embedding: &M, + point: M::Target, + ) -> M::Target { + crate::algebra::mixed_univariate_evaluate(embedding, &self.data, point) + } + + fn mixed_dot>( + &self, + embedding: &M, + other: &CpuBuffer, + ) -> M::Target { + crate::algebra::mixed_dot(embedding, &other.data, &self.data) + } + + fn mixed_scalar_mul_add_to>( + &self, + embedding: &M, + accumulator: &mut CpuBuffer, + weight: M::Target, + ) { + crate::algebra::mixed_scalar_mul_add(embedding, &mut accumulator.data, weight, &self.data); + } + + fn linear_forms_rlc( + size: usize, + linear_forms: &mut [Box>], + rlc_coeffs: &[F], + ) -> Self { + assert_eq!(linear_forms.len(), rlc_coeffs.len()); + let mut covector = vec![F::ZERO; size]; + if let Some((first, linear_forms)) = linear_forms.split_first_mut() { + debug_assert_eq!(rlc_coeffs[0], F::ONE); + if let Some(covector_form) = + (first.as_mut() as &mut dyn Any).downcast_mut::>() + { + mem::swap(&mut covector, &mut covector_form.vector); + } else { + first.accumulate(&mut covector, F::ONE); + } + for (rlc_coeff, linear_form) in rlc_coeffs[1..].iter().zip(linear_forms) { + linear_form.accumulate(&mut covector, *rlc_coeff); + } + } + Self { data: covector } + } + + fn mixed_linear_combination>( + embedding: &M, + vectors: &[&Self], + coeffs: &[M::Target], + ) -> CpuBuffer { + assert_eq!(vectors.len(), coeffs.len()); + let Some((first, vectors)) = vectors.split_first() else { + return CpuBuffer { data: Vec::new() }; + }; + debug_assert_eq!(coeffs[0], M::Target::ONE); + let mut accumulator = crate::algebra::lift(embedding, &first.data); + for (coeff, vector) in coeffs[1..].iter().zip(vectors) { + crate::algebra::mixed_scalar_mul_add(embedding, &mut accumulator, *coeff, &vector.data); + } + CpuBuffer { data: accumulator } + } +} + +#[cfg(test)] +mod tests { + use ark_ff::AdditiveGroup; + + use super::*; + use crate::algebra::{fields::Field64, geometric_accumulate}; + + type F = Field64; + + #[test] + fn scalar_mul_multiplies_in_place() { + let values = vec![F::from(1u64), F::from(2u64), F::from(3u64), F::from(4u64)]; + let weight = F::from(5u64); + let mut buffer = CpuBuffer::from_vec(values.clone()); + buffer.scalar_mul(weight); + let expected: Vec = values.iter().map(|&v| v * weight).collect(); + assert_eq!(buffer.to_slice(), expected.as_slice()); + } + + #[test] + fn accumulate_matches_geometric_accumulate_over_prefix() { + let len = 8usize; + let points = vec![F::from(3u64), F::from(5u64)]; + let scalars = vec![F::from(2u64), F::from(9u64)]; + + // Full-length and prefix accumulation. + for size in [len, 5] { + let mut buffer = CpuBuffer::from_vec(vec![F::ZERO; len]); + let evaluators: Vec<_> = points + .iter() + .map(|&point| UnivariateEvaluation::new(point, size)) + .collect(); + buffer.accumulate_univariate_evaluations(&evaluators, &scalars); + + // Reference: accumulate Σ_j scalars[j]·points[j]^i into the prefix + // of a plain vector. + let mut reference = vec![F::ZERO; len]; + geometric_accumulate(&mut reference[..size], scalars.clone(), &points); + + assert_eq!( + buffer.to_slice(), + reference.as_slice(), + "accumulate mismatch for evaluator size {size}" + ); + } + } +} diff --git a/src/buffer/mod.rs b/src/buffer/mod.rs new file mode 100644 index 00000000..9f320111 --- /dev/null +++ b/src/buffer/mod.rs @@ -0,0 +1,151 @@ +//! Backend-agnostic buffers for protocol data. +//! +//! Protocol code uses the [`ActiveBuffer`] alias to select the backend at +//! compile time. Buffers are owned, backend-managed storage: on the CPU +//! backend they wrap a `Vec`, on an accelerator backend they would own +//! device memory and only [`BufferOps::to_slice`] (and the other readback +//! methods) force a host copy. +//! +//! The trait split follows the element type. [`BufferOps`] is generic over +//! any element and also covers [`struct@Hash`] buffers for Merkle tree +//! nodes. [`Buffer`] adds field arithmetic used by the protocols. +//! +//! [`DefaultRs`] selects the Reed-Solomon encoder for the active backend. + +pub mod cpu; + +use ark_ff::Field; +use ark_std::rand::{ + distributions::{Distribution, Standard}, + CryptoRng, RngCore, +}; +pub use cpu::CpuBuffer; + +use crate::{ + algebra::{ + embedding::Embedding, + linear_form::{LinearForm, UnivariateEvaluation}, + }, + engines::EngineId, + hash::Hash, + protocols::{matrix_commit::Encodable, merkle_tree}, +}; + +pub type ActiveBuffer = CpuBuffer; +pub type DefaultRs = crate::algebra::ntt::NttEngine; + +/// Owned buffer operations over any element type. +/// +/// This trait is not field-specific, so it also covers hash buffers and +/// Merkle node layers. Field arithmetic lives on [`Buffer`]. +pub trait BufferOps { + /// Same-backend buffer type used for Merkle tree nodes. + type Nodes: BufferOps; + + fn from_vec(source: Vec) -> Self; + fn from_slice(source: &[T]) -> Self; + /// Read back the buffer contents as a host slice. + fn to_slice(&self) -> &[T]; + fn len(&self) -> usize; + fn is_empty(&self) -> bool { + self.len() == 0 + } + /// Gather full rows `indices[i] * num_cols .. (indices[i] + 1) * num_cols`. + fn read_rows(&self, num_cols: usize, indices: &[usize]) -> Vec; + /// Gather elements at arbitrary indices. + fn gather_at_indices(&self, indices: &[usize]) -> Vec; + /// Hash rows of width `num_cols` and build a Merkle tree. + fn merklize( + &self, + num_cols: usize, + leaf_hash: EngineId, + merkle: &merkle_tree::Config, + ) -> (Self::Nodes, Hash) + where + T: Encodable + Send + Sync; +} + +/// Field operations on owned buffers. +pub trait Buffer: BufferOps + Clone { + /// Same-backend owned buffer over another field. + /// + /// Used by the `mixed_*` operations that lift base-field data into an + /// extension field through an [`Embedding`]. + type TargetBuffer: Buffer; + + fn zeros(length: usize) -> Self; + + fn random(rng: &mut R, length: usize) -> Self + where + R: RngCore + CryptoRng, + Standard: Distribution; + + /// Inner product with another buffer of the same length. + fn dot(&self, other: &Self) -> F; + + /// Sumcheck round coefficients `(c0, c2)` for `dot(self, other)`. + fn sumcheck_polynomial(&self, other: &Self) -> (F, F); + + fn fold(&mut self, weight: F); + + fn fold_pair(&mut self, other: &mut Self, weight: F) { + self.fold(weight); + other.fold(weight); + } + + fn fold_pair_sumcheck_polynomial(&mut self, other: &mut Self, weight: F) -> (F, F) { + self.fold_pair(other, weight); + self.sumcheck_polynomial(other) + } + + /// 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: &[F], + ); + + /// Univariate evaluation at a target-field point. + fn mixed_univariate_evaluate>( + &self, + embedding: &M, + point: M::Target, + ) -> M::Target; + + /// Inner product with a target-field buffer. + fn mixed_dot>( + &self, + embedding: &M, + other: &Self::TargetBuffer, + ) -> M::Target; + + /// `accumulator += weight * self`, lifted into the target field. + fn mixed_scalar_mul_add_to>( + &self, + embedding: &M, + accumulator: &mut Self::TargetBuffer, + weight: M::Target, + ); + + /// Random linear combination of linear forms into a covector buffer. + fn linear_forms_rlc( + size: usize, + linear_forms: &mut [Box>], + rlc_coeffs: &[F], + ) -> Self; + + /// Random linear combination of vectors, lifted into the target field. + fn mixed_linear_combination>( + embedding: &M, + vectors: &[&Self], + coeffs: &[M::Target], + ) -> Self::TargetBuffer; +} diff --git a/src/lib.rs b/src/lib.rs index 7d1d8b45..4c1a5c3d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,6 +1,7 @@ pub mod algebra; pub mod ark_serde; pub mod bits; +pub mod buffer; pub mod cmdline_utils; pub mod engines; pub mod hash; diff --git a/src/protocols/basecase.rs b/src/protocols/basecase.rs index ec22a628..ce9f56e8 100644 --- a/src/protocols/basecase.rs +++ b/src/protocols/basecase.rs @@ -6,10 +6,8 @@ use serde::{Deserialize, Serialize}; use spongefish::{Decoding, VerificationResult}; use crate::{ - algebra::{ - dot, embedding::Identity, multilinear_extend, random_vector, scalar_mul_add_new, - univariate_evaluate, - }, + algebra::{embedding::Identity, multilinear_extend, univariate_evaluate}, + buffer::{ActiveBuffer, Buffer, BufferOps}, hash::Hash, protocols::{irs_commit, proof_of_work, sumcheck}, transcript::{ @@ -99,9 +97,9 @@ impl Config { pub fn prove( &self, prover_state: &mut ProverState, - mut vector: Vec, + mut vector: ActiveBuffer, witness: &irs_commit::Witness, - mut covector: Vec, + mut covector: ActiveBuffer, mut sum: F, ) -> Opening where @@ -118,7 +116,7 @@ impl Config { assert_eq!(self.commit.num_vectors(), 1); assert_eq!(self.commit.vector_size(), self.sumcheck.initial_size()); assert_eq!(self.sumcheck.final_size(), 1.min(self.commit.vector_size())); - debug_assert_eq!(dot(&vector, &covector), sum); + debug_assert_eq!(vector.dot(&covector), sum); if self.size() == 0 { return Opening { evaluation_points: Vec::new(), @@ -141,11 +139,14 @@ impl Config { // Negligible event over a challenge-sized field; without it the verifier // cannot derive `l(r) = sum / vector_mle(r)`. - assert!(!vector[0].is_zero(), "Proof failed"); + assert!( + !vector.to_slice().first().expect("Proof failed").is_zero(), + "Proof failed" + ); Opening { evaluation_points: point, - linear_form_evaluation: covector[0], + linear_form_evaluation: *covector.to_slice().first().expect("Proof failed"), } } @@ -155,9 +156,9 @@ impl Config { fn maybe_blind_prove( &self, prover_state: &mut ProverState, - vector: &mut Vec, + vector: &mut ActiveBuffer, witness: &irs_commit::Witness, - covector: &[F], + covector: &ActiveBuffer, sum: &mut F, ) -> Option> where @@ -171,14 +172,14 @@ impl Config { { match self.mode { BasecaseMode::Standard => { - prover_state.prover_messages(vector); - prover_state.prover_messages(&witness.masks); + prover_state.prover_messages(vector.to_slice()); + prover_state.prover_messages(witness.masks.to_slice()); None } BasecaseMode::ZeroKnowledge => { - let blinding_vector = random_vector(prover_state.rng(), vector.len()); + let mut blinding_vector = ActiveBuffer::random(prover_state.rng(), vector.len()); let blinding_witness = self.commit.commit(prover_state, &[&blinding_vector]); - let blinding_inner_product = dot(&blinding_vector, covector); + let blinding_inner_product = blinding_vector.dot(covector); prover_state.prover_message(&blinding_inner_product); // Grind the Theorem 7.1 γ-combination gap before γ is sampled. @@ -187,15 +188,21 @@ impl Config { let combination_randomness = prover_state.verifier_message::(); assert!(!combination_randomness.is_zero(), "Proof failed"); - *vector = scalar_mul_add_new(&blinding_vector, combination_randomness, vector); - prover_state.prover_messages(vector); + vector.mixed_scalar_mul_add_to( + &Identity::::new(), + &mut blinding_vector, + combination_randomness, + ); + *vector = blinding_vector; + prover_state.prover_messages(vector.to_slice()); - let combined_irs_randomness = scalar_mul_add_new( - &blinding_witness.masks, + let mut combined_irs_randomness = blinding_witness.masks.clone(); + witness.masks.mixed_scalar_mul_add_to( + &Identity::::new(), + &mut combined_irs_randomness, combination_randomness, - &witness.masks, ); - prover_state.prover_messages(&combined_irs_randomness); + prover_state.prover_messages(combined_irs_randomness.to_slice()); *sum = blinding_inner_product + combination_randomness * *sum; Some(blinding_witness) @@ -338,9 +345,9 @@ mod tests { .session(&format!("Test at {}:{}", file!(), line!())) .instance(&instance); let mut rng = StdRng::seed_from_u64(seed); - let vector = random_vector(&mut rng, config.size()); - let covector = random_vector(&mut rng, config.size()); - let sum = dot(&vector, &covector); + let vector = ActiveBuffer::random(&mut rng, config.size()); + let covector = ActiveBuffer::random(&mut rng, config.size()); + let sum = vector.dot(&covector); let mut prover_state = ProverState::new_std(&ds); let witness = config.commit.commit(&mut prover_state, &[&vector]); @@ -352,7 +359,7 @@ mod tests { sum, ); assert_eq!( - multilinear_extend(&covector, &prover_result.evaluation_points), + multilinear_extend(covector.to_slice(), &prover_result.evaluation_points), prover_result.linear_form_evaluation ); let proof = prover_state.proof(); diff --git a/src/protocols/code_switch.rs b/src/protocols/code_switch.rs index cfaf67b1..e0db0460 100644 --- a/src/protocols/code_switch.rs +++ b/src/protocols/code_switch.rs @@ -15,8 +15,11 @@ use crate::{ algebra::{ dot, embedding::{Embedding, Identity}, - eq_weights, geometric_accumulate, lift, mixed_dot, scalar_mul, univariate_evaluate, + eq_weights, lift, + linear_form::UnivariateEvaluation, + mixed_dot, }, + buffer::{ActiveBuffer, Buffer, BufferOps}, hash::Hash, protocols::{ geometric_challenge::geometric_challenge, @@ -53,13 +56,21 @@ pub struct Config { #[must_use] #[derive(Clone, Debug)] pub struct Witness { - pub message: Vec, + pub message: ActiveBuffer, pub target_witness: IrsWitness, } /// Verifier output from the code-switch. pub type Commitment = IrsCommitment; +#[inline] +fn univariate_evaluators(points: &[F], size: usize) -> Vec> { + points + .iter() + .map(|&point| UnivariateEvaluation::new(point, size)) + .collect() +} + impl Config { /// Create a code-switch config. pub fn new( @@ -210,11 +221,11 @@ impl Config { pub fn prove( &self, prover_state: &mut ProverState, - message: Vec, + message: ActiveBuffer, witness: &IrsWitness, - covector: &mut [M::Target], + covector: &mut ActiveBuffer, folding_randomness: &[M::Target], - mask: &[M::Target], + mask: &ActiveBuffer, ) -> Witness where H: DuplexSpongeInterface, @@ -236,7 +247,6 @@ impl Config { folding_randomness.len(), self.source.interleaving_depth(), ); - // Step 1: g := Enc_{C'}(f, r') — Construction 9.7 Step 1, p.55 let target_witness = self.target.commit(prover_state, &[&message]); @@ -258,9 +268,9 @@ impl Config { 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); - // Covector update — sl' from Completeness proof (p.55-56) + // Covector update — sl' from Completeness proof (p.55-56). + covector.scalar_mul(original_sl_coeff); let eval_points = lift(self.source.embedding(), &source_evaluations.points); - scalar_mul(covector, original_sl_coeff); self.update_covector( covector, ood_rlc_coeffs, @@ -280,8 +290,8 @@ impl Config { fn maybe_send_ood_answers( &self, prover_state: &mut ProverState, - message: &[M::Target], - mask: &[M::Target], + message: &ActiveBuffer, + mask: &ActiveBuffer, ood_points: &[M::Target], ) where H: DuplexSpongeInterface, @@ -290,11 +300,12 @@ impl Config { { let msg_len = message.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 { .. } => { - let mask_eval = univariate_evaluate(mask, point); + let mask_eval = + mask.mixed_univariate_evaluate(&Identity::::new(), point); let shift = point.pow([msg_len as u64]); f_eval + shift * mask_eval } @@ -308,7 +319,7 @@ impl Config { /// the full `[f; r; s]` and in-domain over the `[f; r]` prefix only. fn update_covector( &self, - covector: &mut [M::Target], + covector: &mut ActiveBuffer, ood_rlc_coeffs: &[M::Target], ood_points: &[M::Target], in_domain_rlc_coeffs: &[M::Target], @@ -316,22 +327,22 @@ impl Config { ) { match &self.mode { CodeSwitchMode::Standard => { - let all_points: Vec<_> = - ood_points.iter().chain(in_domain_points).copied().collect(); - let pows: Vec<_> = ood_rlc_coeffs + let mut evaluators = univariate_evaluators(ood_points, covector.len()); + evaluators.extend(univariate_evaluators(in_domain_points, covector.len())); + let scalars = ood_rlc_coeffs .iter() .chain(in_domain_rlc_coeffs) .copied() - .collect(); - geometric_accumulate(covector, pows, &all_points); + .collect::>(); + covector.accumulate_univariate_evaluations(&evaluators, &scalars); } 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(), - in_domain_points, - ); + let ood_evaluators = univariate_evaluators(ood_points, covector.len()); + covector.accumulate_univariate_evaluations(&ood_evaluators, ood_rlc_coeffs); + let in_domain_evaluators = + univariate_evaluators(in_domain_points, self.source.masked_message_length()); + covector + .accumulate_univariate_evaluations(&in_domain_evaluators, in_domain_rlc_coeffs); } } } @@ -410,8 +421,8 @@ impl Config { let num_ood = self.out_domain_samples; let num_in_domain = source_evaluations.points.len(); let coeffs = geometric_challenge(verifier_state, 1 + num_ood + num_in_domain); - let (&original_sl_coeff, all_rlc_coeffs) = coeffs.split_first().unwrap(); - let (ood_rlc_coeffs, in_domain_rlc_coeffs) = all_rlc_coeffs.split_at(num_ood); + let (&original_sl_coeff, constraint_rlc_coeffs) = coeffs.split_first().unwrap(); + let (ood_rlc_coeffs, in_domain_rlc_coeffs) = constraint_rlc_coeffs.split_at(num_ood); *sum = original_sl_coeff * *sum + dot(ood_rlc_coeffs, &ood_answers) @@ -584,7 +595,7 @@ mod tests { } // Lift ι parallel masks (total length source.mask_length × ι) and fold // chunks of length source.mask_length down to a single chunk. - let raw = lift(config.source.embedding(), &source_witness.masks); + let raw = lift(config.source.embedding(), source_witness.masks.to_slice()); let mut mask = fold_chunks(&raw, config.source.mask_length(), folding_randomness); // Append fresh padding s of length message_mask_length - source.mask_length. mask.extend(random_vector::( @@ -607,19 +618,24 @@ mod tests { let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); covector.resize(config.covector_length(), F::ZERO); + let mut covector = ActiveBuffer::from_vec(covector); let instance = U64(seed); let ds = DomainSeparator::protocol(config) .session(&format!("Test at {}:{}", file!(), line!())) .instance(&instance); let mut prover_state = ProverState::new_std(&ds); - let source_witness = config.source.commit(&mut prover_state, &[&f_full]); + let f_full_buffer = ActiveBuffer::from_slice(&f_full); + let source_witness = config.source.commit(&mut prover_state, &[&f_full_buffer]); // Sample γ for sumcheck folding (length log2(ι)). let folding_randomness = sample_folding_randomness(config, &mut rng); // Post-fold message Fold(f_full, γ) of length message_length. - let folded_message = - fold_chunks(&f_full, config.source.message_length(), &folding_randomness); + let folded_message = ActiveBuffer::from_vec(fold_chunks( + &f_full, + config.source.message_length(), + &folding_randomness, + )); let mask_msg = build_mask_msg(config, &source_witness, &folding_randomness, &mut rng); let witness = config.prove( @@ -628,7 +644,7 @@ mod tests { &source_witness, &mut covector, &folding_randomness, - &mask_msg, + &ActiveBuffer::from_slice(&mask_msg), ); let proof = prover_state.proof(); @@ -660,33 +676,41 @@ mod tests { let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); covector.resize(config.covector_length(), F::ZERO); + let mut covector = ActiveBuffer::from_vec(covector); let instance = U64(seed); let ds = DomainSeparator::protocol(config) .session(&format!("Test at {}:{}", file!(), line!())) .instance(&instance); let mut prover_state = ProverState::new_std(&ds); - let source_witness = config.source.commit(&mut prover_state, &[&f_full]); + let f_full_buffer = ActiveBuffer::from_slice(&f_full); + let source_witness = config.source.commit(&mut prover_state, &[&f_full_buffer]); let folding_randomness = sample_folding_randomness(config, &mut rng); - let folded_message = - fold_chunks(&f_full, config.source.message_length(), &folding_randomness); + let folded_message = ActiveBuffer::from_vec(fold_chunks( + &f_full, + config.source.message_length(), + &folding_randomness, + )); let mask_msg = build_mask_msg(config, &source_witness, &folding_randomness, &mut rng); // h is the post-fold polynomial whose inner product with covector // should equal the verifier sum: // - non-ZK: h = folded_message (length message_length) // - ZK: h = [folded_message; mask_msg] (length message_length + l_zk) - let h: Vec = if mask_msg.is_empty() { + let h: ActiveBuffer = if mask_msg.is_empty() { folded_message.clone() } else { - folded_message - .iter() - .chain(mask_msg.iter()) - .copied() - .collect() + ActiveBuffer::from_vec( + folded_message + .to_slice() + .iter() + .chain(mask_msg.iter()) + .copied() + .collect(), + ) }; - let initial_mu = dot(&h, &covector); + let initial_mu = h.dot(&covector); let _witness = config.prove( &mut prover_state, @@ -694,7 +718,7 @@ mod tests { &source_witness, &mut covector, &folding_randomness, - &mask_msg, + &ActiveBuffer::from_slice(&mask_msg), ); let proof = prover_state.proof(); @@ -714,7 +738,7 @@ mod tests { .unwrap(); verifier_state.check_eof().unwrap(); - assert_eq!(dot(&h, &covector), verifier_sum); + assert_eq!(h.dot(&covector), verifier_sum); } fn test_tampered_ood_config>(seed: u64, config: &Config>) @@ -731,27 +755,31 @@ mod tests { let mut covector: Vec = random_vector(&mut rng, config.source.message_length()); covector.resize(config.covector_length(), F::ZERO); + let mut covector = ActiveBuffer::from_vec(covector); // Commit honest f_full, fold to get the honest post-fold message. let mut prover_state = ProverState::new_std(&ds); - let source_witness = config.source.commit(&mut prover_state, &[&f_full]); + let f_full_buffer = ActiveBuffer::from_slice(&f_full); + let source_witness = config.source.commit(&mut prover_state, &[&f_full_buffer]); let folding_randomness = sample_folding_randomness(config, &mut rng); let folded_message = fold_chunks(&f_full, config.source.message_length(), &folding_randomness); // For non-ZK and source.mask_length == 0, h = folded_message and identity holds. - let initial_mu = dot(&folded_message, &covector); + let folded_message_buffer = ActiveBuffer::from_slice(&folded_message); + let initial_mu = folded_message_buffer.dot(&covector); // Tamper the post-fold message before proving. - let mut tampered = folded_message.clone(); + let mut tampered = folded_message; tampered[0] += F::ONE; + let tampered = ActiveBuffer::from_vec(tampered); let _witness = config.prove( &mut prover_state, tampered, &source_witness, &mut covector, &folding_randomness, - &[], + &ActiveBuffer::from_vec(Vec::new()), ); let proof = prover_state.proof(); @@ -772,7 +800,7 @@ mod tests { verifier_state.check_eof().unwrap(); // Sum diverges — downstream sumcheck would reject - assert_ne!(dot(&folded_message, &covector), verifier_sum); + assert_ne!(folded_message_buffer.dot(&covector), verifier_sum); } fn test + 'static>() diff --git a/src/protocols/irs_commit.rs b/src/protocols/irs_commit.rs index 2d0141c1..c33ca6c4 100644 --- a/src/protocols/irs_commit.rs +++ b/src/protocols/irs_commit.rs @@ -23,8 +23,9 @@ use tracing::instrument; use crate::{ algebra::{ dot, embedding::Embedding, fields::FieldWithSize, lift, linear_form::UnivariateEvaluation, - mixed_univariate_evaluate, ntt, random_vector, + ntt, }, + buffer::{ActiveBuffer, Buffer, BufferOps}, engines::EngineId, hash::Hash, protocols::{ @@ -37,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)] @@ -99,8 +100,8 @@ pub struct Config { #[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Debug, Hash, Default, Serialize, Deserialize)] #[must_use] pub struct Witness { - pub masks: Vec, - pub matrix: Vec, + pub masks: ActiveBuffer, + pub matrix: ActiveBuffer, pub matrix_witness: matrix_commit::Witness, } @@ -336,7 +337,7 @@ impl Config { pub fn commit( &self, prover_state: &mut ProverState, - vectors: &[&[M::Source]], + vectors: &[&ActiveBuffer], ) -> Witness where Standard: Distribution, @@ -354,15 +355,12 @@ impl Config { assert_eq!(vectors.len(), self.num_vectors); assert!(vectors.iter().all(|p| p.len() == self.vector_size)); - // Generate random mask - let masks = random_vector(prover_state.rng(), self.mask_length() * self.num_messages()); - - // Interleaved RS Encode the vectors - let messages = vectors - .iter() - .flat_map(|v| chunks_exact_or_empty(v, self.message_length(), self.interleaving_depth)) - .collect::>(); - let matrix = ntt::interleaved_rs_encode(&messages, &masks, self.codeword_length); + let masks = ActiveBuffer::::random( + prover_state.rng(), + self.mask_length() * self.num_messages(), + ); + 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); @@ -400,7 +398,7 @@ impl Config { pub fn commit_with_ood( &self, prover_state: &mut ProverState, - vectors: &[&[M::Source]], + vectors: &[&ActiveBuffer], out_domain_samples: usize, ) -> (Witness, Evaluations) where @@ -415,7 +413,7 @@ impl Config { let mut matrix = Vec::with_capacity(out_domain_samples * vectors.len()); for &point in &points { for &vector in vectors { - let value = mixed_univariate_evaluate(&*self.embedding, vector, point); + let value = vector.mixed_univariate_evaluate(&*self.embedding, point); prover_state.prover_message(&value); matrix.push(value); } @@ -472,18 +470,15 @@ impl Config { // and collect them in the evaluation matrix. let stride = witnesses.len() * self.num_cols(); let mut matrix = vec![M::Source::ZERO; indices.len() * stride]; - let mut submatrix = Vec::with_capacity(indices.len() * self.num_cols()); let mut matrix_col_offset = 0; for witness in witnesses { - submatrix.clear(); - for (point_index, &code_index) in indices.iter().enumerate() { - let row = &witness.matrix - [code_index * self.num_cols()..(code_index + 1) * self.num_cols()]; - submatrix.extend_from_slice(row); - - let matrix_row = &mut matrix[point_index * stride..(point_index + 1) * stride]; - matrix_row[matrix_col_offset..matrix_col_offset + self.num_cols()] - .copy_from_slice(row); + let submatrix = witness.matrix.read_rows(self.num_cols(), &indices); + if self.num_cols() != 0 { + for (point_index, row) in submatrix.chunks_exact(self.num_cols()).enumerate() { + let matrix_row = &mut matrix[point_index * stride..(point_index + 1) * stride]; + matrix_row[matrix_col_offset..matrix_col_offset + self.num_cols()] + .copy_from_slice(row); + } } prover_state.prover_hint_ark(&submatrix); self.matrix_commit @@ -752,10 +747,12 @@ pub(crate) mod tests { // Prover let mut prover_state = ProverState::new_std(&ds); - let witness = config.commit( - &mut prover_state, - &vectors.iter().map(|p| p.as_slice()).collect::>(), - ); + let vector_buffers = vectors + .iter() + .map(|v| ActiveBuffer::from_slice(v)) + .collect::>(); + let vector_refs = vector_buffers.iter().collect::>(); + let witness = config.commit(&mut prover_state, &vector_refs); let in_domain_evals = config.open(&mut prover_state, &[&witness]); if config.deduplicate_in_domain { // Sorting is over index order, not points diff --git a/src/protocols/mask_proximity.rs b/src/protocols/mask_proximity.rs index 5f285872..3ed33a33 100644 --- a/src/protocols/mask_proximity.rs +++ b/src/protocols/mask_proximity.rs @@ -45,7 +45,8 @@ use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, R use serde::{Deserialize, Serialize}; use crate::{ - algebra::{embedding::Identity, random_vector, scalar_mul_add_new, univariate_evaluate}, + algebra::{embedding::Identity, scalar_mul_add, univariate_evaluate}, + buffer::{ActiveBuffer, Buffer, BufferOps}, hash::Hash, protocols::{ irs_commit::{Commitment as IrsCommitment, Config as IrsConfig, Witness as IrsWitness}, @@ -75,7 +76,7 @@ pub struct Config { #[must_use] pub struct Witness { pub(crate) mask_witness: IrsWitness, - pub(crate) fresh_msgs: Vec>, + pub(crate) fresh_msgs: Vec>, } /// Verifier output from the commit phase. @@ -129,7 +130,7 @@ impl Config { pub fn commit( &self, prover_state: &mut ProverState, - original_msgs: &[Vec], + original_msgs: &[&ActiveBuffer], ) -> Witness where F: Codec<[H::U]>, @@ -144,15 +145,15 @@ impl Config { } // Sample fresh mask-of-masks - let fresh_msgs: Vec> = (0..self.num_masks) - .map(|_| random_vector(prover_state.rng(), self.c_zk_commit.vector_size())) + let fresh_msgs: Vec> = (0..self.num_masks) + .map(|_| ActiveBuffer::random(prover_state.rng(), self.c_zk_commit.vector_size())) .collect(); // Tree layout: [originals..., freshes...] - let all_vectors: Vec<&[F]> = original_msgs + let all_vectors: Vec<&ActiveBuffer> = original_msgs .iter() + .copied() .chain(fresh_msgs.iter()) - .map(|v| v.as_slice()) .collect(); let mask_witness = self.c_zk_commit.commit(prover_state, &all_vectors); @@ -181,7 +182,7 @@ impl Config { &self, prover_state: &mut ProverState, witness: &Witness, - original_msgs: &[Vec], + original_msgs: &[&ActiveBuffer], ) where F: Codec<[H::U]>, H: DuplexSpongeInterface, @@ -204,27 +205,32 @@ impl Config { // Step 2: compute and send combined polynomials + IRS randomness let irs_masks_per_vector = self.c_zk_commit.mask_length() * self.c_zk_commit.interleaving_depth(); - assert_eq!( - witness.mask_witness.masks.len(), - 2 * self.num_masks * irs_masks_per_vector - ); + let irs_masks = &witness.mask_witness.masks; + assert_eq!(irs_masks.len(), 2 * self.num_masks * irs_masks_per_vector); for (i, (orig_msg, fresh_msg)) in original_msgs .iter() + .copied() .zip(witness.fresh_msgs.iter()) .enumerate() { // ξ*_i = s_i + γ · ξ_i - let combined_msg = scalar_mul_add_new(fresh_msg, gamma, orig_msg); - prover_state.prover_messages(&combined_msg); + let mut combined_msg = fresh_msg.clone(); + orig_msg.mixed_scalar_mul_add_to(&Identity::::new(), &mut combined_msg, gamma); + prover_state.prover_messages(combined_msg.to_slice()); // r*_i = r'_i + γ · r_i + // Combined on the host: the result is transcript data anyway, and + // the IRS randomness is only a few elements per vector. if irs_masks_per_vector > 0 { - let orig_r = &witness.mask_witness.masks - [i * irs_masks_per_vector..(i + 1) * irs_masks_per_vector]; - let fresh_r = &witness.mask_witness.masks[(self.num_masks + i) - * irs_masks_per_vector - ..(self.num_masks + i + 1) * irs_masks_per_vector]; - let combined_r = scalar_mul_add_new(fresh_r, gamma, orig_r); + let base = i * irs_masks_per_vector; + let fresh_base = (self.num_masks + i) * irs_masks_per_vector; + let masks = irs_masks.to_slice(); + let mut combined_r = masks[fresh_base..fresh_base + irs_masks_per_vector].to_vec(); + scalar_mul_add( + &mut combined_r, + gamma, + &masks[base..base + irs_masks_per_vector], + ); prover_state.prover_messages(&combined_r); } } @@ -370,10 +376,15 @@ mod tests { let original_msgs: Vec> = (0..config.num_masks) .map(|_| random_vector(&mut rng, config.c_zk_commit.vector_size())) .collect(); + let original_buffers = original_msgs + .iter() + .map(|msg| ActiveBuffer::from_slice(msg)) + .collect::>(); + let original_refs = original_buffers.iter().collect::>(); let mut prover_state = ProverState::new_std(&ds); - let witness = config.commit(&mut prover_state, &original_msgs); - config.prove(&mut prover_state, &witness, &original_msgs); + let witness = config.commit(&mut prover_state, &original_refs); + config.prove(&mut prover_state, &witness, &original_refs); let proof = prover_state.proof(); let mut verifier_state = VerifierState::new_std(&ds, &proof); @@ -444,13 +455,23 @@ mod tests { let original_msgs: Vec> = (0..config.num_masks) .map(|_| random_vector(&mut rng, config.c_zk_commit.vector_size())) .collect(); + let original_buffers = original_msgs + .iter() + .map(|msg| ActiveBuffer::from_slice(msg)) + .collect::>(); + let original_refs = original_buffers.iter().collect::>(); let mut prover_state = ProverState::new_std(&ds); - let witness = config.commit(&mut prover_state, &original_msgs); + let witness = config.commit(&mut prover_state, &original_refs); let mut tampered_msgs = original_msgs; tampered_msgs[0][0] += F::ONE; - config.prove(&mut prover_state, &witness, &tampered_msgs); + let tampered_buffers = tampered_msgs + .iter() + .map(|msg| ActiveBuffer::from_slice(msg)) + .collect::>(); + let tampered_refs = tampered_buffers.iter().collect::>(); + config.prove(&mut prover_state, &witness, &tampered_refs); let proof = prover_state.proof(); let mut verifier_state = VerifierState::new_std(&ds, &proof); @@ -475,32 +496,45 @@ mod tests { let original_msgs: Vec> = (0..config.num_masks) .map(|_| random_vector(&mut rng, config.c_zk_commit.vector_size())) .collect(); + let original_buffers = original_msgs + .iter() + .map(|msg| ActiveBuffer::from_slice(msg)) + .collect::>(); + let original_refs = original_buffers.iter().collect::>(); let mut prover_state = ProverState::new_std(&ds); - let witness = config.commit(&mut prover_state, &original_msgs); + let witness = config.commit(&mut prover_state, &original_refs); let gamma: F = prover_state.verifier_message(); let irs_masks_per_vector = config.c_zk_commit.mask_length() * config.c_zk_commit.interleaving_depth(); + let irs_masks = &witness.mask_witness.masks; - for (i, (orig_msg, fresh_msg)) in original_msgs + for (i, (orig_msg, fresh_msg)) in original_refs .iter() + .copied() .zip(witness.fresh_msgs.iter()) .enumerate() { - let mut combined_msg = scalar_mul_add_new(fresh_msg, gamma, orig_msg); + let mut combined_msg = fresh_msg.clone(); + orig_msg.mixed_scalar_mul_add_to(&Identity::::new(), &mut combined_msg, gamma); if i == 0 { - combined_msg[0] += F::ONE; + let mut tampered = combined_msg.to_slice().to_vec(); + tampered[0] += F::ONE; + combined_msg = ActiveBuffer::from_slice(&tampered); } - prover_state.prover_messages(&combined_msg); + prover_state.prover_messages(combined_msg.to_slice()); if irs_masks_per_vector > 0 { - let orig_r = &witness.mask_witness.masks - [i * irs_masks_per_vector..(i + 1) * irs_masks_per_vector]; - let fresh_r = &witness.mask_witness.masks[(config.num_masks + i) - * irs_masks_per_vector - ..(config.num_masks + i + 1) * irs_masks_per_vector]; - let combined_r = scalar_mul_add_new(fresh_r, gamma, orig_r); + let base = i * irs_masks_per_vector; + let fresh_base = (config.num_masks + i) * irs_masks_per_vector; + let masks = irs_masks.to_slice(); + let mut combined_r = masks[fresh_base..fresh_base + irs_masks_per_vector].to_vec(); + scalar_mul_add( + &mut combined_r, + gamma, + &masks[base..base + irs_masks_per_vector], + ); prover_state.prover_messages(&combined_r); } } diff --git a/src/protocols/matrix_commit.rs b/src/protocols/matrix_commit.rs index 1e823162..8ccef4f4 100644 --- a/src/protocols/matrix_commit.rs +++ b/src/protocols/matrix_commit.rs @@ -12,6 +12,7 @@ use tracing::instrument; use zerocopy::{Immutable, IntoBytes}; use crate::{ + buffer::{ActiveBuffer, BufferOps}, engines::EngineId, hash::{self, Hash}, protocols::merkle_tree, @@ -168,7 +169,7 @@ impl Encoder for ZeroCopyEncoder { } } -impl Config { +impl Config { /// Create a new matrix commit configuration with the recommended hash function. pub fn new(num_rows: usize, num_cols: usize) -> Self { // Select a leaf hash function. @@ -208,7 +209,11 @@ impl Config { /// Commit the matrix (in row-major order). #[cfg_attr(feature = "tracing", instrument(skip_all, fields(self = %self, size = matrix.len(), engine)))] - pub fn commit(&self, prover_state: &mut ProverState, matrix: &[T]) -> Witness + pub fn commit( + &self, + prover_state: &mut ProverState, + matrix: &ActiveBuffer, + ) -> Witness where H: DuplexSpongeInterface, R: RngCore + CryptoRng, @@ -216,19 +221,9 @@ impl Config { { assert_eq!(matrix.len(), self.num_rows() * self.num_cols); - let engine = hash::ENGINES - .retrieve(self.leaf_hash_id) - .expect("Failed to retrieve hash engine"); - #[cfg(feature = "tracing")] - tracing::Span::current().record("engine", engine.name().as_ref()); - - // Compute leaf hashes - let mut leaves = Vec::with_capacity(self.merkle_tree.num_nodes()); - leaves.resize(self.merkle_tree.num_leaves, Hash::default()); - hash_rows(&*engine, matrix, &mut leaves[..self.num_rows()]); - - // Commit the leaf hashes - self.merkle_tree.commit(prover_state, leaves) + let (nodes, root) = matrix.merklize(self.num_cols, self.leaf_hash_id, &self.merkle_tree); + prover_state.prover_message(&root); + merkle_tree::Witness::new(nodes) } #[cfg_attr(feature = "tracing", instrument(skip_all, fields(self = %self)))] @@ -292,14 +287,14 @@ impl Config { } } -impl fmt::Display for Config { +impl fmt::Display for Config { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { write!(f, "MatrixCommit({} x {})", self.num_rows(), self.num_cols) } } #[cfg(not(feature = "parallel"))] -fn hash_rows( +pub fn hash_rows( engine: &dyn hash::HashEngine, matrix: &[T], out: &mut [Hash], @@ -308,7 +303,7 @@ fn hash_rows( } #[cfg(feature = "parallel")] -fn hash_rows( +pub fn hash_rows( engine: &dyn hash::HashEngine, matrix: &[T], out: &mut [Hash], @@ -406,7 +401,7 @@ pub(crate) mod tests { num_cols: usize, indices: &[usize], ) where - T: Clone + TypeInfo + Encodable + Send + Sync, + T: Copy + TypeInfo + Encodable + Send + Sync, Standard: Distribution, { crate::tests::init(); @@ -428,22 +423,9 @@ pub(crate) mod tests { .instance(&Empty); // Instance - let matrix: Vec = (0..config.size()).map(|_| rng.gen()).collect(); - let submatrix: Vec = if num_cols > 0 { - indices - .iter() - .flat_map(|&index| { - matrix - .chunks_exact(num_cols) - .nth(index) - .unwrap() - .iter() - .cloned() - }) - .collect::>() - } else { - Vec::new() - }; + let matrix: ActiveBuffer = + ActiveBuffer::from_vec((0..config.size()).map(|_| rng.gen()).collect()); + let submatrix: Vec = matrix.read_rows(num_cols, indices); // Prover let mut prover_state = ProverState::new_std(&ds); @@ -462,7 +444,7 @@ pub(crate) mod tests { fn proptest() where - T: Clone + TypeInfo + Encodable + Send + Sync, + T: Copy + TypeInfo + Encodable + Send + Sync, Standard: Distribution, { let hashes = [hash::COPY, hash::SHA2, hash::SHA3, hash::BLAKE3]; diff --git a/src/protocols/merkle_tree.rs b/src/protocols/merkle_tree.rs index fcf44ba0..dc18e22d 100644 --- a/src/protocols/merkle_tree.rs +++ b/src/protocols/merkle_tree.rs @@ -12,6 +12,7 @@ use tracing::{instrument, span, Level}; use zerocopy::IntoBytes; use crate::{ + buffer::{ActiveBuffer, BufferOps}, engines::EngineId, hash::{self, Hash, HashEngine, ENGINES}, transcript::{ @@ -58,7 +59,7 @@ pub struct Commitment { #[must_use] pub struct Witness { /// The nodes in the Merkle tree, starting with the leaf hash layer. - nodes: Vec, + nodes: ActiveBuffer, } impl Config { @@ -78,13 +79,19 @@ impl Config { (1 << (self.layers.len() + 1)) - 1 } - #[cfg_attr(feature = "tracing", instrument(skip(prover_state, leaves), fields(self = %self)))] - pub fn commit(&self, prover_state: &mut ProverState, leaves: Vec) -> Witness - where - H: DuplexSpongeInterface, - R: RngCore + CryptoRng, - Hash: ProverMessage<[H::U]>, - { + /// Build the full node array from leaf hashes, without touching the transcript. + /// + /// The returned vector holds the leaf layer first and the root last (at `num_nodes() - 1`). + /// + /// # Caller obligation + /// + /// This does **not** commit to anything: the caller must send the root + /// (`nodes[num_nodes() - 1]`) to the transcript before using the tree as + /// a commitment. In production this is funneled through + /// [`matrix_commit::Config::commit`](crate::protocols::matrix_commit::Config::commit), + /// which sends the root returned by `merklize`. A tree whose root was + /// never sent is not bound by the proof. + pub fn build_nodes(&self, leaves: Vec) -> Vec { assert_eq!( leaves.len(), self.num_leaves, @@ -120,10 +127,7 @@ impl Config { remaining = next_remaining; } - // Commit to the root hash. - prover_state.prover_message(&previous[0]); - - Witness { nodes } + nodes } pub fn receive_commitment( @@ -155,33 +159,9 @@ impl Config { assert_eq!(witness.nodes.len(), self.num_nodes()); assert!(indices.iter().all(|&i| i < self.num_leaves)); - // Abstract execution of verify algorithm writing required hashes. - let mut indices = indices.to_vec(); - indices.sort_unstable(); - indices.dedup(); - let (mut layer, mut remaining) = witness.nodes.split_at(1 << self.layers.len()); - while layer.len() > 1 { - let mut next_indices = Vec::with_capacity(indices.len()); - let mut iter = indices.iter().copied().peekable(); - loop { - match (iter.next(), iter.peek()) { - (Some(a), Some(&b)) if b == a ^ 1 => { - // Neighboring indices, merging branches. - next_indices.push(a >> 1); - iter.next(); // Skip the next index. - } - (Some(a), _) => { - // Single index, pushing the neighbor hash. - prover_state.prover_hint(&layer[a ^ 1]); - next_indices.push(a >> 1); - } - (None, _) => break, - } - } - indices = next_indices; - let (next_layer, next_remaining) = remaining.split_at(layer.len() / 2); - layer = next_layer; - remaining = next_remaining; + let node_indices = opening_sibling_indices(self.num_leaves, self.layers.len(), indices); + for hint in witness.nodes.gather_at_indices(&node_indices) { + prover_state.prover_hint(&hint); } } @@ -273,7 +253,12 @@ impl Config { } impl Witness { - pub const fn num_nodes(&self) -> usize { + /// Wrap a fully built node buffer (leaf layer first, root last) as a witness. + pub const fn new(nodes: ActiveBuffer) -> Self { + Self { nodes } + } + + pub fn num_nodes(&self) -> usize { self.nodes.len() } } @@ -303,6 +288,40 @@ fn parallel_hash(engine: &dyn HashEngine, size: usize, input: &[u8], output: &mu } } +/// Flat node indices of sibling hashes required to open at `indices`. +pub fn opening_sibling_indices(num_leaves: usize, layers: usize, indices: &[usize]) -> Vec { + debug_assert!(indices.iter().all(|&i| i < num_leaves)); + + let mut indices = indices.to_vec(); + indices.sort_unstable(); + indices.dedup(); + + let mut node_indices = Vec::new(); + let mut layer_offset = 0usize; + let mut layer_len = 1usize << layers; + while layer_len > 1 { + let mut next_indices = Vec::with_capacity(indices.len()); + let mut iter = indices.iter().copied().peekable(); + loop { + match (iter.next(), iter.peek()) { + (Some(a), Some(&b)) if b == a ^ 1 => { + next_indices.push(a >> 1); + iter.next(); + } + (Some(a), _) => { + node_indices.push(layer_offset + (a ^ 1)); + next_indices.push(a >> 1); + } + (None, _) => break, + } + } + indices = next_indices; + layer_offset += layer_len; + layer_len /= 2; + } + node_indices +} + #[cfg(test)] pub(crate) mod tests { use proptest::{collection::vec, prelude::Strategy}; @@ -340,8 +359,10 @@ pub(crate) mod tests { // Prover let mut prover_state = ProverState::new_std(&ds); - let tree = config.commit(&mut prover_state, leaves); - config.open(&mut prover_state, &tree, &[13, 42]); + let nodes = config.build_nodes(leaves); + prover_state.prover_message(&nodes[config.num_nodes() - 1]); + let witness = Witness::new(ActiveBuffer::from_vec(nodes)); + config.open(&mut prover_state, &witness, &[13, 42]); let proof = prover_state.proof(); // Verifier diff --git a/src/protocols/sumcheck.rs b/src/protocols/sumcheck.rs index 06e2b2e0..41b19161 100644 --- a/src/protocols/sumcheck.rs +++ b/src/protocols/sumcheck.rs @@ -9,11 +9,8 @@ use serde::{Deserialize, Serialize}; use tracing::instrument; use crate::{ - algebra::{ - dot, - sumcheck::{compute_sumcheck_polynomial, fold, fold_and_compute_polynomial}, - univariate_evaluate, - }, + algebra::univariate_evaluate, + buffer::{ActiveBuffer, Buffer, BufferOps}, protocols::proof_of_work, transcript::{ codecs::U64, Codec, Decoding, DuplexSpongeInterface, ProverState, VerificationResult, @@ -151,8 +148,8 @@ impl Config { pub fn prove( &self, prover_state: &mut ProverState, - a: &mut Vec, - b: &mut Vec, + a: &mut ActiveBuffer, + b: &mut ActiveBuffer, sum: &mut F, masks: &[F], ) -> SumcheckOpening @@ -168,9 +165,8 @@ impl Config { ); assert_eq!(a.len(), self.initial_size); assert_eq!(b.len(), self.initial_size); - debug_assert_eq!(dot(a, b), *sum); + debug_assert_eq!(a.dot(b), *sum); assert_eq!(masks.len(), self.num_rounds * self.mask_length()); - let half = F::from(2).inverse().unwrap(); let polynomial_len = self.mask_length().max(3); @@ -184,9 +180,9 @@ impl Config { { // Fold and compute sumcheck polynomial in one pass. let (c0, c2) = if let Some(w) = prev_round_challenge { - fold_and_compute_polynomial(a, b, w) + a.fold_pair_sumcheck_polynomial(b, w) } else { - compute_sumcheck_polynomial(a, b) + a.sumcheck_polynomial(b) }; let c1 = *sum - c0.double() - c2; @@ -217,8 +213,7 @@ impl Config { } if let Some(w) = prev_round_challenge { // Final fold of the inputs (no polynomial computation). - fold(a, w); - fold(b, w); + a.fold_pair(b, w); } *sum = mask_sum + mask_rlc * *sum; @@ -368,9 +363,11 @@ mod tests { use super::*; use crate::{ algebra::{ + dot, fields::{self, Field64}, multilinear_extend, random_vector, }, + buffer::ActiveBuffer, transcript::DomainSeparator, }; @@ -418,8 +415,8 @@ mod tests { let masks = random_vector(&mut rng, config.mask_length() * config.num_rounds); // Prover - let mut vector = initial_vector.clone(); - let mut covector = initial_covector.clone(); + let mut vector = ActiveBuffer::from_slice(&initial_vector); + let mut covector = ActiveBuffer::from_slice(&initial_covector); let mut sum = initial_sum; let mut prover_state = ProverState::new_std(&ds); let SumcheckOpening { @@ -435,8 +432,14 @@ mod tests { assert_eq!(vector.len(), config.final_size()); assert_eq!(covector.len(), config.final_size()); if config.final_size() == 1 { - assert_eq!(multilinear_extend(&initial_vector, &point), vector[0]); - assert_eq!(multilinear_extend(&initial_covector, &point), covector[0]); + assert_eq!( + multilinear_extend(&initial_vector, &point), + vector.to_slice()[0] + ); + assert_eq!( + multilinear_extend(&initial_covector, &point), + covector.to_slice()[0] + ); } else { // TODO: Check correct folding. } @@ -446,7 +449,10 @@ mod tests { .zip(&point) .map(|(m, x)| univariate_evaluate(m, *x)) .sum(); - assert_eq!(sum, expected_mask_sum + mask_rlc * dot(&vector, &covector)); + assert_eq!( + sum, + expected_mask_sum + mask_rlc * dot(vector.to_slice(), covector.to_slice()) + ); let proof = prover_state.proof(); diff --git a/src/protocols/whir/mod.rs b/src/protocols/whir/mod.rs index 74912af0..68a950ce 100644 --- a/src/protocols/whir/mod.rs +++ b/src/protocols/whir/mod.rs @@ -17,6 +17,7 @@ use crate::{ embedding::{Embedding, Identity}, linear_form::LinearForm, }, + buffer::ActiveBuffer, hash::Hash, protocols::{irs_commit, proof_of_work, sumcheck}, transcript::{ @@ -112,11 +113,11 @@ impl Config { /// After the IRS commit, runs the legacy WHIR OOD step: samples /// `initial_out_domain_samples` random points from the verifier and sends /// each vector's evaluation at each point. - #[cfg_attr(feature = "tracing", instrument(skip_all, fields(size = vectors.first().unwrap().len())))] + #[cfg_attr(feature = "tracing", instrument(skip_all, fields(size = vectors.first().map_or(0, |v| crate::buffer::BufferOps::len(*v)))))] pub fn commit( &self, prover_state: &mut ProverState, - vectors: &[&[M::Source]], + vectors: &[&ActiveBuffer], ) -> Witness where Standard: Distribution, @@ -178,6 +179,7 @@ mod tests { linear_form::{Covector, Evaluate, LinearForm, MultilinearExtension}, random_vector, }, + buffer::{ActiveBuffer, BufferOps}, hash, parameters::ProtocolParameters, protocols::params::DecodingRegime, @@ -285,15 +287,16 @@ mod tests { let mut prover_state = ProverState::new_std(&ds); // Commit to the polynomial and generate auxiliary witness data - let witness = params.commit(&mut prover_state, &[&vector]); + let vector_buffer = ActiveBuffer::from_slice(&vector); + let witness = params.commit(&mut prover_state, &[&vector_buffer]); let prove_linear_forms = build_prove_forms(&points, num_variables, true); // Generate a proof for the given statement and witness let _ = params.prove( &mut prover_state, - vec![Cow::from(vector)], - vec![Cow::Owned(witness)], + &[&vector_buffer], + vec![&witness], prove_linear_forms, Cow::Borrowed(evaluations.as_slice()), ); @@ -426,7 +429,7 @@ mod tests { vec![F::from((i + 1) as u64); num_coeffs] }) .collect(); - let vec_refs = vectors.iter().map(|v| v.as_slice()).collect::>(); + let vec_refs = vectors.iter().collect::>(); let points: Vec<_> = (0..num_points_per_poly) .map(|_| random_vector(thread_rng(), num_variables)) @@ -447,7 +450,7 @@ mod tests { .flat_map(|linear_form| { vec_refs .iter() - .map(|vec| linear_form.evaluate(params.embedding(), vec)) + .map(|&vec| linear_form.evaluate(params.embedding(), vec)) }) .collect::>(); @@ -458,8 +461,12 @@ mod tests { let mut prover_state = ProverState::new_std(&ds); // Commit to each polynomial and generate witnesses + let vector_buffers = vectors + .iter() + .map(|v| ActiveBuffer::from_slice(v)) + .collect::>(); let mut witnesses = Vec::new(); - for &vec in &vec_refs { + for vec in &vector_buffers { let witness = params.commit(&mut prover_state, &[vec]); witnesses.push(witness); } @@ -469,11 +476,8 @@ mod tests { // Batch prove all polynomials together let _ = params.prove( &mut prover_state, - vectors - .iter() - .map(|v| Cow::Borrowed(v.as_slice())) - .collect(), - witnesses.into_iter().map(Cow::Owned).collect(), + &vector_buffers.iter().collect::>(), + witnesses.iter().collect(), prove_linear_forms, Cow::Borrowed(evaluations.as_slice()), ); @@ -616,16 +620,19 @@ mod tests { .instance(&Empty); let mut prover_state = ProverState::new_std(&ds); - let witness1 = params.commit(&mut prover_state, &[&vec1]); - let witness2 = params.commit(&mut prover_state, &[&vec2]); + let vec1_buffer = ActiveBuffer::from_slice(&vec1); + let vec2_buffer = ActiveBuffer::from_slice(&vec2); + let witness1 = params.commit(&mut prover_state, &[&vec1_buffer]); + let witness2 = params.commit(&mut prover_state, &[&vec2_buffer]); let prove_linear_forms = build_prove_forms(&constraint_points, num_variables, false); // Generate proof with mismatched polynomials + let vec_wrong_buffer = ActiveBuffer::from_vec(vec_wrong); let _ = params.prove( &mut prover_state, - vec![Cow::Borrowed(vec1.as_slice()), Cow::from(vec_wrong)], - vec![Cow::Owned(witness1), Cow::Owned(witness2)], + &[&vec1_buffer, &vec_wrong_buffer], + vec![&witness1, &witness2], prove_linear_forms, Cow::Borrowed(evaluations.as_slice()), ); @@ -698,7 +705,7 @@ mod tests { let all_vectors: Vec> = (0..num_witnesses * batch_size) .map(|i| vec![F::from((i + 1) as u64); num_coeffs]) .collect::>(); - let vec_refs = all_vectors.iter().map(|p| p.as_slice()).collect::>(); + let vec_refs = all_vectors.iter().collect::>(); let points: Vec<_> = (0..num_points_per_poly) .map(|_| random_vector(thread_rng(), num_variables)) @@ -719,7 +726,7 @@ mod tests { .flat_map(|linear_form| { vec_refs .iter() - .map(|vec| linear_form.evaluate(params.embedding(), vec)) + .map(|&vec| linear_form.evaluate(params.embedding(), vec)) }) .collect::>(); @@ -730,8 +737,13 @@ mod tests { let mut prover_state = ProverState::new_std(&ds); // Commit using commit_batch (stacks batch_size polynomials per witness) + let vector_buffers = all_vectors + .iter() + .map(|v| ActiveBuffer::from_slice(v)) + .collect::>(); + let buffer_refs = vector_buffers.iter().collect::>(); let mut witnesses = Vec::new(); - for witness_polys in vec_refs.chunks(batch_size) { + for witness_polys in buffer_refs.chunks(batch_size) { let witness = params.commit(&mut prover_state, witness_polys); witnesses.push(witness); } @@ -741,11 +753,8 @@ mod tests { // Batch prove all witnesses together let _ = params.prove( &mut prover_state, - all_vectors - .iter() - .map(|v| Cow::Borrowed(v.as_slice())) - .collect(), - witnesses.into_iter().map(Cow::Owned).collect(), + &vector_buffers.iter().collect::>(), + witnesses.iter().collect(), prove_linear_forms, Cow::Borrowed(evaluations.as_slice()), ); @@ -840,7 +849,7 @@ mod tests { let vectors: Vec> = (0..batch_size) .map(|_| random_vector(thread_rng(), num_coeffs)) .collect(); - let vec_refs = vectors.iter().map(|v| v.as_slice()).collect::>(); + let vec_refs = vectors.iter().collect::>(); // Generate `num_points` random points in the multilinear domain let points: Vec<_> = (0..num_points) @@ -856,7 +865,12 @@ mod tests { let mut prover_state = ProverState::new_std(&ds); // Create a commitment to the polynomial and generate auxiliary witness data - let batched_witness = params.commit(&mut prover_state, &vec_refs); + let vector_buffers = vectors + .iter() + .map(|v| ActiveBuffer::from_slice(v)) + .collect::>(); + let buffer_refs = vector_buffers.iter().collect::>(); + let batched_witness = params.commit(&mut prover_state, &buffer_refs); // Create a weights matrix and evaluations for each polynomial let mut linear_forms: Vec>>> = Vec::new(); @@ -886,11 +900,8 @@ mod tests { .collect::>(); let _ = params.prove( &mut prover_state, - vectors - .iter() - .map(|v| Cow::Borrowed(v.as_slice())) - .collect(), - vec![Cow::Owned(batched_witness)], + &buffer_refs, + vec![&batched_witness], prove_linear_forms, Cow::Borrowed(values.as_slice()), ); diff --git a/src/protocols/whir/prover.rs b/src/protocols/whir/prover.rs index 598ed31d..0710d0ee 100644 --- a/src/protocols/whir/prover.rs +++ b/src/protocols/whir/prover.rs @@ -1,4 +1,4 @@ -use std::{any::Any, borrow::Cow, mem}; +use std::borrow::Cow; use ark_ff::{AdditiveGroup, Field}; use ark_std::rand::{distributions::Standard, prelude::Distribution, CryptoRng, RngCore}; @@ -7,15 +7,8 @@ use tracing::instrument; use super::{Config, Witness}; use crate::{ - algebra::{ - dot, - embedding::Embedding, - eq_weights, lift, - linear_form::{Covector, Evaluate, LinearForm, UnivariateEvaluation}, - mixed_scalar_mul_add, - sumcheck::fold, - tensor_product, - }, + algebra::{dot, embedding::Embedding, eq_weights, linear_form::LinearForm, tensor_product}, + buffer::{ActiveBuffer, Buffer, BufferOps}, hash::Hash, protocols::{geometric_challenge::geometric_challenge, irs_commit, whir::FinalClaim}, transcript::{ @@ -26,7 +19,7 @@ use crate::{ }; enum RoundWitness<'a, F: Field, M: Embedding> { - Initial(Vec>>), + Initial(Vec<&'a Witness>), Round(irs_commit::Witness), } @@ -50,8 +43,8 @@ impl Config { pub fn prove<'a, H, R>( &self, prover_state: &mut ProverState, - vectors: Vec>, - witnesses: Vec>>, + vectors: &[&ActiveBuffer], + witnesses: Vec<&'a Witness>, linear_forms: Vec>>, evaluations: Cow<'a, [M::Target]>, ) -> FinalClaim @@ -73,7 +66,7 @@ impl Config { witnesses.len() * self.initial_committer.num_vectors() ); assert_eq!(evaluations.len(), num_vectors * linear_forms.len()); - for vector in &vectors { + for vector in vectors { assert_eq!(vector.len(), self.initial_size()); } for linear_form in &linear_forms { @@ -85,8 +78,14 @@ impl Config { { use crate::algebra::linear_form::Covector; let covector = Covector::from(&**linear_form); - for (vector, evaluation) in zip_strict(&vectors, evaluations) { - debug_assert_eq!(covector.evaluate(self.embedding(), vector), *evaluation); + for (vector, evaluation) in zip_strict(vectors, evaluations) { + debug_assert_eq!( + vector.mixed_dot( + self.embedding(), + &ActiveBuffer::from_slice(covector.vector.as_slice()) + ), + *evaluation + ); } } if vectors.is_empty() { @@ -110,12 +109,13 @@ impl Config { if j >= vector_offset && j < oods_row.len() + vector_offset { debug_assert_eq!( oods_row[j - vector_offset], - oods_eval.evaluate(self.embedding(), vector) + vector.mixed_univariate_evaluate(self.embedding(), oods_eval.point) ); oods_matrix.push(oods_row[j - vector_offset]); } else { - let eval = oods_eval.evaluate(self.embedding(), vector); + let eval = + vector.mixed_univariate_evaluate(self.embedding(), oods_eval.point); prover_state.prover_message(&eval); oods_matrix.push(eval); } @@ -130,16 +130,11 @@ impl Config { // Random linear combination of the vectors. let mut vector_rlc_coeffs: Vec = geometric_challenge(prover_state, num_vectors); assert_eq!(vector_rlc_coeffs[0], M::Target::ONE); - // Recycle the first input as the accumulator (its coefficient is always ONE). - let mut vectors = vectors.into_iter(); - let first = vectors.next().expect("non-empty"); - let mut vector = match first { - Cow::Borrowed(slice) => lift(self.embedding(), slice), - Cow::Owned(vec) => self.embedding().map_vec(vec), - }; - for (rlc_coeff, input_vector) in zip_strict(&vector_rlc_coeffs[1..], vectors) { - mixed_scalar_mul_add(self.embedding(), &mut vector, *rlc_coeff, &input_vector); - } + let mut vector = ActiveBuffer::::mixed_linear_combination( + self.embedding(), + vectors, + &vector_rlc_coeffs, + ); let mut prev_witness: RoundWitness<'a, M::Target, M> = RoundWitness::Initial(witnesses); @@ -149,26 +144,16 @@ impl Config { let has_constraints = !constraint_rlc_coeffs.is_empty(); let (initial_forms_rlc_coeffs, oods_rlc_coeffs) = constraint_rlc_coeffs.split_at(linear_forms.len()); - // Try to recycle the first linear form as Covector. - let mut covector = vec![]; let mut linear_forms = linear_forms; - if let Some((first, linear_forms)) = linear_forms.split_first_mut() { - debug_assert_eq!(initial_forms_rlc_coeffs[0], M::Target::ONE); - if let Some(covector_form) = - (first.as_mut() as &mut dyn Any).downcast_mut::>() - { - mem::swap(&mut covector, &mut covector_form.vector); - } else { - covector.resize(self.initial_size(), M::Target::ZERO); - first.accumulate(&mut covector, M::Target::ONE); - } - for (rlc_coeff, linear_form) in zip_strict(&initial_forms_rlc_coeffs[1..], linear_forms) - { - linear_form.accumulate(&mut covector, *rlc_coeff); - } - } else if has_constraints { - covector.resize(self.initial_size(), M::Target::ZERO); - } + let mut covector = if has_constraints { + ActiveBuffer::::linear_forms_rlc( + self.initial_size(), + &mut linear_forms, + initial_forms_rlc_coeffs, + ) + } else { + ActiveBuffer::::zeros(0) + }; drop(linear_forms); // Compute "The Sum" @@ -180,17 +165,17 @@ impl Config { .sum(); drop(evaluations); - debug_assert!(!has_constraints || dot(&vector, &covector) == the_sum); + debug_assert!(!has_constraints || vector.dot(&covector) == the_sum); // Add OODS constraints - UnivariateEvaluation::accumulate_many(&oods_evals, &mut covector, oods_rlc_coeffs); + covector.accumulate_univariate_evaluations(&oods_evals, oods_rlc_coeffs); the_sum += zip_strict(oods_rlc_coeffs, oods_matrix.chunks_exact(num_vectors)) .map(|(poly_coeff, row)| *poly_coeff * dot(&vector_rlc_coeffs, row)) .sum::(); drop(oods_evals); drop(oods_matrix); - debug_assert!(!has_constraints || dot(&vector, &covector) == the_sum); + debug_assert!(!has_constraints || vector.dot(&covector) == the_sum); // Run initial sumcheck on batched vectors with combined statement let mut folding_randomness = if has_constraints { @@ -207,15 +192,15 @@ impl Config { self.initial_skip_pow.prove(prover_state); // Fold vector for &f in &folding_randomness { - fold(&mut vector, f); + vector.fold(f); } // Covector must be all zeros. - covector = vec![M::Target::ZERO; self.initial_sumcheck.final_size()]; + covector = ActiveBuffer::::zeros(self.initial_sumcheck.final_size()); folding_randomness }; let mut evaluation_point = folding_randomness.clone(); - debug_assert_eq!(dot(&vector, &covector), the_sum); + debug_assert_eq!(vector.dot(&covector), the_sum); // Execute standard WHIR rounds on the batched vectors for (round_index, round_config) in self.round_configs.iter().enumerate() { @@ -258,13 +243,9 @@ impl Config { ))) .collect::>(); let stir_rlc_coeffs = geometric_challenge(prover_state, stir_challenges.len()); - UnivariateEvaluation::accumulate_many( - &stir_challenges, - &mut covector, - &stir_rlc_coeffs, - ); + covector.accumulate_univariate_evaluations(&stir_challenges, &stir_rlc_coeffs); the_sum += dot(&stir_rlc_coeffs, &stir_evaluations); - debug_assert_eq!(dot(&vector, &covector), the_sum); + debug_assert_eq!(vector.dot(&covector), the_sum); // Run sumcheck for this round folding_randomness = round_config @@ -273,7 +254,7 @@ impl Config { .round_challenges; evaluation_point.extend(folding_randomness.iter().copied()); - debug_assert_eq!(dot(&vector, &covector), the_sum); + debug_assert_eq!(vector.dot(&covector), the_sum); prev_witness = RoundWitness::Round(new_witness); vector_rlc_coeffs = vec![M::Target::ONE]; @@ -281,7 +262,7 @@ impl Config { // Directly send the vector to the verifier. assert_eq!(vector.len(), self.final_sumcheck.initial_size()); - for coeff in &vector { + for coeff in vector.to_slice() { prover_state.prover_message(coeff); } diff --git a/src/protocols/whir_zk/committer.rs b/src/protocols/whir_zk/committer.rs index e895bf85..08ce6ba7 100644 --- a/src/protocols/whir_zk/committer.rs +++ b/src/protocols/whir_zk/committer.rs @@ -1,3 +1,5 @@ +// IGNORE CHANGES TO THIS FILE - NOT FULLY PORTED TO PROPERLY USE BUFFER ABSTRACTION. + use ark_ff::Field; use ark_std::rand::{distributions::Standard, prelude::Distribution}; #[cfg(feature = "tracing")] @@ -6,6 +8,7 @@ use tracing::instrument; use super::{utils::BlindingPolynomials, Config}; use crate::{ algebra::embedding::Identity, + buffer::{ActiveBuffer, BufferOps}, hash::Hash, protocols::whir, transcript::{ @@ -44,7 +47,7 @@ impl Config { pub fn commit( &self, prover_state: &mut ProverState, - polynomials: &[&[F]], + polynomials: &[&ActiveBuffer], ) -> Witness where Standard: Distribution, @@ -65,6 +68,7 @@ impl Config { let num_blinding_variables = self.num_blinding_variables(); let num_witness_variables = self.num_witness_variables(); for &poly in polynomials { + let poly = poly.to_slice(); let blinding = BlindingPolynomials::sample( prover_state.rng(), num_blinding_variables, @@ -88,9 +92,10 @@ impl Config { .zip(mask.iter().cycle()) .map(|(&coeff, &m)| coeff + m) .collect::>(); + let f_hat_buffer = ActiveBuffer::from_slice(&f_hat_vec); let witness = self .blinded_commitment - .commit(prover_state, &[f_hat_vec.as_slice()]); + .commit(prover_state, &[&f_hat_buffer]); f_hat_vectors.push(f_hat_vec); f_hat_witnesses.push(witness); blinding_polynomials.push(blinding); @@ -112,10 +117,11 @@ impl Config { ); blinding_vectors.extend(layout); } - let blinding_vector_refs = blinding_vectors + let blinding_buffers = blinding_vectors .iter() - .map(Vec::as_slice) + .map(|v| ActiveBuffer::from_slice(v)) .collect::>(); + let blinding_vector_refs = blinding_buffers.iter().collect::>(); let blinding_witness = self .blinding_commitment .commit(prover_state, &blinding_vector_refs); diff --git a/src/protocols/whir_zk/mod.rs b/src/protocols/whir_zk/mod.rs index 8275226e..b4c5eedb 100644 --- a/src/protocols/whir_zk/mod.rs +++ b/src/protocols/whir_zk/mod.rs @@ -1,3 +1,5 @@ +// IGNORE CHANGES TO THIS FILE - NOT FULLY PORTED TO PROPERLY USE BUFFER ABSTRACTION. + #![cfg(feature = "rs_in_order")] // TODO: Support permuted. mod committer; mod prover; @@ -257,6 +259,7 @@ mod tests { linear_form::{Covector, Evaluate, LinearForm, MultilinearExtension}, random_vector, }, + buffer::{ActiveBuffer, BufferOps}, hash, parameters::ProtocolParameters, protocols::params::DecodingRegime, @@ -355,13 +358,15 @@ mod tests { .session(&tag) .instance(&Empty); let mut prover_state = ProverState::new_std(&ds); - let witness = params.commit(&mut prover_state, vectors); + let vector_buffers = vectors + .iter() + .map(|v| ActiveBuffer::from_slice(v)) + .collect::>(); + let vector_refs = vector_buffers.iter().collect::>(); + let witness = params.commit(&mut prover_state, &vector_refs); let _ = params.prove( &mut prover_state, - vectors - .iter() - .map(|&v| Cow::Borrowed(v)) - .collect::>(), + &vector_refs, witness, prove_forms, Cow::Borrowed(evaluations), @@ -447,13 +452,15 @@ mod tests { .session(&format!("zk-stage1-negative {}:{}", file!(), line!())) .instance(&Empty); let mut prover_state = ProverState::new_std(&ds); - let witness = params.commit(&mut prover_state, &vectors); + let vector_buffers = vectors + .iter() + .map(|v| ActiveBuffer::from_slice(v)) + .collect::>(); + let vector_refs = vector_buffers.iter().collect::>(); + let witness = params.commit(&mut prover_state, &vector_refs); let _ = params.prove( &mut prover_state, - vectors - .iter() - .map(|&v| Cow::Borrowed(v)) - .collect::>(), + &vector_refs, witness, prove_forms, Cow::Borrowed(&evaluations), @@ -501,13 +508,15 @@ mod tests { .session(&format!("zk-stage1-tamper {}:{}", file!(), line!())) .instance(&Empty); let mut prover_state = ProverState::new_std(&ds); - let witness = params.commit(&mut prover_state, &vectors); + let vector_buffers = vectors + .iter() + .map(|v| ActiveBuffer::from_slice(v)) + .collect::>(); + let vector_refs = vector_buffers.iter().collect::>(); + let witness = params.commit(&mut prover_state, &vector_refs); let _ = params.prove( &mut prover_state, - vectors - .iter() - .map(|&v| Cow::Borrowed(v)) - .collect::>(), + &vector_refs, witness, prove_forms, Cow::Borrowed(&evaluations), @@ -562,10 +571,11 @@ mod tests { let outcome = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { let mut prover_state = ProverState::new_std(&ds); - let witness = params.commit(&mut prover_state, &[&vector]); + let vector_buffer = ActiveBuffer::from_slice(&vector); + let witness = params.commit(&mut prover_state, &[&vector_buffer]); let _ = params.prove( &mut prover_state, - vec![Cow::Borrowed(&vector)], + &[&vector_buffer], witness, prove_forms, Cow::Owned(vec![wrong_evaluation]), diff --git a/src/protocols/whir_zk/prover.rs b/src/protocols/whir_zk/prover.rs index 549e5f36..8bd23ac9 100644 --- a/src/protocols/whir_zk/prover.rs +++ b/src/protocols/whir_zk/prover.rs @@ -1,3 +1,5 @@ +// IGNORE CHANGES TO THIS FILE - NOT FULLY PORTED TO PROPERLY USE BUFFER ABSTRACTION. + use std::borrow::Cow; use ark_ff::Field; @@ -14,6 +16,7 @@ use crate::{ linear_form::{Covector, Evaluate, LinearForm}, mixed_dot, scalar_mul_add, }, + buffer::{ActiveBuffer, BufferOps}, hash::Hash, protocols::{ whir::FinalClaim, @@ -206,13 +209,13 @@ impl Config { /// inner witness-side WHIR prover. #[allow(clippy::too_many_lines)] #[cfg_attr(feature = "tracing", instrument(skip_all))] - pub fn prove<'a, H, R>( + pub fn prove( &self, prover_state: &mut ProverState, - vectors: Vec>, + vectors: &[&ActiveBuffer], witness: Witness, linear_forms: Vec>>, - evaluations: Cow<'a, [F]>, + evaluations: Cow<'_, [F]>, ) -> FinalClaim where Standard: Distribution, @@ -266,7 +269,6 @@ impl Config { let num_witness_variables = self.num_witness_variables(); let num_blinding_variables = self.num_blinding_variables(); let num_witness_variables_plus_1 = num_witness_variables + 1; - drop(vectors); // TODO: These are never touched? // Compute w_folded evaluations of all blinding vectors before rho for binding. let (w_folded_weights, m_evals, w_folded_blinding_evals) = { @@ -397,10 +399,16 @@ impl Config { let result = { #[cfg(feature = "tracing")] let _span = tracing::info_span!("inner_blinded_prove").entered(); + let f_hat_buffers = f_hat_vectors + .iter() + .map(|v| ActiveBuffer::from_slice(v)) + .collect::>(); + let f_hat_refs = f_hat_buffers.iter().collect::>(); + let f_hat_witness_refs = f_hat_witnesses.iter().collect::>(); self.blinded_commitment.prove( prover_state, - f_hat_vectors.into_iter().map(Cow::Owned).collect(), - f_hat_witnesses.into_iter().map(Cow::Owned).collect(), + &f_hat_refs, + f_hat_witness_refs, linear_forms, Cow::Owned(modified_evaluations), ) @@ -424,10 +432,15 @@ impl Config { .concat(); // Blinding sub-proof result is discarded: the blinding WHIR's // evaluation point is not needed by the outer protocol. + let blinding_buffers = blinding_vectors + .iter() + .map(|v| ActiveBuffer::from_slice(v)) + .collect::>(); + let blinding_refs = blinding_buffers.iter().collect::>(); let _ = self.blinding_commitment.prove( prover_state, - blinding_vectors.into_iter().map(Cow::Owned).collect(), - vec![Cow::Owned(blinding_witness)], + &blinding_refs, + vec![&blinding_witness], blinding_forms, Cow::Owned(all_blinding_claims), );