diff --git a/crates/circuit/src/constraints.rs b/crates/circuit/src/constraints.rs index 5b45f0bc..61d0b6b9 100644 --- a/crates/circuit/src/constraints.rs +++ b/crates/circuit/src/constraints.rs @@ -6,16 +6,17 @@ //! rank-1 constraints `(A z) * (B z) = C z` over that integer witness. Every //! integer coefficient is an arbitrary-precision signed [`BigInt`]. +use num_bigint::BigInt; +use num_traits::{One, Zero}; +use rayon::prelude::*; use std::array; -use std::collections::{BTreeMap, BTreeSet, btree_map::Entry}; +use std::cmp::Ordering; use std::error::Error; use std::fmt::{self, Display}; use std::iter::Sum; +use std::mem; use std::ops::{Add, AddAssign, Mul, Neg, Sub, SubAssign}; -use num_bigint::BigInt; -use num_traits::{One, Zero}; - use crate::witgen::PackedWitness; use crate::{BoolWitness, Circuit, HintResult, PackedBits, ScalarBits, WitnessContext}; @@ -94,12 +95,18 @@ impl SparseMatrix { pub const fn column_count(&self) -> usize { self.columns } +} - fn map_values_with(self, map: &mut impl FnMut(C) -> D) -> SparseMatrix { +impl SparseMatrix { + fn map_values_with(self, map: M) -> SparseMatrix + where + D: Send + Sync, + M: Fn(C) -> D + Send + Sync, + { SparseMatrix { rows: self .rows - .into_iter() + .into_par_iter() .map(|row| SparseRow { entries: row .entries @@ -273,7 +280,7 @@ impl Display for ConstraintMatrixShapeError { impl Error for ConstraintMatrixShapeError {} -impl ConstraintMatrices { +impl ConstraintMatrices { /// Checks that A, B, and C share a shape and consume the assignment /// produced by M. pub fn validate_shape(&self) -> Result<(), ConstraintMatrixShapeError> { @@ -313,12 +320,16 @@ impl ConstraintMatrices { /// Consumes the matrices and maps every A/B/C coefficient. /// /// The Boolean `M` matrix and sparse topology are moved unchanged. - pub fn map_coefficients(self, mut map: impl FnMut(C) -> D) -> ConstraintMatrices { + pub fn map_coefficients(self, map: M) -> ConstraintMatrices + where + D: Send + Sync, + M: Fn(C) -> D + Send + Sync, + { ConstraintMatrices { m: self.m, - a: self.a.map_values_with(&mut map), - b: self.b.map_values_with(&mut map), - c: self.c.map_values_with(&mut map), + a: self.a.map_values_with(&map), + b: self.b.map_values_with(&map), + c: self.c.map_values_with(&map), } } } @@ -424,7 +435,8 @@ fn evaluate_integer_row(row: &SparseRow, witness: &[BigInt]) -> BigInt { #[derive(Clone, Debug, Default, Eq, PartialEq)] pub struct BoolLinearCombination { constant: bool, - witnesses: BTreeSet, + /// Witness indices, sorted and deduplicated. + witnesses: Vec, } impl BoolLinearCombination { @@ -433,25 +445,26 @@ impl BoolLinearCombination { self.constant } - /// Zero-based Boolean witness indices with coefficient one. - pub fn witnesses(&self) -> &BTreeSet { + /// Sorted zero-based Boolean witness indices with coefficient one. + pub fn witnesses(&self) -> &[usize] { &self.witnesses } pub(crate) fn witness(index: usize) -> Self { Self { constant: false, - witnesses: BTreeSet::from([index]), + witnesses: vec![index], } } pub(crate) fn xor(mut self, rhs: Self) -> Self { self.constant ^= rhs.constant; - for witness in rhs.witnesses { - if !self.witnesses.insert(witness) { - self.witnesses.remove(&witness); - } - } + self.witnesses = merge_sorted_vecs( + self.witnesses, + rhs.witnesses, + |lhs, rhs| lhs.cmp(rhs), + |_lhs, _rhs| None, // Drop overlaps + ); self } } @@ -460,7 +473,7 @@ impl From for BoolLinearCombination { fn from(constant: bool) -> Self { Self { constant, - witnesses: BTreeSet::new(), + witnesses: Vec::new(), } } } @@ -473,7 +486,9 @@ impl BoolWitness for BoolLinearCombination { #[derive(Clone, Debug, Eq, PartialEq)] pub struct LinearCombination { constant: BigInt, - witnesses: BTreeMap, + /// Indices of witness elements with nonzero coefficients. + /// Sorted and deduplicated by witness index. + witnesses: Vec<(usize, BigInt)>, } impl LinearCombination { @@ -482,47 +497,31 @@ impl LinearCombination { &self.constant } - /// Nonzero coefficients keyed by zero-based integer witness index. - pub fn witnesses(&self) -> &BTreeMap { + /// Nonzero coefficients sorted by zero-based integer witness index. + pub fn witnesses(&self) -> &[(usize, BigInt)] { &self.witnesses } fn witness(index: usize) -> Self { Self { constant: BigInt::zero(), - witnesses: BTreeMap::from([(index, BigInt::one())]), - } - } - - fn add_term(&mut self, index: usize, coefficient: BigInt) { - if coefficient.is_zero() { - return; - } - match self.witnesses.entry(index) { - Entry::Vacant(entry) => { - entry.insert(coefficient); - } - Entry::Occupied(mut entry) => { - *entry.get_mut() += coefficient; - if entry.get().is_zero() { - entry.remove(); - } - } + witnesses: vec![(index, BigInt::one())], } } fn into_sparse_row(self) -> SparseRow { - let mut entries = - Vec::with_capacity(self.witnesses.len() + usize::from(!self.constant.is_zero())); - if !self.constant.is_zero() { - entries.push((0, self.constant)); + let Self { + constant, + mut witnesses, + } = self; + for (column, _) in &mut witnesses { + *column += 1; } - entries.extend( - self.witnesses - .into_iter() - .map(|(witness, coefficient)| (witness + 1, coefficient)), - ); - SparseRow { entries } + if !constant.is_zero() { + witnesses.reserve_exact(1); + witnesses.insert(0, (0, constant)); + } + SparseRow { entries: witnesses } } } @@ -530,7 +529,7 @@ impl From for LinearCombination { fn from(constant: BigInt) -> Self { Self { constant, - witnesses: BTreeMap::new(), + witnesses: Vec::new(), } } } @@ -549,10 +548,7 @@ impl Add for LinearCombination { type Output = Self; fn add(mut self, rhs: Self) -> Self::Output { - self.constant += rhs.constant; - for (witness, coefficient) in rhs.witnesses { - self.add_term(witness, coefficient); - } + self += rhs; self } } @@ -560,9 +556,19 @@ impl Add for LinearCombination { impl AddAssign for LinearCombination { fn add_assign(&mut self, rhs: Self) { self.constant += rhs.constant; - for (witness, coefficient) in rhs.witnesses { - self.add_term(witness, coefficient); - } + self.witnesses = merge_sorted_vecs( + mem::take(&mut self.witnesses), + rhs.witnesses, + |(lhs_idx, _), (rhs_idx, _)| lhs_idx.cmp(rhs_idx), + |(wit_idx, lhs_coeff), (_, rhs_coeff)| { + let coeff = lhs_coeff + rhs_coeff; + if coeff.is_zero() { + None + } else { + Some((wit_idx, coeff)) + } + }, + ); } } @@ -571,11 +577,9 @@ impl Neg for LinearCombination { fn neg(mut self) -> Self::Output { self.constant = -self.constant; - self.witnesses = self - .witnesses - .into_iter() - .map(|(witness, coefficient)| (witness, -coefficient)) - .collect(); + for (_, coefficient) in &mut self.witnesses { + *coefficient = -mem::take(coefficient); + } self } } @@ -598,15 +602,13 @@ impl Mul for LinearCombination { type Output = Self; fn mul(mut self, rhs: BigInt) -> Self::Output { - self.constant *= rhs.clone(); - self.witnesses = self - .witnesses - .into_iter() - .filter_map(|(witness, coefficient)| { - let coefficient = coefficient * rhs.clone(); - (!coefficient.is_zero()).then_some((witness, coefficient)) - }) - .collect(); + if rhs.is_zero() { + return Self::zero(); + } + self.constant *= &rhs; + for (_, coefficient) in &mut self.witnesses { + *coefficient *= &rhs; + } self } } @@ -617,6 +619,54 @@ impl Sum for LinearCombination { } } +/// Merges two sorted vectors with no duplicates. Compares elements using `cmp` function, +/// and if two elements are equal, element produced by the `merge` function is added instead. +/// +/// Concatenation will only reserve as much extra space as needed. +fn merge_sorted_vecs( + mut lhs: Vec, + mut rhs: Vec, + cmp: impl Fn(&T, &T) -> Ordering, + merge: impl Fn(T, T) -> Option, +) -> Vec { + let (Some(lhs_first), Some(lhs_last)) = (lhs.first(), lhs.last()) else { + return rhs; + }; + let (Some(rhs_first), Some(rhs_last)) = (rhs.first(), rhs.last()) else { + return lhs; + }; + // Disjoint ranges concatenate without a merge. + if cmp(lhs_last, rhs_first).is_lt() { + lhs.reserve_exact(rhs.len()); + lhs.extend(rhs); + return lhs; + } + if cmp(rhs_last, lhs_first).is_lt() { + rhs.reserve_exact(lhs.len()); + rhs.extend(lhs); + return rhs; + } + let mut merged = Vec::with_capacity(lhs.len() + rhs.len()); + let mut lhs = lhs.into_iter().peekable(); + let mut rhs = rhs.into_iter().peekable(); + while let (Some(left), Some(right)) = (lhs.peek(), rhs.peek()) { + match cmp(left, right) { + Ordering::Less => merged.extend(lhs.next()), + Ordering::Greater => merged.extend(rhs.next()), + Ordering::Equal => { + let left = lhs.next().expect("impossible"); + let right = rhs.next().expect("impossible"); + if let Some(new) = merge(left, right) { + merged.push(new); + } + } + } + } + merged.extend(lhs); + merged.extend(rhs); + merged +} + /// Circuit backend that records sparse M/A/B/C matrices without evaluating hints. #[derive(Clone, Debug)] pub struct ConstraintGenerator { @@ -667,6 +717,13 @@ impl ConstraintGenerator { } } + /// Records `value` as the next M row and returns its integer witness index. + fn record_bitz(&mut self, value: BoolLinearCombination) -> usize { + let witness = self.m_rows.len(); + self.m_rows.push(value); + witness + } + /// Finishes generation and materializes the four sparse matrices. pub fn into_matrices(self) -> ConstraintMatrices { let Self { @@ -712,12 +769,20 @@ impl ConstraintGenerator { } fn bool_sparse_row(value: BoolLinearCombination) -> SparseBoolRow { - let mut positions = Vec::with_capacity(value.witnesses.len() + usize::from(value.constant)); - if value.constant { - positions.push(0); + let BoolLinearCombination { + constant, + mut witnesses, + } = value; + for position in &mut witnesses { + *position += 1; + } + if constant { + witnesses.reserve_exact(1); + witnesses.insert(0, 0); + } + SparseBoolRow { + positions: witnesses, } - positions.extend(value.witnesses.into_iter().map(|witness| witness + 1)); - SparseBoolRow { positions } } impl Circuit for ConstraintGenerator { @@ -757,9 +822,29 @@ impl Circuit for ConstraintGenerator { } fn bitz(&mut self, value: BoolLinearCombination) -> LinearCombination { - let witness = self.m_rows.len(); - self.m_rows.push(value); - LinearCombination::witness(witness) + LinearCombination::witness(self.record_bitz(value)) + } + + fn bitz_unsigned( + &mut self, + bits_le: &ScalarBits, + ) -> (LinearCombination, LinearCombination) { + assert!(LOW <= N, "low part cannot be wider than the input"); + // Same layout as the default, one `bitz` per bit with coefficient + // 2^index, but each sum is built once with exact capacity. + let mut witnesses = Vec::with_capacity(N); + for (index, bit) in bits_le.0.iter().enumerate() { + witnesses.push((self.record_bitz(bit.clone()), BigInt::one() << index)); + } + let low = LinearCombination { + constant: BigInt::zero(), + witnesses: witnesses[..LOW].to_vec(), + }; + let full = LinearCombination { + constant: BigInt::zero(), + witnesses, + }; + (full, low) } fn assert_r1c( @@ -788,6 +873,13 @@ mod tests { use super::*; use crate::witgen::Witgen; + fn terms(pairs: &[(usize, i64)]) -> Vec<(usize, BigInt)> { + pairs + .iter() + .map(|&(witness, coefficient)| (witness, BigInt::from(coefficient))) + .collect() + } + #[test] fn materializes_freigen_matrix_conventions_and_checks_witnesses() { let mut generator = ConstraintGenerator::new(2); @@ -886,4 +978,74 @@ mod tests { )) ); } + + #[test] + fn linear_combination_terms_stay_sorted_and_merge_shared_witnesses() { + let term = |index: usize, coefficient: i64| { + LinearCombination::witness(index) * BigInt::from(coefficient) + }; + let exact = |value: &LinearCombination| value.witnesses.capacity() == value.witnesses.len(); + + // Disjoint ranges append on either side without spare capacity. + let ascending = term(1, 1) + term(3, 1); + assert_eq!(ascending.witnesses(), terms(&[(1, 1), (3, 1)])); + assert!(exact(&ascending)); + let descending = term(3, 1) + term(1, 1); + assert_eq!(descending.witnesses(), terms(&[(1, 1), (3, 1)])); + assert!(exact(&descending)); + + // Interleaved ranges merge, shared witnesses sum, cancellations vanish. + let interleaved = (term(0, 1) + term(2, 1)) + (term(1, 1) + term(3, 1) + term(2, 3)); + assert_eq!( + interleaved.witnesses(), + terms(&[(0, 1), (1, 1), (2, 4), (3, 1)]) + ); + let cancelled = (term(1, 1) + term(2, 1)) - term(1, 1); + assert_eq!(cancelled.witnesses(), terms(&[(2, 1)])); + assert!((term(1, 2) - term(1, 2)).is_zero()); + assert!((term(5, 7) * BigInt::zero()).is_zero()); + + let scaled = -(term(1, 2) + LinearCombination::from(BigInt::from(3))) * BigInt::from(5); + assert_eq!(*scaled.constant(), BigInt::from(-15)); + assert_eq!(scaled.witnesses(), terms(&[(1, -10)])); + } + + #[test] + fn bitz_unsigned_lifts_bits_with_powers_of_two_into_exact_rows() { + let mut generator = ConstraintGenerator::new(3); + let bits = ScalarBits(generator.inputs::<3>()); + let (full, low) = generator.bitz_unsigned::<1, 3, 1, 2>(&bits); + + assert_eq!(generator.m_rows.len(), 3); + assert_eq!(generator.m_rows[2].witnesses(), &[2]); + assert!(full.constant().is_zero()); + assert_eq!(full.witnesses(), terms(&[(0, 1), (1, 2), (2, 4)])); + assert_eq!(full.witnesses.capacity(), 3); + assert_eq!(low.witnesses(), terms(&[(0, 1), (1, 2)])); + assert_eq!(low.witnesses.capacity(), 2); + } + + #[test] + fn bool_linear_combination_xor_is_a_sorted_symmetric_difference() { + let bit = BoolLinearCombination::witness; + + let ascending = bit(1).xor(bit(3)); + assert_eq!(ascending.witnesses(), &[1, 3]); + assert_eq!(bit(3).xor(bit(1)).witnesses(), &[1, 3]); + assert_eq!( + ascending.clone().xor(bit(3).xor(bit(5))).witnesses(), + &[1, 5] + ); + assert_eq!( + bit(0).xor(bit(2)).xor(bit(1).xor(bit(3))).witnesses(), + &[0, 1, 2, 3] + ); + + let cancelled = ascending + .clone() + .xor(ascending) + .xor(BoolLinearCombination::from(true)); + assert!(cancelled.witnesses().is_empty()); + assert!(cancelled.constant()); + } } diff --git a/crates/spartan/src/matrix.rs b/crates/spartan/src/matrix.rs index 17e97f6b..8215b6bc 100644 --- a/crates/spartan/src/matrix.rs +++ b/crates/spartan/src/matrix.rs @@ -10,10 +10,13 @@ use num_traits::{Signed, ToPrimitive}; use poly::DenseMultilinearExtension; use rayon::prelude::*; use sha2::{Digest, Sha256}; +use std::sync::LazyLock; use transcript::Encoding; use crate::sumcheck::R1csProductMles; +static FQ_DEFAULT_MODULUS: LazyLock = LazyLock::new(|| BigInt::from(Q100)); + /// Failures while preparing or evaluating Spartan's R1CS matrices. #[derive(Clone, Copy, Debug, Eq, PartialEq)] pub enum SpartanMatrixError { @@ -153,8 +156,8 @@ where /// Reduces a signed integer canonically modulo Q100. pub fn bigint_to_fq(value: &BigInt) -> FqDefault { - let modulus = BigInt::from(Q100); - let mut reduced = value % &modulus; + let modulus = &*FQ_DEFAULT_MODULUS; + let mut reduced = value % modulus; if reduced.is_negative() { reduced += modulus; } @@ -397,7 +400,7 @@ where Ok(evaluation) } -pub(crate) fn r1cs_num_vars( +pub(crate) fn r1cs_num_vars( matrices: &ConstraintMatrices, ) -> Result<(usize, usize), SpartanMatrixError> { matrices