Merge pull request #232 from WizardOfMenlo/recmo/zkwhirv3

zkWHIRv3 part 1: Single commitment base case

Author: recmo

Cargo.toml

@@ -74,6 +74,9 @@ parallel = ["dep:rayon"]
 rayon = ["dep:rayon"]
 asm = ["ark-ff/asm"]
 tracing = ["dep:tracing"]
+# Do not permute evaluation pointsin RS
+# This flag is to be removed after whir_zk supports it.
+rs_in_order = []
 
 [[bench]]
 name = "expand_from_coeff"

benches/expand_from_coeff.rs

@@ -1,5 +1,5 @@
 use divan::{black_box, AllocProfiler, Bencher};
-use whir::algebra::{fields::Field64, ntt};
+use whir::algebra::{fields::Field64, ntt, random_vector};
 
 #[global_allocator]
 static ALLOC: AllocProfiler = AllocProfiler::system();
@@ -10,7 +10,7 @@ static ALLOC: AllocProfiler = AllocProfiler::system();
 //      - RS code expansion factors, and
 //      - interleaved bloc size exponent of 2
 //
-const TEST_CASES: &[(u32, usize, usize)] = &[
+const TEST_CASES: &[(usize, usize, usize)] = &[
     (16, 2, 2),
     (18, 2, 2),
     (20, 2, 3),
@@ -21,20 +21,24 @@ const TEST_CASES: &[(u32, usize, usize)] = &[
 ];
 
 #[divan::bench(args = TEST_CASES)]
-fn interleaved_rs_encode(bencher: Bencher, case: &(u32, usize, usize)) {
+fn interleaved_rs_encode(bencher: Bencher, case: &(usize, usize, usize)) {
     bencher
         .with_inputs(|| {
             let (exp, expansion, coset_sz) = *case;
-
-            let size = 1 << exp;
-            let coeffs: Vec<_> = (0..size).map(Field64::from).collect();
+            let message_length = 1 << (exp - coset_sz);
+            let num_messages = 1 << coset_sz;
+            let mut rng = ark_std::rand::thread_rng();
+            let coeffs: Vec<Vec<Field64>> = (0..num_messages)
+                .map(|_| random_vector(&mut rng, message_length))
+                .collect();
             (coeffs, expansion, coset_sz)
         })
-        .bench_values(|(coeffs, expansion, coset_sz)| {
+        .bench_values(|(coeffs, expansion, _coset_sz)| {
+            let coeffs_refs = coeffs.iter().map(|v| v.as_slice()).collect::<Vec<_>>();
             black_box(ntt::interleaved_rs_encode(
-                &[&coeffs],
-                (coeffs.len() >> coset_sz) * expansion,
-                1 << coset_sz,
+                &coeffs_refs,
+                &[],
+                coeffs[0].len() * expansion,
             ))
         });
 }

proptest-regressions/protocols/basecase.txt

@@ -0,0 +1,7 @@
+# Seeds for failure cases proptest has generated in the past. It is
+# automatically read and these particular cases re-run before any
+# novel cases are generated.
+#
+# It is recommended to check this file in to source control so that
+# everyone who runs the test benefits from these saved cases.
+cc 6fd50adb2eacc5a44429f8aeb8bae55b4a05ae2ac74432e184fa91729f375276 # shrinks to seed = 0, config = Config { commit: Config { embedding: Identity { field: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 } }, num_vectors: 1, vector_size: 0, mask_length: 0, codeword_length: 1, interleaving_depth: 1, matrix_commit: Config { element_type: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 }, num_cols: 1, leaf_hash_id: 09459020f451874a1b399819d079632cc0f9263b1486c423173c6e15d8e2d61d, merkle_tree: Config { num_leaves: 1, layers: [] } }, johnson_slack: 0.0, in_domain_samples: 0, out_domain_samples: 0, deduplicate_in_domain: false }, sumcheck: Config { field: FieldInfo { characteristic: [255, 255, 255, 255, 0, 0, 0, 1], extension_degree: 1 }, initial_size: 0, round_pow: Config { hash_id: 03e01749ebcc0477924254eb482066b864a8dd4d77252464ca6f5b6f5cc05b4c, threshold: 18446744073709551615 }, num_rounds: 0 }, masked: false }

src/algebra/fields.rs

@@ -65,6 +65,8 @@ impl<F: Field> TypeInfo for F {
 #[derive(MontConfig)]
 #[modulus = "21888242871839275222246405745257275088548364400416034343698204186575808495617"]
 #[generator = "5"]
+#[small_subgroup_base = "3"]
+#[small_subgroup_power = "2"]
 pub struct BN254Config;
 pub type Field256 = Fp256<MontBackend<BN254Config, 4>>;
 
@@ -83,6 +85,8 @@ pub type Field128 = Fp128<MontBackend<FrConfig128, 2>>;
 #[derive(MontConfig)]
 #[modulus = "18446744069414584321"]
 #[generator = "7"]
+#[small_subgroup_base = "15"]
+#[small_subgroup_power = "1"]
 pub struct FConfig64;
 pub type Field64 = Fp64<MontBackend<FConfig64, 1>>;
 

src/algebra/mod.rs

@@ -7,6 +7,7 @@ pub mod ntt;
 pub mod sumcheck;
 
 use ark_ff::{AdditiveGroup, Field};
+use ark_std::rand::{distributions::Standard, prelude::Distribution, Rng};
 pub use multilinear::{eval_eq, mixed_multilinear_extend, multilinear_extend};
 pub use multilinear_point::MultilinearPoint;
 #[cfg(feature = "parallel")]
@@ -15,6 +16,15 @@ use rayon::prelude::*;
 use self::embedding::Embedding;
 #[cfg(feature = "parallel")]
 use crate::utils::workload_size;
+use crate::utils::zip_strict;
+
+pub fn random_vector<R, F: Field>(rng: &mut R, length: usize) -> Vec<F>
+where
+    R: Rng,
+    Standard: Distribution<F>,
+{
+    (0..length).map(|_| rng.gen()).collect::<Vec<F>>()
+}
 
 pub fn geometric_sequence<F: Field>(base: F, length: usize) -> Vec<F> {
     let mut result = Vec::with_capacity(length);
@@ -51,6 +61,21 @@ pub fn lift<M: Embedding>(embedding: &M, source: &[M::Source]) -> Vec<M::Target>
     result
 }
 
+pub fn scalar_mul<F: Field>(vector: &mut [F], weight: F) {
+    for value in vector.iter_mut() {
+        *value *= weight;
+    }
+}
+
+/// Scalar mul add into new vector.
+///
+/// Returns r such that r[i] = a[i] + c · b[i].
+pub fn scalar_mul_add_new<F: Field>(a: &[F], c: F, b: &[F]) -> Vec<F> {
+    zip_strict(a.iter(), b.iter())
+        .map(|(a, b)| *a + c * *b)
+        .collect::<Vec<F>>()
+}
+
 pub fn scalar_mul_add<F: Field>(accumulator: &mut [F], weight: F, vector: &[F]) {
     mixed_scalar_mul_add(
         &embedding::Identity::<F>::new(),

src/algebra/multilinear.rs

@@ -1,4 +1,4 @@
-use ark_ff::Field;
+use ark_ff::{AdditiveGroup, Field};
 
 use crate::algebra::embedding::{Embedding, Identity};
 
@@ -8,75 +8,146 @@ pub fn multilinear_extend<F: Field>(evals: &[F], point: &[F]) -> F {
 }
 
 /// Evaluate the multi-linear extension of `evals` in `point`.
+///
+/// Supports implicit zero-padding: when `evals.len() < 1 << point.len()`,
+/// the missing tail entries are treated as zeros.
+#[allow(clippy::too_many_lines)]
 pub fn mixed_multilinear_extend<M: Embedding>(
     embedding: &M,
     evals: &[M::Source],
     point: &[M::Target],
 ) -> M::Target {
-    assert_eq!(evals.len(), 1 << point.len());
-
-    // Helper to compute (a + (b - a) * c) efficiently with a, b in source field.
-    let mixed = |a, b, c| embedding.mixed_add(embedding.mixed_mul(c, b - a), a);
-
-    match point {
-        [] => embedding.map(evals[0]),
-        [x] => mixed(evals[0], evals[1], *x),
-        [x0, x1] => {
-            let a0 = mixed(evals[0], evals[1], *x1);
-            let a1 = mixed(evals[2], evals[3], *x1);
-            a0 + (a1 - a0) * *x0
+    #[inline]
+    fn eval_exact<M: Embedding>(
+        embedding: &M,
+        evals: &[M::Source],
+        point: &[M::Target],
+    ) -> M::Target {
+        debug_assert_eq!(evals.len(), 1 << point.len());
+
+        // Helper to compute (a + (b - a) * c) efficiently with a, b in source field.
+        let mixed = |a, b, c| embedding.mixed_add(embedding.mixed_mul(c, b - a), a);
+
+        match point {
+            [] => embedding.map(evals[0]),
+            [x] => mixed(evals[0], evals[1], *x),
+            [x0, x1] => {
+                let a0 = mixed(evals[0], evals[1], *x1);
+                let a1 = mixed(evals[2], evals[3], *x1);
+                a0 + (a1 - a0) * *x0
+            }
+            [x0, x1, x2] => {
+                let a00 = mixed(evals[0], evals[1], *x2);
+                let a01 = mixed(evals[2], evals[3], *x2);
+                let a10 = mixed(evals[4], evals[5], *x2);
+                let a11 = mixed(evals[6], evals[7], *x2);
+                let a0 = a00 + (a01 - a00) * *x1;
+                let a1 = a10 + (a11 - a10) * *x1;
+                a0 + (a1 - a0) * *x0
+            }
+            [x0, x1, x2, x3] => {
+                let a000 = mixed(evals[0], evals[1], *x3);
+                let a001 = mixed(evals[2], evals[3], *x3);
+                let a010 = mixed(evals[4], evals[5], *x3);
+                let a011 = mixed(evals[6], evals[7], *x3);
+                let a100 = mixed(evals[8], evals[9], *x3);
+                let a101 = mixed(evals[10], evals[11], *x3);
+                let a110 = mixed(evals[12], evals[13], *x3);
+                let a111 = mixed(evals[14], evals[15], *x3);
+                let a00 = a000 + (a001 - a000) * *x2;
+                let a01 = a010 + (a011 - a010) * *x2;
+                let a10 = a100 + (a101 - a100) * *x2;
+                let a11 = a110 + (a111 - a110) * *x2;
+                let a0 = a00 + (a01 - a00) * *x1;
+                let a1 = a10 + (a11 - a10) * *x1;
+                a0 + (a1 - a0) * *x0
+            }
+            [x, tail @ ..] => {
+                let (f0, f1) = evals.split_at(evals.len() / 2);
+                #[cfg(not(feature = "parallel"))]
+                let (f0, f1) = (
+                    eval_exact(embedding, f0, tail),
+                    eval_exact(embedding, f1, tail),
+                );
+
+                #[cfg(feature = "parallel")]
+                let (f0, f1) = {
+                    use crate::utils::workload_size;
+                    if evals.len() > workload_size::<M::Source>() {
+                        rayon::join(
+                            || eval_exact(embedding, f0, tail),
+                            || eval_exact(embedding, f1, tail),
+                        )
+                    } else {
+                        (
+                            eval_exact(embedding, f0, tail),
+                            eval_exact(embedding, f1, tail),
+                        )
+                    }
+                };
+
+                f0 + (f1 - f0) * *x
+            }
         }
-        [x0, x1, x2] => {
-            let a00 = mixed(evals[0], evals[1], *x2);
-            let a01 = mixed(evals[2], evals[3], *x2);
-            let a10 = mixed(evals[4], evals[5], *x2);
-            let a11 = mixed(evals[6], evals[7], *x2);
-            let a0 = a00 + (a01 - a00) * *x1;
-            let a1 = a10 + (a11 - a10) * *x1;
-            a0 + (a1 - a0) * *x0
+    }
+
+    #[inline]
+    fn eval_partial<M: Embedding>(
+        embedding: &M,
+        evals: &[M::Source],
+        point: &[M::Target],
+    ) -> M::Target {
+        let size = 1 << point.len();
+        debug_assert!(evals.len() <= size);
+        if evals.is_empty() {
+            return M::Target::ZERO;
         }
-        [x0, x1, x2, x3] => {
-            let a000 = mixed(evals[0], evals[1], *x3);
-            let a001 = mixed(evals[2], evals[3], *x3);
-            let a010 = mixed(evals[4], evals[5], *x3);
-            let a011 = mixed(evals[6], evals[7], *x3);
-            let a100 = mixed(evals[8], evals[9], *x3);
-            let a101 = mixed(evals[10], evals[11], *x3);
-            let a110 = mixed(evals[12], evals[13], *x3);
-            let a111 = mixed(evals[14], evals[15], *x3);
-            let a00 = a000 + (a001 - a000) * *x2;
-            let a01 = a010 + (a011 - a010) * *x2;
-            let a10 = a100 + (a101 - a100) * *x2;
-            let a11 = a110 + (a111 - a110) * *x2;
-            let a0 = a00 + (a01 - a00) * *x1;
-            let a1 = a10 + (a11 - a10) * *x1;
-            a0 + (a1 - a0) * *x0
+        if evals.len() == size {
+            return eval_exact(embedding, evals, point);
         }
-        [x, tail @ ..] => {
-            let (f0, f1) = evals.split_at(evals.len() / 2);
-            #[cfg(not(feature = "parallel"))]
-            let (f0, f1) = (
-                mixed_multilinear_extend(embedding, f0, tail),
-                mixed_multilinear_extend(embedding, f1, tail),
-            );
-            #[cfg(feature = "parallel")]
-            let (f0, f1) = {
-                use crate::utils::workload_size;
-                if evals.len() > workload_size::<M::Source>() {
-                    rayon::join(
-                        || mixed_multilinear_extend(embedding, f0, tail),
-                        || mixed_multilinear_extend(embedding, f1, tail),
-                    )
-                } else {
-                    (
-                        mixed_multilinear_extend(embedding, f0, tail),
-                        mixed_multilinear_extend(embedding, f1, tail),
-                    )
+
+        match point {
+            [] => embedding.map(evals[0]),
+            [x, tail @ ..] => {
+                let half = size / 2;
+
+                // Only low half has data; high half is all implicit zeros.
+                if evals.len() <= half {
+                    let f0 = eval_partial(embedding, evals, tail);
+                    return f0 * (M::Target::ONE - *x);
                 }
-            };
-            f0 + (f1 - f0) * *x
+
+                // Low subtree is exact/full, high subtree is partial.
+                let (low, high) = evals.split_at(half);
+
+                #[cfg(not(feature = "parallel"))]
+                let (f0, f1) = (
+                    eval_exact(embedding, low, tail),
+                    eval_partial(embedding, high, tail),
+                );
+
+                #[cfg(feature = "parallel")]
+                let (f0, f1) = {
+                    use crate::utils::workload_size;
+                    if evals.len() > workload_size::<M::Source>() {
+                        rayon::join(
+                            || eval_exact(embedding, low, tail),
+                            || eval_partial(embedding, high, tail),
+                        )
+                    } else {
+                        (
+                            eval_exact(embedding, low, tail),
+                            eval_partial(embedding, high, tail),
+                        )
+                    }
+                };
+
+                f0 + (f1 - f0) * *x
+            }
         }
     }
+
+    eval_partial(embedding, evals, point)
 }
 
 /// Accumulates a scaled evaluation of the equality function.
@@ -88,7 +159,7 @@ pub fn mixed_multilinear_extend<M: Embedding>(
 /// eq(X) = ∏ (1 - X_i + 2X_i z_i)
 /// ```
 ///
-/// where `z_i` are the  points.
+/// where `z_i` are the points.
 pub fn eval_eq<F: Field>(accumulator: &mut [F], point: &[F], scalar: F) {
     assert_eq!(accumulator.len(), 1 << point.len());
     if let [x0, xs @ ..] = point {
@@ -110,3 +181,45 @@ pub fn eval_eq<F: Field>(accumulator: &mut [F], point: &[F], scalar: F) {
         accumulator[0] += scalar;
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use ark_std::rand::{rngs::StdRng, SeedableRng};
+    use proptest::proptest;
+
+    use super::*;
+    use crate::algebra::{random_vector, sumcheck::tests::zero_pad};
+
+    pub type F = crate::algebra::fields::Field64;
+
+    #[test]
+    fn test_multilinear_zero_extend() {
+        proptest!(|(seed:u64, length in 0_usize..(1 << 14))| {
+            let mut rng = StdRng::seed_from_u64(seed);
+            let vector: Vec<F> = random_vector(&mut rng, length);
+            let extended_vector = zero_pad(&vector);
+            let num_variables = length.next_power_of_two().trailing_zeros();
+            let point = random_vector(&mut rng, num_variables as usize);
+            assert_eq!(
+                multilinear_extend(&vector, &point),
+                multilinear_extend(&extended_vector, &point)
+            );
+        });
+    }
+
+    #[test]
+    fn test_multilinear_extra_variables() {
+        proptest!(|(seed:u64, length in 0_usize..(1 << 10), excess_variables in 0_usize..3)| {
+            let mut rng = StdRng::seed_from_u64(seed);
+            let vector: Vec<F> = random_vector(&mut rng, length);
+            let num_variables = length.next_power_of_two().trailing_zeros() as usize + excess_variables;
+            let point = random_vector(&mut rng, num_variables);
+            let mut extended_vector = vector.clone();
+            extended_vector.resize(1 << num_variables, F::ZERO);
+            assert_eq!(
+                multilinear_extend(&vector, &point),
+                multilinear_extend(&extended_vector, &point)
+            );
+        });
+    }
+}