diff --git a/crates/circuit/src/matrix_products.rs b/crates/circuit/src/matrix_products.rs index 93ed24c..19fc7e2 100644 --- a/crates/circuit/src/matrix_products.rs +++ b/crates/circuit/src/matrix_products.rs @@ -45,19 +45,20 @@ impl RuntimeModulus { } pub(crate) fn reduce(&self, value: &StoredInteger) -> [u64; PRIME_LIMBS] { - if value.words.is_empty() { + let words = value.words(); + if words.is_empty() { return [0; PRIME_LIMBS]; } if PRIME_LIMBS == 2 && self.words[1] >> 63 != 0 { let modulus = [self.words[0], self.words[1]]; - let reduced = reduce_normalized_2(&value.words, modulus); + let reduced = reduce_normalized_2(words, modulus); let reduced = if value.is_negative() { // The stored words are x + 2^(64k), where x is negative and // k is the trimmed storage length. Recover x modulo p without // materializing the two's-complement magnitude. let mut width_power = [1, 0]; - for _ in 0..value.words.len() { + for _ in 0..words.len() { width_power = reduce_step_2(width_power, 0, modulus); } subtract_mod_2(reduced, width_power, modulus) @@ -73,10 +74,10 @@ impl RuntimeModulus { let negative = value.is_negative(); let magnitude; let words = if negative { - magnitude = twos_complement_magnitude(&value.words); + magnitude = twos_complement_magnitude(words); magnitude.as_slice() } else { - value.words.as_ref() + words }; let mut reduced = [0; PRIME_LIMBS]; @@ -290,27 +291,71 @@ pub struct MatrixProducts { /// An exact signed two's-complement integer with redundant sign limbs removed. /// -/// Every matrix row has a materialized [`StoredInteger`]. Zero uses an empty -/// boxed slice, so it does not perform a separate limb allocation. +/// Every matrix row has a materialized [`StoredInteger`]. A value of at most +/// one word is held inline; only wider values allocate their limbs. #[derive(Clone, Debug, Eq, Hash, PartialEq)] pub struct StoredInteger { - words: Box<[u64]>, + words: Words, +} + +/// Normalized little-endian two's-complement words. +/// +/// The derived equality and hash need one representation per integer: +/// [`StoredInteger::from_normalized`] is the only place that picks a variant. +#[derive(Clone, Debug, Eq, Hash, PartialEq)] +enum Words { + /// Zero, or the only word of a one-word value. + Inline(u64), + /// Two words or more. + Heap(Box<[u64]>), } impl StoredInteger { + const ZERO: Self = Self { + words: Words::Inline(0), + }; + + /// Stores normalized little-endian two's-complement words; zero is empty. + fn from_normalized(words: W) -> Self + where + W: AsRef<[u64]> + Into>, + { + let words = match *words.as_ref() { + [] => Words::Inline(0), + [word] => Words::Inline(word), + _ => Words::Heap(words.into()), + }; + Self { words } + } + /// Normalized little-endian two's-complement words; zero is empty. pub fn words(&self) -> &[u64] { - &self.words + match &self.words { + Words::Inline(0) => &[], + Words::Inline(word) => std::slice::from_ref(word), + Words::Heap(words) => words, + } } /// Whether this integer is exactly zero. pub fn is_zero(&self) -> bool { - self.words.is_empty() + matches!(self.words, Words::Inline(0)) } /// Whether this integer is negative. pub fn is_negative(&self) -> bool { - self.words.last().is_some_and(|word| word >> 63 == 1) + match &self.words { + Words::Inline(word) => word >> 63 == 1, + Words::Heap(words) => words.last().is_some_and(|word| word >> 63 == 1), + } + } + + /// Bytes of limb storage allocated outside this value. + pub fn heap_bytes(&self) -> usize { + match &self.words { + Words::Inline(_) => 0, + Words::Heap(words) => size_of_val(words.as_ref()), + } } } @@ -319,9 +364,7 @@ impl From<&num_bigint::BigInt> for StoredInteger { fn from(value: &num_bigint::BigInt) -> Self { use num_bigint::{BigUint, Sign}; if value.is_zero() { - return Self { - words: Box::default(), - }; + return Self::ZERO; } let (sign, bytes) = value.to_bytes_le(); @@ -351,9 +394,7 @@ impl From<&num_bigint::BigInt> for StoredInteger { break; } } - Self { - words: magnitude.into_boxed_slice(), - } + Self::from_normalized(magnitude) } } @@ -378,9 +419,7 @@ impl From<&Integer> for StoredInteger { /// Stores a gadget-local fixed integer without changing its value. fn from(value: &Integer) -> Self { if value.is_zero() { - return Self { - words: Box::default(), - }; + return Self::ZERO; } let words = value.words(); @@ -394,9 +433,7 @@ impl From<&Integer> for StoredInteger { break; } } - Self { - words: words[..len].into(), - } + Self::from_normalized(&words[..len]) } } @@ -616,6 +653,25 @@ mod tests { } } + #[test] + fn only_values_wider_than_one_word_allocate_limbs() { + assert_eq!(size_of::(), 2 * size_of::()); + for value in [0_i128, 1, -1, i128::from(i64::MAX), i128::from(i64::MIN)] { + let stored = StoredInteger::from(&value); + assert_eq!(stored.heap_bytes(), 0, "{value}"); + assert_eq!( + StoredInteger::from(&R::from(value)).heap_bytes(), + 0, + "{value}" + ); + } + for value in [1_i128 << 63, -(1 << 63) - 1, 1 << 100, -(1 << 100)] { + let stored = StoredInteger::from(&value); + assert!(stored.words().len() >= 2, "{value}"); + assert_eq!(stored.heap_bytes(), size_of_val(stored.words()), "{value}"); + } + } + #[test] fn arbitrary_bigints_round_trip_through_stored_integers() { let boundary = R::one() << 128_usize; diff --git a/crates/circuit/src/matrix_sparse.rs b/crates/circuit/src/matrix_sparse.rs index fa29a98..589704c 100644 --- a/crates/circuit/src/matrix_sparse.rs +++ b/crates/circuit/src/matrix_sparse.rs @@ -12,7 +12,7 @@ use common::BitzConstraintRing; use crypto_bigint::modular::{FixedMontyForm, FixedMontyParams}; use crypto_bigint::{Odd, U128}; use rayon::prelude::*; -use std::mem::{size_of, size_of_val}; +use std::mem::size_of; use thiserror::Error; const PARALLEL_NNZ_THRESHOLD: usize = 1 << 15; @@ -173,7 +173,7 @@ impl MaterializedAbc { + self .coefficients .iter() - .map(|coefficient| size_of_val(coefficient.words())) + .map(StoredInteger::heap_bytes) .sum::() } diff --git a/crates/circuit/src/matrix_transpose.rs b/crates/circuit/src/matrix_transpose.rs index 54edd64..e4fe513 100644 --- a/crates/circuit/src/matrix_transpose.rs +++ b/crates/circuit/src/matrix_transpose.rs @@ -10,7 +10,7 @@ use std::sync::atomic::{AtomicU64, Ordering}; use crate::witgen::Z; use crate::{BoolWitness, Circuit, HintResult, PackedBits, ScalarBits, WitnessContext}; use common::{TransposedWeights, VirtualMap, VirtualMapError}; -use field::F128; +use field::{F128, Wide256}; use rayon::prelude::*; use thiserror::Error; @@ -265,10 +265,14 @@ impl MTransposeRecorder { } } + let column_offsets = column_offsets.into_boxed_slice(); + let row_indices = row_indices.into_boxed_slice(); + let digest = csc_digest(row_count, &column_offsets, &row_indices); MaterializedMTranspose { row_count, - column_offsets: column_offsets.into_boxed_slice(), - row_indices: row_indices.into_boxed_slice(), + column_offsets, + row_indices, + digest, } } } @@ -279,6 +283,8 @@ pub struct MaterializedMTranspose { row_count: usize, column_offsets: Box<[u32]>, row_indices: Box<[u32]>, + /// Precomputed [`csc_digest`] of the fields above. + digest: [u8; 32], } impl MaterializedMTranspose { @@ -344,11 +350,7 @@ impl MaterializedMTranspose { } F128::new(lo, hi) }; - let parallel = force_parallel.unwrap_or_else(|| { - rayon::current_num_threads() > 1 - && self.row_indices.len() >= PARALLEL_MATRIX_NNZ_THRESHOLD - }); - if parallel { + if self.should_gathers_in_parallel(force_parallel) { output .par_iter_mut() .enumerate() @@ -361,6 +363,56 @@ impl MaterializedMTranspose { } Ok(()) } + + /// Whether a gather splits its columns across threads, unless forced. + fn should_gathers_in_parallel(&self, force_parallel: Option) -> bool { + force_parallel.unwrap_or_else(|| { + rayon::current_num_threads() > 1 + && self.row_indices.len() >= PARALLEL_MATRIX_NNZ_THRESHOLD + }) + } + + /// Computes `(row_weights (x) column_weights) * M` without writing the + /// tensor out: one unreduced product per nonzero entry in M, one reduction per column. + fn transpose_inner( + &self, + row_weights: &[F128], + column_weights: &[F128], + force_parallel: Option, + ) -> Result { + let rows_per_column = row_weights.len(); + // A bit count past `usize` covers every row of `M`. + let covers_m = rows_per_column + .checked_mul(column_weights.len()) + .is_none_or(|bits| bits >= self.row_count); + if !rows_per_column.is_power_of_two() || !covers_m { + return Err(VirtualMapError::WeightCountMismatch); + } + let log_rows = rows_per_column.trailing_zeros(); + let row_mask = rows_per_column - 1; + + let evaluate = |column: usize| { + let start = self.column_offsets[column] as usize; + let end = self.column_offsets[column + 1] as usize; + let mut sum = Wide256::zero(); + for &index in &self.row_indices[start..end] { + let index = index as usize; + sum += Wide256::mul( + row_weights[index & row_mask], + column_weights[index >> log_rows], + ); + } + sum.reduce() + }; + // Column zero is `M`'s constant column. + let bit_columns = 1..self.column_count(); + let on_bits = if self.should_gathers_in_parallel(force_parallel) { + bit_columns.into_par_iter().map(evaluate).collect() + } else { + bit_columns.map(evaluate).collect() + }; + Ok(TransposedWeights::new(on_bits, evaluate(0))) + } } /// Domain separation for the digest. @@ -368,37 +420,37 @@ const DIGEST_DOMAIN: &[u8] = b"bitz/virtual-map/csc/v1"; const DIGEST_CHUNK: usize = 1 << 12; -impl VirtualMap for MaterializedMTranspose { - fn transpose(&self, weights: &[F128]) -> Result { - let weights = weights - .get(..self.row_count) - .ok_or(VirtualMapError::WeightCountMismatch)?; - let mut transposed = self - .apply(weights) - .map_err(|_| VirtualMapError::WeightCountMismatch)?; +/// Hashes the shape and the CSC payload of `M`. +fn csc_digest(row_count: usize, column_offsets: &[u32], row_indices: &[u32]) -> [u8; 32] { + let mut hasher = blake3::Hasher::new(); + hasher.update(DIGEST_DOMAIN); + for count in [row_count, column_offsets.len() - 1, row_indices.len()] { + hasher.update(&(count as u64).to_le_bytes()); + } + // Hash the offsets too: two matrices can share every row index and + // still split the columns differently. + for indices in [column_offsets, row_indices] { + let mut bytes = Vec::with_capacity(DIGEST_CHUNK * size_of::()); + for chunk in indices.chunks(DIGEST_CHUNK) { + bytes.clear(); + bytes.extend(chunk.iter().flat_map(|index| index.to_le_bytes())); + hasher.update(&bytes); + } + } + *hasher.finalize().as_bytes() +} - // Column zero is `M`'s constant column. - let on_bits = transposed.split_off(1); - Ok(TransposedWeights::new(on_bits, transposed[0])) +impl VirtualMap for MaterializedMTranspose { + fn transpose( + &self, + row_weights: &[F128], + column_weights: &[F128], + ) -> Result { + self.transpose_inner(row_weights, column_weights, None) } fn digest(&self) -> [u8; 32] { - let mut hasher = blake3::Hasher::new(); - hasher.update(DIGEST_DOMAIN); - for count in [self.row_count, self.column_count(), self.nonzero_count()] { - hasher.update(&(count as u64).to_le_bytes()); - } - // Hash the offsets too: two matrices can share every row index and - // still split the columns differently. - for indices in [&self.column_offsets, &self.row_indices] { - let mut bytes = Vec::with_capacity(DIGEST_CHUNK * size_of::()); - for chunk in indices.chunks(DIGEST_CHUNK) { - bytes.clear(); - bytes.extend(chunk.iter().flat_map(|index| index.to_le_bytes())); - hasher.update(&bytes); - } - } - *hasher.finalize().as_bytes() + self.digest } fn h_len(&self) -> usize { @@ -552,6 +604,14 @@ mod tests { .collect() } + /// Row and column factors that exactly cover `bits` virtual bits rounded up + /// to whole columns of `rows_per_column`. + fn factors(bits: usize, rows_per_column: usize) -> (Vec, Vec) { + let mut row_weights = challenges(rows_per_column + bits.div_ceil(rows_per_column)); + let column_weights = row_weights.split_off(rows_per_column); + (row_weights, column_weights) + } + /// `` must equal ``. #[test] fn the_transposed_claim_matches_the_claim_on_the_integer_witness() { @@ -570,12 +630,14 @@ mod tests { assert_eq!(map.f_len(), f.bit_len() + 1, "M's columns are 1 || f"); assert!(h.bit(0), "M's first row is the constant one"); - let weights = challenges(map.h_len()); + let (row_weights, column_weights) = factors(map.h_len(), 4); let direct = (0..h.bit_len()) .filter(|index| h.bit(*index)) - .fold(F128::new(0, 0), |sum, index| sum + weights[index]); + .fold(F128::new(0, 0), |sum, index| { + sum + row_weights[index % 4] * column_weights[index / 4] + }); - let transposed = map.transpose(&weights).unwrap(); + let transposed = map.transpose(&row_weights, &column_weights).unwrap(); let through_f = (0..f.bit_len()) .filter(|index| f.bit(*index)) .fold(transposed.constant_weight(), |sum, index| { @@ -592,9 +654,28 @@ mod tests { example_circuit(&mut materializer, &inputs); let map = materializer.finish(); - let short = map.h_len() - 1; + let (row_weights, column_weights) = factors(map.h_len(), 4); + let short = column_weights.len() - 1; + assert_eq!( + map.transpose(&row_weights, &column_weights[..short]), + Err(VirtualMapError::WeightCountMismatch) + ); + } + + #[test] + fn row_counts_that_are_not_a_power_of_two_are_rejected() { + let mut materializer = MTransposeGenerator::new(3); + let inputs = materializer.take_boxed_inputs(); + example_circuit(&mut materializer, &inputs); + let map = materializer.finish(); + + let (row_weights, column_weights) = factors(map.h_len(), 3); assert_eq!( - map.transpose(&challenges(short)), + map.transpose(&row_weights, &column_weights), + Err(VirtualMapError::WeightCountMismatch) + ); + assert_eq!( + map.transpose(&[], &column_weights), Err(VirtualMapError::WeightCountMismatch) ); } @@ -605,15 +686,15 @@ mod tests { let inputs = materializer.take_boxed_inputs(); example_circuit(&mut materializer, &inputs); let map = materializer.finish(); - let weights = challenges(map.h_len()); - let expected = map.transpose(&weights).unwrap(); + let (row_weights, column_weights) = factors(map.h_len(), 4); + let expected = map.transpose(&row_weights, &column_weights).unwrap(); for padded_len in [ - map.h_len().next_power_of_two(), - 2 * map.h_len().next_power_of_two(), + column_weights.len().next_power_of_two(), + 2 * column_weights.len().next_power_of_two(), ] { - let mut padded = weights.clone(); + let mut padded = column_weights.clone(); padded.resize(padded_len, F128::new(13, 17)); - assert_eq!(map.transpose(&padded).unwrap(), expected); + assert_eq!(map.transpose(&row_weights, &padded).unwrap(), expected); } } @@ -696,6 +777,35 @@ mod tests { assert_eq!(parallel, sequential); } + /// The factored gather must agree with `r * M` on the tensor written out. + #[test] + fn factored_and_dense_matrix_gathers_agree() { + const ROWS_PER_COLUMN: usize = 128; + let mut materializer = MTransposeGenerator::new(COMPRESSION_INPUT_BITS); + let matrix_inputs = materializer.take_boxed_inputs(); + let _ = compression_circuit(&mut materializer, &matrix_inputs); + let transpose = materializer.finish(); + + let (row_weights, column_weights) = factors(transpose.row_count(), ROWS_PER_COLUMN); + let r: Vec<_> = (0..transpose.row_count()) + .map(|index| { + row_weights[index % ROWS_PER_COLUMN] * column_weights[index / ROWS_PER_COLUMN] + }) + .collect(); + let mut dense = transpose.apply(&r).unwrap(); + let on_bits = dense.split_off(1); + let expected = TransposedWeights::new(on_bits, dense[0]); + + for parallel in [false, true] { + assert_eq!( + transpose + .transpose_inner(&row_weights, &column_weights, Some(parallel)) + .unwrap(), + expected + ); + } + } + #[test] fn materialized_supports_can_exceed_the_inline_capacity() { let mut recorder = MTransposeRecorder::with_witnesses(6); diff --git a/crates/circuit/src/matrix_wengert.rs b/crates/circuit/src/matrix_wengert.rs index 21e3f1e..5776d95 100644 --- a/crates/circuit/src/matrix_wengert.rs +++ b/crates/circuit/src/matrix_wengert.rs @@ -9,7 +9,7 @@ use std::collections::HashMap; use std::iter::Sum; -use std::mem::{size_of, size_of_val}; +use std::mem::size_of; use std::ops::{Add, AddAssign, Mul, Neg, Sub, SubAssign}; use std::sync::{Arc, Mutex}; @@ -622,7 +622,7 @@ impl WengertTape { + self .coefficients .iter() - .map(|coefficient| size_of_val(coefficient.words())) + .map(StoredInteger::heap_bytes) .sum::() } diff --git a/crates/common/src/params.rs b/crates/common/src/params.rs index 60d1211..070310d 100644 --- a/crates/common/src/params.rs +++ b/crates/common/src/params.rs @@ -259,6 +259,7 @@ mod tests { fn transpose( &self, _: &[F128], + _: &[F128], ) -> Result { unreachable!("the parameter gates never apply the map") } diff --git a/crates/common/src/virtual_map.rs b/crates/common/src/virtual_map.rs index 0908d3c..756f62b 100644 --- a/crates/common/src/virtual_map.rs +++ b/crates/common/src/virtual_map.rs @@ -10,11 +10,10 @@ //! ``` //! //! `M` stays with the caller. This crate needs only `M^T v` and a digest. +//! `v` is passed as separate row and column factors. use field::F128; use num_traits::ConstZero; -#[cfg(feature = "parallel")] -use rayon::prelude::*; use crate::{ BitZParams, BitzClaimField, ClaimError, LinearClaim, OpeningQuery, Shape, VirtualParams, @@ -48,7 +47,8 @@ pub enum VirtualMapError { PointLengthMismatch, /// The inner-product coefficients do not cover the padded virtual witness. ClaimWeightCountMismatch, - /// The input omits coordinates of `h`, or the output does not have one weight per bit of `f`. + /// The input factors omit coordinates of `h` or have no power-of-two row count, + /// or the output does not have one weight per bit of `f`. WeightCountMismatch, /// The map's weights on `f` outnumber the committed bits. CommittedShapeTooSmall, @@ -58,10 +58,19 @@ pub enum VirtualMapError { pub trait VirtualMap { /// Computes `M^T weights`, splitting off the constant-one coordinate. /// - /// `weights[i]` multiplies bit `h[i]`. At least `h_len()` weights are required; - /// any remaining weights multiply zero-padded bits of `h` and are ignored. + /// `weights` here are `(row_weights (x) column_weights)`. + /// + /// `row_weights[b] * column_weights[c]` multiplies bit `h[c * row_count + b]`, + /// with `row_count = row_weights.len()` a power of two. The factors must cover + /// `h_len()` bits; remaining weights multiply zero-padded bits of `h` and are + /// ignored. + /// /// The result has `f_len() - 1` bit weights and the constant-column weight. - fn transpose(&self, weights: &[F128]) -> Result; + fn transpose( + &self, + row_weights: &[F128], + column_weights: &[F128], + ) -> Result; /// Binds the map into the statement frame. /// @@ -156,19 +165,24 @@ impl<'a, F: BitzClaimField, M: VirtualMap> VirtualStatement<'a, F, M> { /// Rewrites a claim on padded virtual bits into an opening on committed bits. /// - /// Coefficients use `column * row_count + row` order. The map drops weights on - /// virtual padding. The opening subtracts the constant-column weight from the - /// target and zero-pads the remaining weights to the commitment size, refusing - /// more weights than it has. A single-column `InnerProduct` holds the dense - /// weights with column weight one. + /// Coefficients use `column * row_count + row` order, so an `Mle` point splits + /// into its row and column coordinates. The map drops weights on virtual + /// padding. The opening subtracts the constant-column weight from the target + /// and zero-pads the remaining weights to the commitment size, refusing more + /// weights than it has. A single-column `InnerProduct` holds the dense weights + /// with column weight one. pub fn transpose_query(&self, query: OpeningQuery) -> Result { let shape = self.params.claim().shape(); - let (weights, target) = match query { + let (transposed, target) = match query { OpeningQuery::Mle { point, target } => { if point.len() != shape.log_bits() { return Err(VirtualMapError::PointLengthMismatch); } - (poly::eq_table(&point), target) + let (row_point, column_point) = point.split_at(shape.log_rows()); + let eq_row = poly::eq_table(row_point); + let eq_col = poly::eq_table(column_point); + let transposed = self.map.transpose(&eq_row, &eq_col)?; + (transposed, target) } OpeningQuery::InnerProduct { claim } => { if claim.row_weights().len() != shape.rows() @@ -176,14 +190,12 @@ impl<'a, F: BitzClaimField, M: VirtualMap> VirtualStatement<'a, F, M> { { return Err(VirtualMapError::ClaimWeightCountMismatch); } - ( - flatten(claim.row_weights(), claim.column_weights()), - claim.target(), - ) + let transposed = self + .map + .transpose(claim.row_weights(), claim.column_weights())?; + (transposed, claim.target()) } }; - let transposed = self.map.transpose(&weights)?; - drop(weights); if Some(transposed.weights().len()) != self.map.f_len().checked_sub(1) { return Err(VirtualMapError::WeightCountMismatch); } @@ -202,28 +214,6 @@ impl<'a, F: BitzClaimField, M: VirtualMap> VirtualStatement<'a, F, M> { } } -/// `row_weights (x) column_weights` written out per bit of `h`, one column -/// after another. -fn flatten(row_weights: &[F128], column_weights: &[F128]) -> Vec { - let mut weights = vec![F128::ZERO; row_weights.len() * column_weights.len()]; - let process_column = |(slot, &scale): (&mut [F128], &F128)| { - for (weight, &row) in slot.iter_mut().zip(row_weights) { - *weight = scale * row; - } - }; - #[cfg(feature = "parallel")] - weights - .par_chunks_mut(row_weights.len()) - .zip(column_weights) - .for_each(process_column); - #[cfg(not(feature = "parallel"))] - weights - .chunks_mut(row_weights.len()) - .zip(column_weights) - .for_each(process_column); - weights -} - #[cfg(test)] mod tests { use super::*; @@ -240,15 +230,21 @@ mod tests { } impl VirtualMap for DenseMap { - fn transpose(&self, weights: &[F128]) -> Result { - if weights.len() < self.h_len() { + fn transpose( + &self, + row_weights: &[F128], + column_weights: &[F128], + ) -> Result { + let row_count = row_weights.len(); + if row_count * column_weights.len() < self.h_len() { return Err(VirtualMapError::WeightCountMismatch); } let mut transposed = vec![F128::ZERO; self.columns]; - for (row, weight) in self.rows.iter().zip(weights) { + for (index, row) in self.rows.iter().enumerate() { + let weight = row_weights[index % row_count] * column_weights[index / row_count]; for (column, set) in row.iter().enumerate() { if *set { - transposed[column] += *weight; + transposed[column] += weight; } } } @@ -306,7 +302,7 @@ mod tests { .filter(|(bit, _)| **bit) .fold(F128::ZERO, |sum, (_, weight)| sum + *weight); - let transposed = map.transpose(&weights).unwrap(); + let transposed = map.transpose(&weights, &[F128::ONE]).unwrap(); let through_f = f .iter() .zip(transposed.weights()) @@ -335,7 +331,7 @@ mod tests { #[test] fn missing_virtual_witness_weights_are_rejected() { assert_eq!( - map().transpose(&[F128::ONE; 3]), + map().transpose(&[F128::ONE; 3], &[F128::ONE]), Err(VirtualMapError::WeightCountMismatch) ); } @@ -344,9 +340,15 @@ mod tests { fn weights_on_zero_padding_do_not_change_the_transpose() { let map = map(); let weights = [F128::new(7, 11); 4]; + let expected = map.transpose(&weights, &[F128::ONE]); let mut padded = weights.to_vec(); padded.extend([F128::new(13, 17); 4]); - assert_eq!(map.transpose(&weights), map.transpose(&padded)); + assert_eq!(map.transpose(&padded, &[F128::ONE]), expected); + // Padding as whole columns of `h`. + assert_eq!( + map.transpose(&weights, &[F128::ONE, F128::new(13, 17)]), + expected + ); } fn params() -> BitZParams { @@ -398,6 +400,39 @@ mod tests { assert_eq!(claim.target(), target - rows[0] * columns[0]); } + /// The split point must weigh virtual bit `i` by entry `i` of the equality + /// table over all the virtual bits. + #[test] + fn mle_queries_weigh_virtual_bits_by_the_equality_table() { + let params = params(); + let row_count = params.shape().rows(); + // `h_1 = f_0` in column zero, `h_far = f_0 + f_1` in column two. + let far = 2 * row_count + 5; + let mut map = DenseMap { + rows: vec![vec![false; 3]; far + 1], + columns: 3, + }; + map.rows[1][1] = true; + map.rows[far][1] = true; + map.rows[far][2] = true; + let input_claim = input_claim(¶ms); + let statement = VirtualStatement::new(params, *params.shape(), &map, &input_claim).unwrap(); + let point: Vec<_> = (0..params.shape().log_bits() as u64) + .map(|index| F128::new(7 + index, 11)) + .collect(); + let table = poly::eq_table(&point); + let target = F128::new(13, 17); + let OpeningQuery::InnerProduct { claim } = statement + .transpose_query(OpeningQuery::Mle { point, target }) + .unwrap() + else { + panic!("transposition must produce an inner-product opening"); + }; + assert_eq!(claim.row_weights()[0], table[1] + table[far]); + assert_eq!(claim.row_weights()[1], table[far]); + assert_eq!(claim.target(), target); + } + #[test] fn opening_queries_must_cover_the_configured_virtual_shape() { let map = map(); diff --git a/crates/tests/tests/virtual_prove.rs b/crates/tests/tests/virtual_prove.rs index d93a68d..a8144c2 100644 --- a/crates/tests/tests/virtual_prove.rs +++ b/crates/tests/tests/virtual_prove.rs @@ -24,13 +24,20 @@ type F = field::Fq; struct Map(u8); impl VirtualMap for Map { - fn transpose(&self, weights: &[F128]) -> Result { - if weights.len() < self.h_len() { + fn transpose( + &self, + row_weights: &[F128], + column_weights: &[F128], + ) -> Result { + let row_count = row_weights.len(); + if row_count * column_weights.len() < self.h_len() { return Err(VirtualMapError::WeightCountMismatch); } + let weight = + |index: usize| row_weights[index % row_count] * column_weights[index / row_count]; Ok(TransposedWeights::new( - vec![weights[1] + weights[129], weights[128] + weights[129]], - weights[0], + vec![weight(1) + weight(129), weight(128) + weight(129)], + weight(0), )) }