Merge pull request #2531 from o1-labs/dw/add-monomial-mvpoly-trait

MVPoly: add_monomial in trait

mvpoly/src/lib.rs

@@ -76,4 +76,11 @@ pub trait MVPoly<F: PrimeField, const N: usize, const D: usize>:
     /// Evaluate the polynomial at the vector point `x` and the extra variable
     /// `u` using its homogeneous form of degree D.
     fn homogeneous_eval(&self, x: &[F; N], u: F) -> F;
+
+    /// Add the monomial `coeff * x_1^{e_1} * ... * x_N^{e_N}` to the
+    /// polynomial, where `e_i` are the values given by the array `exponents`.
+    ///
+    /// For instance, to add the monomial `3 * x_1^2 * x_2^3` to the polynomial,
+    /// one would call `add_monomial([2, 3], 3)`.
+    fn add_monomial(&mut self, exponents: [usize; N], coeff: F);
 }

mvpoly/src/monomials.rs

@@ -352,13 +352,6 @@ impl<const N: usize, const D: usize, F: PrimeField> Sparse<F, N, D> {
             .and_modify(|c| *c = coeff)
             .or_insert(coeff);
     }
-
-    pub fn add_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> {
@@ -468,6 +461,13 @@ impl<const N: usize, const D: usize, F: PrimeField> MVPoly<F, N, D> for Sparse<F
             })
             .sum()
     }
+
+    fn add_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> From<prime::Dense<F, N, D>>

mvpoly/src/prime.rs

@@ -423,6 +423,21 @@ impl<F: PrimeField, const N: usize, const D: usize> MVPoly<F, N, D> for Dense<F,
                 acc + c * monomial * u.pow([u_power as u64])
             })
     }
+
+    fn add_monomial(&mut self, exponents: [usize; N], coeff: F) {
+        let mut prime_gen = PrimeNumberGenerator::new();
+        let primes = prime_gen.get_first_nth_primes(N);
+        let normalized_index = exponents
+            .iter()
+            .zip(primes.iter())
+            .fold(1, |acc, (d, p)| acc * p.pow(*d as u32));
+        let inv_idx = self
+            .normalized_indices
+            .iter()
+            .position(|&x| x == normalized_index)
+            .unwrap();
+        self.coeff[inv_idx] += coeff;
+    }
 }
 
 impl<F: PrimeField, const N: usize, const D: usize> Dense<F, N, D> {

mvpoly/tests/monomials.rs

@@ -430,3 +430,42 @@ fn test_homogeneous_eval() {
             + u * u * Fp::from(42)
     );
 }
+
+#[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/tests/prime.rs

@@ -830,3 +830,43 @@ fn test_homogeneous_eval() {
     let homogenous_eval = p1.homogeneous_eval(&random_eval, u);
     assert_eq!(homogenous_eval, u * u);
 }
+
+#[test]
+fn test_add_monomial() {
+    let mut rng = o1_utils::tests::make_test_rng(None);
+
+    // Adding constant monomial one to zero
+    let mut p1 = Dense::<Fp, 4, 2>::zero();
+    p1.add_monomial([0, 0, 0, 0], Fp::one());
+    assert_eq!(p1, Dense::<Fp, 4, 2>::one());
+
+    // Adding random constant monomial one to zero
+    let mut p2 = Dense::<Fp, 4, 2>::zero();
+    let random_c = Fp::rand(&mut rng);
+    p2.add_monomial([0, 0, 0, 0], random_c);
+    assert_eq!(p2, Dense::<Fp, 4, 2>::from(random_c));
+
+    let mut p3 = Dense::<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 = Dense::<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);
+}