Merge pull request #2872 from o1-labs/mvpoly/param-compute-nested-loo…

…p-degree-bound mvpoly: introduce optional param upper bound to compute nested loops
o1-labs · Dec 11, 2024 · 644bf15 · 644bf15
2 parents 155b329 + ac2500f
commit 644bf15
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 15 deletions.
diff --git a/mvpoly/src/monomials.rs b/mvpoly/src/monomials.rs
@@ -293,7 +293,8 @@ impl<const N: usize, const D: usize, F: PrimeField> MVPoly<F, N, D> for Sparse<F
     unsafe fn random<RNG: RngCore>(rng: &mut RNG, max_degree: Option<usize>) -> Self {
         let degree = max_degree.unwrap_or(D);
         // Generating all monomials with degree <= degree^N
-        let nested_loops_indices: Vec<Vec<usize>> = compute_indices_nested_loop(vec![degree; N]);
+        let nested_loops_indices: Vec<Vec<usize>> =
+            compute_indices_nested_loop(vec![degree; N], max_degree);
         // Filtering the monomials with degree <= degree
         let exponents: Vec<Vec<usize>> = nested_loops_indices
             .into_iter()
@@ -409,8 +410,10 @@ impl<const N: usize, const D: usize, F: PrimeField> MVPoly<F, N, D> for Sparse<F
             // 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(non_zero_exponents.iter().map(|d| *d + 1).collect());
+            let indices = compute_indices_nested_loop(
+                non_zero_exponents.iter().map(|d| *d + 1).collect(),
+                None,
+            );
             for i in 0..=u_degree {
                 // Add the binomial from the homogeneisation
                 // i.e (u_degree choose i)

diff --git a/mvpoly/src/utils.rs b/mvpoly/src/utils.rs
@@ -220,7 +220,8 @@ pub fn compute_all_two_factors_decomposition(
     }
 }
 
-/// Compute the list of indices to perform N nested loops of different size each.
+/// Compute the list of indices to perform N nested loops of different size
+/// each, whose sum is less than or equal to an optional upper bound.
 /// In other words, if we have to perform the 3 nested loops:
 /// ```rust
 /// let n1 = 3;
@@ -254,10 +255,13 @@ pub fn compute_all_two_factors_decomposition(
 ///
 /// In the case of an empty loop (i.e. one value in the input list is 0), the
 /// expected output is the empty list.
-pub fn compute_indices_nested_loop(nested_loop_sizes: Vec<usize>) -> Vec<Vec<usize>> {
+pub fn compute_indices_nested_loop(
+    nested_loop_sizes: Vec<usize>,
+    upper_bound: Option<usize>,
+) -> Vec<Vec<usize>> {
     let n = nested_loop_sizes.iter().product();
     (0..n)
-        .map(|i| {
+        .filter_map(|i| {
             let mut div = 1;
             // Compute indices for the loop, step i
             let indices: Vec<usize> = nested_loop_sizes
@@ -268,7 +272,15 @@ pub fn compute_indices_nested_loop(nested_loop_sizes: Vec<usize>) -> Vec<Vec<usi
                     k
                 })
                 .collect();
-            indices
+            if let Some(upper_bound) = upper_bound {
+                if indices.iter().sum::<usize>() <= upper_bound {
+                    Some(indices)
+                } else {
+                    None
+                }
+            } else {
+                Some(indices)
+            }
         })
         .collect()
 }
diff --git a/mvpoly/tests/utils.rs b/mvpoly/tests/utils.rs
@@ -232,7 +232,7 @@ pub fn test_compute_indices_nested_loop() {
     // 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);
+    let mut comp_indices = compute_indices_nested_loop(nested_loops, None);
     comp_indices.sort();
     assert_eq!(exp_indices, comp_indices);
 
@@ -247,7 +247,7 @@ pub fn test_compute_indices_nested_loop() {
         vec![2, 1],
     ];
     exp_indices.sort();
-    let mut comp_indices = compute_indices_nested_loop(nested_loops);
+    let mut comp_indices = compute_indices_nested_loop(nested_loops, None);
     comp_indices.sort();
     assert_eq!(exp_indices, comp_indices);
 
@@ -292,38 +292,51 @@ pub fn test_compute_indices_nested_loop() {
         vec![2, 2, 1, 1],
     ];
     exp_indices.sort();
-    let mut comp_indices = compute_indices_nested_loop(nested_loops);
+    let mut comp_indices = compute_indices_nested_loop(nested_loops, None);
     comp_indices.sort();
     assert_eq!(exp_indices, comp_indices);
 
     // Simple and single loop
     let nested_loops = vec![3];
     let exp_indices = vec![vec![0], vec![1], vec![2]];
-    let mut comp_indices = compute_indices_nested_loop(nested_loops);
+    let mut comp_indices = compute_indices_nested_loop(nested_loops, None);
     comp_indices.sort();
     assert_eq!(exp_indices, comp_indices);
 
     // relatively large loops
     let nested_loops = vec![10, 10];
-    let comp_indices = compute_indices_nested_loop(nested_loops);
+    let comp_indices = compute_indices_nested_loop(nested_loops, None);
     // Only checking the length as it would take too long to unroll the result
     assert_eq!(comp_indices.len(), 100);
 
     // Non-uniform loop sizes, relatively large
     let nested_loops = vec![5, 7, 3];
-    let comp_indices = compute_indices_nested_loop(nested_loops);
+    let comp_indices = compute_indices_nested_loop(nested_loops, None);
     assert_eq!(comp_indices.len(), 5 * 7 * 3);
 }
 
 #[test]
 fn test_compute_indices_nested_loop_edge_cases() {
     let nested_loops = vec![];
-    let comp_indices: Vec<Vec<usize>> = compute_indices_nested_loop(nested_loops);
+    let comp_indices: Vec<Vec<usize>> = compute_indices_nested_loop(nested_loops, None);
     let exp_output: Vec<Vec<usize>> = vec![vec![]];
     assert_eq!(comp_indices, exp_output);
 
     // With one empty loop. Should match the documentation
     let nested_loops = vec![3, 0, 2];
-    let comp_indices = compute_indices_nested_loop(nested_loops);
+    let comp_indices = compute_indices_nested_loop(nested_loops, None);
     assert_eq!(comp_indices.len(), 0);
 }
+
+#[test]
+fn test_compute_indices_nested_loops_upper_bound() {
+    let nested_loops = vec![3, 3];
+    let comp_indices = compute_indices_nested_loop(nested_loops.clone(), Some(0));
+    assert_eq!(comp_indices.len(), 1);
+
+    let comp_indices = compute_indices_nested_loop(nested_loops.clone(), Some(1));
+    assert_eq!(comp_indices.len(), 3);
+
+    let comp_indices = compute_indices_nested_loop(nested_loops, Some(2));
+    assert_eq!(comp_indices.len(), 6);
+}