Merge pull request #2522 from o1-labs/dw/edge-case-utils-indices

MVPoly/utils: edge cases for nested loop indices computation

mvpoly/src/utils.rs

@@ -248,6 +248,12 @@ pub fn compute_all_two_factors_decomposition(
 /// ```
 /// For N nested loops, the algorithm is the same, with the division increasing
 /// by the factor `N_k` for the index `i_(k + 1)`
+///
+/// In the case of an empty list, the function will return a list containing a
+/// single element which is the empty list.
+///
+/// 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>> {
     let n = nested_loop_sizes.iter().product();
     (0..n)

mvpoly/tests/utils.rs

@@ -295,4 +295,35 @@ pub fn test_compute_indices_nested_loop() {
     let mut comp_indices = compute_indices_nested_loop(nested_loops);
     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);
+    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);
+    // 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);
+    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 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);
+    assert_eq!(comp_indices.len(), 0);
 }