MVPoly/prime: implement compute_cross_terms

mvpoly/src/prime.rs

@@ -161,7 +161,8 @@ use rand::RngCore;
 use std::ops::{Index, IndexMut};
 
 use crate::utils::{
-    compute_all_two_factors_decomposition, naive_prime_factors, PrimeNumberGenerator,
+    compute_all_two_factors_decomposition, compute_indices_nested_loop, naive_prime_factors,
+    PrimeNumberGenerator,
 };
 
 /// Represents a multivariate polynomial of degree less than `D` in `N` variables.
@@ -439,6 +440,126 @@ impl<F: PrimeField, const N: usize, const D: usize> Dense<F, N, D> {
             });
         result
     }
+
+    /// Compute the cross-terms as described in [Behind Nova: cross-terms
+    /// computation for high degree
+    /// gates](https://hackmd.io/@dannywillems/Syo5MBq90)
+    ///
+    /// The polynomial must not necessarily be homogeneous. For this reason, the
+    /// values `u1` and `u2` represents the extra variable that is used to make
+    /// the polynomial homogeneous.
+    ///
+    /// The homogeneous degree is supposed to be the one defined by the type of
+    /// the polynomial, i.e. `D`.
+    ///
+    /// The output is a map of `D - 1` values that represents the cross-terms
+    /// for each power of `r`.
+    // IMPROVEME: Dummy implementation, a cache can be used to save the
+    // previously computed multiplications and powers.
+    // Maybe using a symbolic form could speed up the computation.
+    pub fn compute_cross_terms(
+        &self,
+        eval1: &[F; N],
+        eval2: &[F; N],
+        u1: F,
+        u2: F,
+    ) -> HashMap<usize, F> {
+        assert!(
+            D >= 2,
+            "The degree of the polynomial must be greater than 2"
+        );
+        let mut cross_terms_by_powers_of_r: HashMap<usize, F> = HashMap::new();
+        let mut prime_gen = PrimeNumberGenerator::new();
+        let primes = prime_gen.get_first_nth_primes(N);
+        // Computing invididual contribution of each monomial
+        // FIXME: handle constant, prime decompo returns empty list of 1.
+        self.coeff.iter().enumerate().for_each(|(i, c)| {
+            // If the coefficient is zero, we skip the computation as the contribution
+            // is null.
+            if *c == F::zero() || i == 0 {
+                // FIXME: handle constant, prime decompo returns empty list of 1.
+            } else {
+                let idx = self.normalized_indices[i];
+                let prime_decomposition = naive_prime_factors(idx, &mut prime_gen);
+                // Fetch individual degree
+                let degrees: Vec<usize> = prime_decomposition.iter().map(|(_, d)| *d).collect();
+                // No cross-terms
+                // FIXME
+                if degrees.len() == 1 && (degrees[0] == 0 || degrees[0] == D) {
+                    return;
+                }
+                // We compute the missing degree to homogeneize the polynomial.
+                // The variable u given as a parameter will be used to
+                // homogeneize the polynomial.
+                let u_degree: usize = D - degrees.iter().sum::<usize>();
+                // Will be used to compute the nested sums
+                // It returns all the indices i_1, ..., i_k for the sums:
+                // Σ_{i_1 = 0}^{n_1} Σ_{i_2 = 0}^{n_2} ... Σ_{i_k = 0}^{n_k}
+                let indices = compute_indices_nested_loop(degrees.iter().map(|d| *d + 1).collect());
+                // We treat the homogeneisation degree separately in the sum. It
+                // eases the code below
+                for i in 0..=u_degree {
+                    // Add the binomial from the homogeneisation
+                    // i.e (u_degree choose i)
+                    let u_binomial_term = binomial(u_degree, i);
+                    indices.iter().for_each(|indices| {
+                        let sum_indices = indices.iter().sum::<usize>() + i;
+                        // If the sum of the indices is 0 or D, we skip the
+                        // computation as the contribution would go in the
+                        // evaluation of the polynomial at each evaluation
+                        // vectors eval1 and eval2
+                        if sum_indices == 0 || sum_indices == D {
+                            return;
+                        }
+                        // Compute
+                        // (n_1 choose i_1) * (n_2 choose i_2) * ... * (n_k choose i_k)
+                        let binomial_term = indices
+                            .iter()
+                            .zip(degrees.iter())
+                            .fold(u_binomial_term, |acc, (i, d)| acc * binomial(*d, *i));
+                        let binomial_term = F::from(binomial_term as u64);
+                        // Compute the product x_k^i_k
+                        let monomial_eval1 = prime_decomposition.iter().zip(indices.iter()).fold(
+                            F::one(),
+                            |acc, ((p, _), i)| {
+                                // Get the evaluation of the corresponding
+                                // variable
+                                let inv_p = primes
+                                    .iter()
+                                    .position(|x| x == p)
+                                    .expect("The variable must be in the list of primes");
+                                acc * eval1[inv_p].pow([*i as u64])
+                            },
+                        );
+                        // Compute the product x'_k^(n_k - i_k)
+                        let monomial_eval2 = prime_decomposition.iter().zip(indices.iter()).fold(
+                            F::one(),
+                            |acc, ((p, d), i)| {
+                                // Get the evaluation of the corresponding
+                                // variable
+                                let inv_p = primes
+                                    .iter()
+                                    .position(|x| x == p)
+                                    .expect("The variable must be in the list of primes");
+                                acc * eval2[inv_p].pow([(d - *i) as u64])
+                            },
+                        );
+                        // u1^i * u2^(u_degree - i)
+                        let u = u1.pow([i as u64]) * u2.pow([(u_degree - i) as u64]);
+                        let res = binomial_term * monomial_eval1 * monomial_eval2 * u;
+                        let res = *c * res;
+                        // power of r is Σ (n_k - i_k)
+                        let power_r: usize = D - sum_indices;
+                        cross_terms_by_powers_of_r
+                            .entry(power_r)
+                            .and_modify(|e| *e += res)
+                            .or_insert(res);
+                    });
+                }
+            }
+        });
+        cross_terms_by_powers_of_r
+    }
 }
 
 impl<F: PrimeField, const N: usize, const D: usize> Default for Dense<F, N, D> {

mvpoly/src/utils.rs

@@ -219,3 +219,54 @@ pub fn compute_all_two_factors_decomposition(
         factors
     }
 }
+
+/// Compute the list of indices to perform N nested loops of different size each.
+/// In other words, if we have to perform the 3 nested loops:
+/// ```rust
+/// let n1 = 3;
+/// let n2 = 3;
+/// let n3 = 5;
+/// for i in 0..n1 {
+///   for j in 0..n2 {
+///     for k in 0..n3 {
+///     }
+///   }
+/// }
+/// ```
+/// the output will be all the possible values of `i`, `j`, and `k`.
+/// The algorithm is as follows:
+/// ```rust
+/// let n1 = 3;
+/// let n2 = 3;
+/// let n3 = 5;
+/// (0..(n1 * n2 * n3)).map(|l| {
+///   let i = l               % n1;
+///   let j = (l / n1)        % n2;
+///   let k = (l / (n1 * n2)) % n3;
+///   (i, j, k)
+/// });
+/// ```
+/// For N nested loops, the algorithm is the same, with the division increasing
+/// by the factor `N_k` for the index `i_(k + 1)`
+pub fn compute_indices_nested_loop(nested_loop_sizes: Vec<usize>) -> Vec<Vec<usize>> {
+    let n = nested_loop_sizes.iter().product();
+    (0..n)
+        .map(|i| {
+            let mut div = 1;
+            // Compute indices for the loop, step i
+            let indices: Vec<usize> = nested_loop_sizes
+                .iter()
+                .map(|n_i| {
+                    let k = (i / div) % n_i;
+                    div *= n_i;
+                    k
+                })
+                .collect();
+            assert!(
+                div == n,
+                "The division must be equal to the number of terms at the end"
+            );
+            indices
+        })
+        .collect()
+}

mvpoly/tests/prime.rs

@@ -802,3 +802,30 @@ fn test_mvpoly_mul_pbt() {
     let p2 = unsafe { Dense::<Fp, 4, 6>::random(&mut rng, Some(max_degree)) };
     assert_eq!(p1.clone() * p2.clone(), p2.clone() * p1.clone());
 }
+
+#[test]
+fn test_mvpoly_compute_cross_terms_eval_zero() {
+    let mut rng = o1_utils::tests::make_test_rng(None);
+    let p1 = unsafe { Dense::<Fp, 4, 2>::random(&mut rng, None) };
+    let u1 = Fp::rand(&mut rng);
+    let u2 = Fp::rand(&mut rng);
+    let eval: [Fp; 4] = [Fp::zero(); 4];
+    let res = p1.compute_cross_terms(&eval, &eval, u1, u2);
+    res.iter()
+        .for_each(|(_power, cross_term)| assert_eq!(*cross_term, Fp::zero()));
+}
+
+// FIXME
+#[test]
+fn test_mvpoly_compute_cross_terms_degree_two() {
+    let mut rng = o1_utils::tests::make_test_rng(None);
+    let p1 = unsafe { Dense::<Fp, 4, 2>::random(&mut rng, None) };
+    let random_eval1: [Fp; 4] = std::array::from_fn(|_| Fp::rand(&mut rng));
+    let random_eval2: [Fp; 4] = std::array::from_fn(|_| Fp::rand(&mut rng));
+    let u1 = Fp::rand(&mut rng);
+    let u2 = Fp::rand(&mut rng);
+    let res = p1.compute_cross_terms(&random_eval1, &random_eval2, u1, u2);
+    // We only have one cross-term in this case
+    assert_eq!(res.len(), 1);
+    println!("{:?}", res);
+}

mvpoly/tests/utils.rs

@@ -1,8 +1,8 @@
 use std::collections::HashMap;
 
 use mvpoly::utils::{
-    compute_all_two_factors_decomposition, get_mapping_with_primes, is_prime, naive_prime_factors,
-    PrimeNumberGenerator,
+    compute_all_two_factors_decomposition, compute_indices_nested_loop, get_mapping_with_primes,
+    is_prime, naive_prime_factors, PrimeNumberGenerator,
 };
 
 pub const FIRST_FIFTY_PRIMES: [usize; 50] = [
@@ -225,3 +225,74 @@ pub fn test_compute_all_two_factors_decomposition_with_multiplicity() {
         .to_vec()
     );
 }
+
+#[test]
+pub fn test_compute_indices_nested_loop() {
+    let nested_loops = vec![2, 2];
+    // sorting to get the same order
+    let mut exp_indices = vec![vec![0, 0], vec![0, 1], vec![1, 0], vec![1, 1]];
+    exp_indices.sort();
+    let mut comp_indices = compute_indices_nested_loop(nested_loops);
+    comp_indices.sort();
+    assert_eq!(exp_indices, comp_indices);
+
+    let nested_loops = vec![3, 2];
+    // sorting to get the same order
+    let mut exp_indices = vec![
+        vec![0, 0],
+        vec![0, 1],
+        vec![1, 0],
+        vec![1, 1],
+        vec![2, 0],
+        vec![2, 1],
+    ];
+    exp_indices.sort();
+    let mut comp_indices = compute_indices_nested_loop(nested_loops);
+    comp_indices.sort();
+    assert_eq!(exp_indices, comp_indices);
+
+    let nested_loops = vec![3, 3, 2, 2];
+    // sorting to get the same order
+    let mut exp_indices = vec![
+        vec![0, 0, 0, 0],
+        vec![0, 0, 0, 1],
+        vec![0, 0, 1, 0],
+        vec![0, 0, 1, 1],
+        vec![0, 1, 0, 0],
+        vec![0, 1, 0, 1],
+        vec![0, 1, 1, 0],
+        vec![0, 1, 1, 1],
+        vec![0, 2, 0, 0],
+        vec![0, 2, 0, 1],
+        vec![0, 2, 1, 0],
+        vec![0, 2, 1, 1],
+        vec![1, 0, 0, 0],
+        vec![1, 0, 0, 1],
+        vec![1, 0, 1, 0],
+        vec![1, 0, 1, 1],
+        vec![1, 1, 0, 0],
+        vec![1, 1, 0, 1],
+        vec![1, 1, 1, 0],
+        vec![1, 1, 1, 1],
+        vec![1, 2, 0, 0],
+        vec![1, 2, 0, 1],
+        vec![1, 2, 1, 0],
+        vec![1, 2, 1, 1],
+        vec![2, 0, 0, 0],
+        vec![2, 0, 0, 1],
+        vec![2, 0, 1, 0],
+        vec![2, 0, 1, 1],
+        vec![2, 1, 0, 0],
+        vec![2, 1, 0, 1],
+        vec![2, 1, 1, 0],
+        vec![2, 1, 1, 1],
+        vec![2, 2, 0, 0],
+        vec![2, 2, 0, 1],
+        vec![2, 2, 1, 0],
+        vec![2, 2, 1, 1],
+    ];
+    exp_indices.sort();
+    let mut comp_indices = compute_indices_nested_loop(nested_loops);
+    comp_indices.sort();
+    assert_eq!(exp_indices, comp_indices);
+}