From 86d9aa6b3caf9ed7189faaa7b98240c97051a7f3 Mon Sep 17 00:00:00 2001
From: Danny Willems <be.danny.willems@gmail.com>
Date: Thu, 5 Sep 2024 11:51:47 -0700
Subject: [PATCH] MVPoly: add is_multilinear

---
 mvpoly/src/lib.rs         |  5 +++++
 mvpoly/src/monomials.rs   |  6 ++++++
 mvpoly/src/prime.rs       | 20 ++++++++++++++++++++
 mvpoly/tests/monomials.rs | 37 ++++++++++++++++++++++++++++++++++++-
 mvpoly/tests/prime.rs     | 39 ++++++++++++++++++++++++++++++++++++++-
 5 files changed, 105 insertions(+), 2 deletions(-)

diff --git a/mvpoly/src/lib.rs b/mvpoly/src/lib.rs
index 87a983cfcb..5f01259531 100644
--- a/mvpoly/src/lib.rs
+++ b/mvpoly/src/lib.rs
@@ -98,4 +98,9 @@ pub trait MVPoly<F: PrimeField, const N: usize, const D: usize>:
         u1: F,
         u2: F,
     ) -> HashMap<usize, F>;
+
+
+    /// Return true if the multi-variate polynomial is multilinear, i.e. if each
+    /// variable in each monomial is of maximum degree 1.
+    fn is_multilinear(&self) -> bool;
 }
diff --git a/mvpoly/src/monomials.rs b/mvpoly/src/monomials.rs
index c6ab2359b0..c96468958b 100644
--- a/mvpoly/src/monomials.rs
+++ b/mvpoly/src/monomials.rs
@@ -557,6 +557,12 @@ impl<const N: usize, const D: usize, F: PrimeField> MVPoly<F, N, D> for Sparse<F
         });
         cross_terms_by_powers_of_r
     }
+
+    fn is_multilinear(&self) -> bool {
+        self.monomials
+            .iter()
+            .all(|(exponents, _)| exponents.iter().all(|&d| d <= 1))
+    }
 }
 
 impl<const N: usize, const D: usize, F: PrimeField> From<prime::Dense<F, N, D>>
diff --git a/mvpoly/src/prime.rs b/mvpoly/src/prime.rs
index 51b0852380..624dada66b 100644
--- a/mvpoly/src/prime.rs
+++ b/mvpoly/src/prime.rs
@@ -388,6 +388,26 @@ impl<F: PrimeField, const N: usize, const D: usize> MVPoly<F, N, D> for Dense<F,
     ) -> HashMap<usize, F> {
         unimplemented!()
     }
+
+    fn is_multilinear(&self) -> bool {
+        if self.is_zero() {
+            return true;
+        }
+        let normalized_indices = self.normalized_indices.clone();
+        let mut prime_gen = PrimeNumberGenerator::new();
+        normalized_indices
+            .iter()
+            .zip(self.coeff.iter())
+            .all(|(idx, c)| {
+                if c.is_zero() {
+                    true
+                } else {
+                    let decomposition_of_i = naive_prime_factors(*idx, &mut prime_gen);
+                    // Each prime number/variable should appear at most once
+                    decomposition_of_i.iter().all(|(_p, d)| *d <= 1)
+                }
+            })
+    }
 }
 
 impl<F: PrimeField, const N: usize, const D: usize> Dense<F, N, D> {
diff --git a/mvpoly/tests/monomials.rs b/mvpoly/tests/monomials.rs
index 0a630a2fce..c40a9339f8 100644
--- a/mvpoly/tests/monomials.rs
+++ b/mvpoly/tests/monomials.rs
@@ -1,7 +1,7 @@
 use ark_ff::{Field, One, UniformRand, Zero};
 use mina_curves::pasta::Fp;
 use mvpoly::{monomials::Sparse, MVPoly};
-use rand::Rng;
+use rand::{seq::SliceRandom, Rng};
 
 #[test]
 fn test_mul_by_one() {
@@ -663,3 +663,38 @@ fn test_mvpoly_compute_cross_terms_degree_seven() {
     };
     assert_eq!(lhs, rhs);
 }
+
+#[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());
+}
diff --git a/mvpoly/tests/prime.rs b/mvpoly/tests/prime.rs
index 98da818d1a..d9c1864825 100644
--- a/mvpoly/tests/prime.rs
+++ b/mvpoly/tests/prime.rs
@@ -5,7 +5,7 @@ use kimchi::circuits::{
 };
 use mina_curves::pasta::Fp;
 use mvpoly::{prime::Dense, utils::PrimeNumberGenerator, MVPoly};
-use rand::Rng;
+use rand::{seq::SliceRandom, Rng};
 
 #[test]
 fn test_vector_space_dimension() {
@@ -919,3 +919,40 @@ fn test_add_monomial() {
         random_c1 * random_eval[0] * random_eval[0] + random_c2 * random_eval[1] * random_eval[1];
     assert_eq!(eval_p4, exp_eval_p4);
 }
+
+#[test]
+fn test_is_multilinear() {
+    let mut rng = o1_utils::tests::make_test_rng(None);
+    let p1 = Dense::<Fp, 6, 2>::zero();
+    assert!(p1.is_multilinear());
+
+    let c = Fp::rand(&mut rng);
+    let p2 = Dense::<Fp, 6, 2>::from(c);
+    assert!(p2.is_multilinear());
+
+    {
+        let mut p = Dense::<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 = Dense::<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 p = unsafe { Dense::<Fp, 6, 4>::random(&mut rng, None) };
+        assert!(!p.is_multilinear());
+    }
+}