Merge pull request #2553 from o1-labs/tests/mvpoly-gen-2542

MVPoly tests generalization finalization.

mvpoly/src/lib.rs

@@ -109,7 +109,8 @@ pub trait MVPoly<F: PrimeField, const N: usize, const D: usize>:
         u2: F,
     ) -> HashMap<usize, F>;
 
-    fn modify_monomial_with_scalar(&mut self, scalar: F);
+    /// Modify the monomial in the polynomial to the new value `coeff`.
+    fn modify_monomial(&mut self, exponents: [usize; N], coeff: F);
 
     /// Return true if the multi-variate polynomial is multilinear, i.e. if each
     /// variable in each monomial is of maximum degree 1.

mvpoly/src/monomials.rs

@@ -346,15 +346,6 @@ impl<const N: usize, const D: usize, F: PrimeField> Zero for Sparse<F, N, D> {
     }
 }
 
-impl<const N: usize, const D: usize, F: PrimeField> Sparse<F, N, D> {
-    pub fn modify_monomial(&mut self, exponents: [usize; N], coeff: F) {
-        self.monomials
-            .entry(exponents)
-            .and_modify(|c| *c = coeff)
-            .or_insert(coeff);
-    }
-}
-
 impl<const N: usize, const D: usize, F: PrimeField> MVPoly<F, N, D> for Sparse<F, N, D> {
     /// Returns the degree of the polynomial.
     ///
@@ -565,8 +556,11 @@ impl<const N: usize, const D: usize, F: PrimeField> MVPoly<F, N, D> for Sparse<F
         cross_terms_by_powers_of_r
     }
 
-    fn modify_monomial_with_scalar(&mut self, scalar: F) {
-        self.modify_monomial([0; N], scalar);
+    fn modify_monomial(&mut self, exponents: [usize; N], coeff: F) {
+        self.monomials
+            .entry(exponents)
+            .and_modify(|c| *c = coeff)
+            .or_insert(coeff);
     }
 
     fn is_multilinear(&self) -> bool {

mvpoly/src/pbt.rs

@@ -18,7 +18,7 @@
 
 use crate::MVPoly;
 use ark_ff::PrimeField;
-use rand::Rng;
+use rand::{seq::SliceRandom, Rng};
 use std::ops::Neg;
 
 pub fn test_mul_by_one<F: PrimeField, const N: usize, const D: usize, T: MVPoly<F, N, D>>() {
@@ -158,7 +158,7 @@ pub fn test_mul_by_scalar<F: PrimeField, const N: usize, const D: usize, T: MVPo
     let p1 = unsafe { T::random(&mut rng, None) };
     let mut p2 = T::zero();
     let c = F::rand(&mut rng);
-    p2.modify_monomial_with_scalar(c);
+    p2.modify_monomial([0; N], c);
     assert_eq!(p2 * p1.clone(), p1.clone().mul_by_scalar(c));
 }
 
@@ -369,3 +369,163 @@ pub fn test_is_zero<F: PrimeField, const N: usize, const D: usize, T: MVPoly<F,
     let p2 = unsafe { T::random(&mut rng, None) };
     assert!(!p2.is_zero());
 }
+
+pub fn test_homogeneous_eval<F: PrimeField, const N: usize, const D: usize, T: MVPoly<F, N, D>>() {
+    let mut rng = o1_utils::tests::make_test_rng(None);
+    let random_eval = std::array::from_fn(|_| F::rand(&mut rng));
+    let u = F::rand(&mut rng);
+
+    // Homogeneous form is u^2
+    let p1 = T::one();
+    let homogenous_eval = p1.homogeneous_eval(&random_eval, u);
+    assert_eq!(homogenous_eval, u * u);
+
+    let mut p2 = T::zero();
+    let mut exp1 = [0; N];
+    exp1[0] = 1;
+    p2.add_monomial(exp1, F::one());
+    let homogenous_eval = p2.homogeneous_eval(&random_eval, u);
+    assert_eq!(homogenous_eval, random_eval[0] * u);
+
+    let mut p3 = T::zero();
+    let mut exp2 = [0; N];
+    exp2[1] = 1;
+    p3.add_monomial(exp2, F::one());
+    let homogenous_eval = p3.homogeneous_eval(&random_eval, u);
+    assert_eq!(homogenous_eval, random_eval[1] * u);
+
+    let mut p4 = T::zero();
+    let mut exp3 = [0; N];
+    exp3[0] = 1;
+    exp3[1] = 1;
+    p4.add_monomial(exp3, F::one());
+    let homogenous_eval = p4.homogeneous_eval(&random_eval, u);
+    assert_eq!(homogenous_eval, random_eval[0] * random_eval[1]);
+
+    let mut p5 = T::zero();
+    let mut exp4 = [0; N];
+    exp4[0] = 2;
+    p5.add_monomial(exp4, F::one());
+    let homogenous_eval = p5.homogeneous_eval(&random_eval, u);
+    assert_eq!(homogenous_eval, random_eval[0] * random_eval[0]);
+
+    let mut p6 = T::zero();
+    let mut exp5a = [0; N];
+    let mut exp5b = [0; N];
+    exp5a[1] = 2;
+    exp5b[0] = 2;
+    p6.add_monomial(exp5a, F::one());
+    p6.add_monomial(exp5b, F::one());
+    let homogenous_eval = p6.homogeneous_eval(&random_eval, u);
+    assert_eq!(
+        homogenous_eval,
+        random_eval[1] * random_eval[1] + random_eval[0] * random_eval[0]
+    );
+
+    let mut p7 = T::zero();
+    let mut exp6a = [0; N];
+    let mut exp6b = [0; N];
+    let mut exp6c = [0; N];
+    exp6a[1] = 2;
+    exp6b[0] = 2;
+    exp6c[0] = 1;
+    p7.add_monomial(exp6a, F::one());
+    p7.add_monomial(exp6b, F::one());
+    p7.add_monomial(exp6c, F::one());
+    p7.add_monomial([0; N], F::from(42u32));
+    let homogenous_eval = p7.homogeneous_eval(&random_eval, u);
+    assert_eq!(
+        homogenous_eval,
+        random_eval[1] * random_eval[1]
+            + random_eval[0] * random_eval[0]
+            + u * random_eval[0]
+            + u * u * F::from(42u32)
+    );
+}
+
+pub fn test_add_monomial<F: PrimeField, const N: usize, const D: usize, T: MVPoly<F, N, D>>() {
+    let mut rng = o1_utils::tests::make_test_rng(None);
+
+    // Adding constant monomial one to zero
+    let mut p1 = T::zero();
+    p1.add_monomial([0; N], F::one());
+    assert_eq!(p1, T::one());
+
+    // Adding random constant monomial one to zero
+    let mut p2 = T::zero();
+    let random_c = F::rand(&mut rng);
+    p2.add_monomial([0; N], random_c);
+    assert_eq!(p2, T::from(random_c));
+
+    let mut p3 = T::zero();
+    let random_c1 = F::rand(&mut rng);
+    let random_c2 = F::rand(&mut rng);
+    // X1 + X2
+    let mut exp1 = [0; N];
+    let mut exp2 = [0; N];
+    exp1[0] = 1;
+    exp2[1] = 1;
+    p3.add_monomial(exp1, random_c1);
+    p3.add_monomial(exp2, random_c2);
+
+    let random_eval = std::array::from_fn(|_| F::rand(&mut rng));
+    let eval_p3 = p3.eval(&random_eval);
+    let exp_eval_p3 = random_c1 * random_eval[0] + random_c2 * random_eval[1];
+    assert_eq!(eval_p3, exp_eval_p3);
+
+    let mut p4 = T::zero();
+    let random_c1 = F::rand(&mut rng);
+    let random_c2 = F::rand(&mut rng);
+    // X1^2 + X2^2
+    let mut exp1 = [0; N];
+    let mut exp2 = [0; N];
+    exp1[0] = 2;
+    exp2[1] = 2;
+    p4.add_monomial(exp1, random_c1);
+    p4.add_monomial(exp2, random_c2);
+    let eval_p4 = p4.eval(&random_eval);
+    let exp_eval_p4 =
+        random_c1 * random_eval[0] * random_eval[0] + random_c2 * random_eval[1] * random_eval[1];
+    assert_eq!(eval_p4, exp_eval_p4);
+}
+
+pub fn test_is_multilinear<F: PrimeField, const N: usize, const D: usize, T: MVPoly<F, N, D>>() {
+    let mut rng = o1_utils::tests::make_test_rng(None);
+
+    // Test with zero polynomial
+    let p1 = T::zero();
+    assert!(p1.is_multilinear());
+
+    // Test with a constant polynomial
+    let c = F::rand(&mut rng);
+    let p2 = T::from(c);
+    assert!(p2.is_multilinear());
+
+    // Test with a polynomial with one variable having a linear monomial
+    {
+        let mut p = T::zero();
+        let c = F::rand(&mut rng);
+        let idx = rng.gen_range(0..N);
+        let monomials_exponents = std::array::from_fn(|i| if i == idx { 1 } else { 0 });
+        p.add_monomial(monomials_exponents, c);
+        assert!(p.is_multilinear());
+    }
+
+    // Test with a multilinear polynomial with random variables
+    {
+        let mut p = T::zero();
+        let c = F::rand(&mut rng);
+        let nb_var = rng.gen_range(0..D);
+        let mut monomials_exponents: [usize; N] =
+            std::array::from_fn(|i| if i <= nb_var { 1 } else { 0 });
+        monomials_exponents.shuffle(&mut rng);
+        p.add_monomial(monomials_exponents, c);
+        assert!(p.is_multilinear());
+    }
+
+    // Test with a random polynomial (very unlikely to be multilinear)
+    {
+        let p = unsafe { T::random(&mut rng, None) };
+        assert!(!p.is_multilinear());
+    }
+}

mvpoly/src/prime.rs

@@ -449,8 +449,18 @@ impl<F: PrimeField, const N: usize, const D: usize> MVPoly<F, N, D> for Dense<F,
         unimplemented!()
     }
 
-    fn modify_monomial_with_scalar(&mut self, scalar: F) {
-        self[0] = scalar;
+    fn modify_monomial(&mut self, exponents: [usize; N], coeff: F) {
+        let mut prime_gen = PrimeNumberGenerator::new();
+        let primes = prime_gen.get_first_nth_primes(N);
+        let index = exponents
+            .iter()
+            .zip(primes.iter())
+            .fold(1, |acc, (exp, &prime)| acc * prime.pow(*exp as u32));
+        if let Some(pos) = self.normalized_indices.iter().position(|&x| x == index) {
+            self.coeff[pos] = coeff;
+        } else {
+            panic!("Exponent combination out of bounds for the given polynomial degree and number of variables.");
+        }
     }
 
     fn is_multilinear(&self) -> bool {

mvpoly/tests/monomials.rs

@@ -1,7 +1,6 @@
 use ark_ff::{Field, One, UniformRand, Zero};
 use mina_curves::pasta::Fp;
 use mvpoly::{monomials::Sparse, MVPoly};
-use rand::{seq::SliceRandom, Rng};
 
 #[test]
 fn test_mul_by_one() {
@@ -209,101 +208,12 @@ fn test_is_zero() {
 
 #[test]
 fn test_homogeneous_eval() {
-    let mut rng = o1_utils::tests::make_test_rng(None);
-    let random_eval = std::array::from_fn(|_| Fp::rand(&mut rng));
-    let u = Fp::rand(&mut rng);
-    // Homogeneous form is u^2
-    let p1 = Sparse::<Fp, 4, 2>::one();
-    let homogenous_eval = p1.homogeneous_eval(&random_eval, u);
-    assert_eq!(homogenous_eval, u * u);
-
-    let mut p2 = Sparse::<Fp, 4, 2>::zero();
-    // X1
-    p2.add_monomial([1, 0, 0, 0], Fp::one());
-    let homogenous_eval = p2.homogeneous_eval(&random_eval, u);
-    assert_eq!(homogenous_eval, random_eval[0] * u);
-
-    let mut p3 = Sparse::<Fp, 4, 2>::zero();
-    // X2
-    p3.add_monomial([0, 1, 0, 0], Fp::one());
-    let homogenous_eval = p3.homogeneous_eval(&random_eval, u);
-    assert_eq!(homogenous_eval, random_eval[1] * u);
-
-    let mut p4 = Sparse::<Fp, 4, 2>::zero();
-    // X1 * X2
-    p4.add_monomial([1, 1, 0, 0], Fp::one());
-    let homogenous_eval = p4.homogeneous_eval(&random_eval, u);
-    assert_eq!(homogenous_eval, random_eval[0] * random_eval[1]);
-
-    let mut p5 = Sparse::<Fp, 4, 2>::zero();
-    // X1^2
-    p5.add_monomial([2, 0, 0, 0], Fp::one());
-    let homogenous_eval = p5.homogeneous_eval(&random_eval, u);
-    assert_eq!(homogenous_eval, random_eval[0] * random_eval[0]);
-
-    let mut p6 = Sparse::<Fp, 4, 2>::zero();
-    // X2^2 + X1^2
-    p6.add_monomial([0, 2, 0, 0], Fp::one());
-    p6.add_monomial([2, 0, 0, 0], Fp::one());
-    let homogenous_eval = p6.homogeneous_eval(&random_eval, u);
-    assert_eq!(
-        homogenous_eval,
-        random_eval[1] * random_eval[1] + random_eval[0] * random_eval[0]
-    );
-
-    let mut p7 = Sparse::<Fp, 4, 2>::zero();
-    // X2^2 + X1^2 + X1 + 42
-    p7.add_monomial([0, 2, 0, 0], Fp::one());
-    p7.add_monomial([2, 0, 0, 0], Fp::one());
-    p7.add_monomial([1, 0, 0, 0], Fp::one());
-    p7.add_monomial([0, 0, 0, 0], Fp::from(42));
-    let homogenous_eval = p7.homogeneous_eval(&random_eval, u);
-    assert_eq!(
-        homogenous_eval,
-        random_eval[1] * random_eval[1]
-            + random_eval[0] * random_eval[0]
-            + u * random_eval[0]
-            + u * u * Fp::from(42)
-    );
+    mvpoly::pbt::test_homogeneous_eval::<Fp, 4, 2, Sparse<Fp, 4, 2>>();
 }
 
 #[test]
 fn test_add_monomial() {
-    let mut rng = o1_utils::tests::make_test_rng(None);
-
-    // Adding constant monomial one to zero
-    let mut p1 = Sparse::<Fp, 4, 2>::zero();
-    p1.add_monomial([0, 0, 0, 0], Fp::one());
-    assert_eq!(p1, Sparse::<Fp, 4, 2>::one());
-
-    // Adding random constant monomial one to zero
-    let mut p2 = Sparse::<Fp, 4, 2>::zero();
-    let random_c = Fp::rand(&mut rng);
-    p2.add_monomial([0, 0, 0, 0], random_c);
-    assert_eq!(p2, Sparse::<Fp, 4, 2>::from(random_c));
-
-    let mut p3 = Sparse::<Fp, 4, 2>::zero();
-    let random_c1 = Fp::rand(&mut rng);
-    let random_c2 = Fp::rand(&mut rng);
-    // X1 + X2
-    p3.add_monomial([1, 0, 0, 0], random_c1);
-    p3.add_monomial([0, 1, 0, 0], random_c2);
-    let random_eval = std::array::from_fn(|_| Fp::rand(&mut rng));
-    let eval_p3 = p3.eval(&random_eval);
-    let exp_eval_p3 = random_c1 * random_eval[0] + random_c2 * random_eval[1];
-    assert_eq!(eval_p3, exp_eval_p3);
-
-    let mut p4 = Sparse::<Fp, 4, 2>::zero();
-    let random_c1 = Fp::rand(&mut rng);
-    let random_c2 = Fp::rand(&mut rng);
-    // X1^2 + X2^2
-    p4.add_monomial([2, 0, 0, 0], random_c1);
-    p4.add_monomial([0, 2, 0, 0], random_c2);
-    let random_eval = std::array::from_fn(|_| Fp::rand(&mut rng));
-    let eval_p4 = p4.eval(&random_eval);
-    let exp_eval_p4 =
-        random_c1 * random_eval[0] * random_eval[0] + random_c2 * random_eval[1] * random_eval[1];
-    assert_eq!(eval_p4, exp_eval_p4);
+    mvpoly::pbt::test_add_monomial::<Fp, 4, 2, Sparse<Fp, 4, 2>>();
 }
 
 #[test]
@@ -502,37 +412,7 @@ fn test_mvpoly_compute_cross_terms_degree_seven() {
 
 #[test]
 fn test_is_multilinear() {
-    let mut rng = o1_utils::tests::make_test_rng(None);
-    let p1 = Sparse::<Fp, 6, 2>::zero();
-    assert!(p1.is_multilinear());
-
-    let c = Fp::rand(&mut rng);
-    let p2 = Sparse::<Fp, 6, 2>::from(c);
-    assert!(p2.is_multilinear());
-
-    {
-        let mut p = Sparse::<Fp, 6, 3>::zero();
-        let c = Fp::rand(&mut rng);
-        let idx = rng.gen_range(0..6);
-        let monomials_exponents = std::array::from_fn(|i| if i == idx { 1 } else { 0 });
-        p.add_monomial(monomials_exponents, c);
-        assert!(p.is_multilinear());
-    }
-
-    {
-        let mut p = Sparse::<Fp, 6, 4>::zero();
-        let c = Fp::rand(&mut rng);
-        let nb_var = rng.gen_range(0..4);
-        let mut monomials_exponents: [usize; 6] =
-            std::array::from_fn(|i| if i <= nb_var { 1 } else { 0 });
-        monomials_exponents.shuffle(&mut rng);
-        p.add_monomial(monomials_exponents, c);
-        assert!(p.is_multilinear());
-    }
-
-    // Very unlikely to have a random polynomial being multilinear
-    let p3 = unsafe { Sparse::<Fp, 6, 4>::random(&mut rng, None) };
-    assert!(!p3.is_multilinear());
+    mvpoly::pbt::test_is_multilinear::<Fp, 6, 2, Sparse<Fp, 6, 2>>();
 }
 
 #[test]