Merge pull request #2514 from o1-labs/dw/compute-indices-nested-loops

MVPoly: util to compute nested loops indices

.github/workflows/ci-nightly.yml

@@ -57,6 +57,8 @@ jobs:
           eval $(opam env)
           make clean
           make nextest-all-with-coverage
+          make test-doc-with-coverage
+          make generate-test-coverage-report
 
       - name: Use shared code coverage summary
         uses: ./.github/actions/coverage-summary-shared

.github/workflows/ci.yml

@@ -148,6 +148,12 @@ jobs:
         run: |
           make install-test-deps
 
+      - name: Doc tests
+        if: ${{ matrix.rust_toolchain_version != env.RUST_TOOLCHAIN_COVERAGE_VERSION }}
+        run: |
+          eval $(opam env)
+          make test-doc
+
       - name: Run non-heavy tests without the code coverage
         if: ${{ matrix.rust_toolchain_version != env.RUST_TOOLCHAIN_COVERAGE_VERSION }}
         run: |
@@ -159,6 +165,8 @@ jobs:
         run: |
           eval $(opam env)
           make nextest-with-coverage
+          make test-doc-with-coverage
+          make generate-test-coverage-report
 
       - name: Use shared code coverage summary
         if: ${{ matrix.rust_toolchain_version == env.RUST_TOOLCHAIN_COVERAGE_VERSION }}
@@ -169,8 +177,3 @@ jobs:
         uses: ./.github/actions/codecov-shared
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
-
-      - name: Doc tests
-        run: |
-          eval $(opam env)
-          cargo test --all-features --release --doc

Makefile

@@ -1,5 +1,10 @@
 # Variables
-COVERAGE_ENV = CARGO_INCREMENTAL=0 RUSTFLAGS='-Cinstrument-coverage' LLVM_PROFILE_FILE=$(shell pwd)/target/profraw/cargo-test-%p-%m.profraw
+# Known coverage limitations and issues:
+# - https://github.com/rust-lang/rust/issues/79417
+# - https://github.com/nextest-rs/nextest/issues/16
+# FIXME: Update or remove the `codecov.yml` file to enable the `patch` coverage report and the corresponding PR check,
+#        once situation with the Rust's Doctests will be improved.
+COVERAGE_ENV = CARGO_INCREMENTAL=0 RUSTFLAGS='-Cinstrument-coverage' RUSTDOCFLAGS="-Cinstrument-coverage" LLVM_PROFILE_FILE=$(shell pwd)/target/profraw/cargo-test-%p-%m.profraw
 # FIXME: In latest 0.8.19+ -t CLI argument can accept comma separated list of custom output types, hence, no need in double invocation
 GRCOV_CALL = grcov ./target/profraw --binary-path ./target/release/deps/ -s . --branch --ignore-not-existing --ignore "**/tests/**"
 
@@ -43,53 +48,54 @@ build:
 release:
 		cargo build --release --all-targets --all-features
 
+# Test the project's docs comments
+test-doc:
+		cargo test --all-features --release --doc
+
+test-doc-with-coverage:
+		$(COVERAGE_ENV) $(MAKE) test-doc
+
 # Test the project with non-heavy tests and using native cargo test runner
 test:
 		cargo test --all-features --release $(CARGO_EXTRA_ARGS) -- --nocapture --skip heavy $(BIN_EXTRA_ARGS)
 
 test-with-coverage:
 		$(COVERAGE_ENV) CARGO_EXTRA_ARGS="$(CARGO_EXTRA_ARGS)" BIN_EXTRA_ARGS="$(BIN_EXTRA_ARGS)" $(MAKE) test
-		$(MAKE) generate-test-coverage-report
 
 # Test the project with heavy tests and using native cargo test runner
 test-heavy:
 		cargo test --all-features --release $(CARGO_EXTRA_ARGS) -- --nocapture heavy $(BIN_EXTRA_ARGS)
 
 test-heavy-with-coverage:
 		$(COVERAGE_ENV) CARGO_EXTRA_ARGS="$(CARGO_EXTRA_ARGS)" BIN_EXTRA_ARGS="$(BIN_EXTRA_ARGS)" $(MAKE) test-heavy
-		$(MAKE) generate-test-coverage-report
 
 # Test the project with all tests and using native cargo test runner
 test-all:
 		cargo test --all-features --release $(CARGO_EXTRA_ARGS) -- --nocapture $(BIN_EXTRA_ARGS)
 
 test-all-with-coverage:
 		$(COVERAGE_ENV) CARGO_EXTRA_ARGS="$(CARGO_EXTRA_ARGS)" BIN_EXTRA_ARGS="$(BIN_EXTRA_ARGS)" $(MAKE) test-all
-		$(MAKE) generate-test-coverage-report
 
 # Test the project with non-heavy tests and using nextest test runner
 nextest:
 		cargo nextest run --all-features --release --profile ci -E "not test(heavy)" $(BIN_EXTRA_ARGS)
 
 nextest-with-coverage:
 		$(COVERAGE_ENV) BIN_EXTRA_ARGS="$(BIN_EXTRA_ARGS)" $(MAKE) nextest
-		$(MAKE) generate-test-coverage-report
 
 # Test the project with heavy tests and using nextest test runner
 nextest-heavy:
 		cargo nextest run --all-features --release --profile ci -E "test(heavy)" $(BIN_EXTRA_ARGS)
 
 nextest-heavy-with-coverage:
 		$(COVERAGE_ENV) BIN_EXTRA_ARGS="$(BIN_EXTRA_ARGS)" $(MAKE) nextest-heavy
-		$(MAKE) generate-test-coverage-report
 
 # Test the project with all tests and using nextest test runner
 nextest-all:
 		cargo nextest run --all-features --release --profile ci $(BIN_EXTRA_ARGS)
 
 nextest-all-with-coverage:
 		$(COVERAGE_ENV) BIN_EXTRA_ARGS="$(BIN_EXTRA_ARGS)" $(MAKE) nextest-all
-		$(MAKE) generate-test-coverage-report
 
 # Format the code
 format:
@@ -113,4 +119,4 @@ generate-test-coverage-report:
 		@echo "The test coverage report is available at: ./target/coverage"
 		@echo ""
 
-.PHONY: all setup install-test-deps clean build release test test-with-coverage test-heavy test-heavy-with-coverage test-all test-all-with-coverage nextest nextest-with-coverage nextest-heavy nextest-heavy-with-coverage nextest-all nextest-all-with-coverage format lint generate-test-coverage-report
+.PHONY: all setup install-test-deps clean build release test-doc test-doc-with-coverage test test-with-coverage test-heavy test-heavy-with-coverage test-all test-all-with-coverage nextest nextest-with-coverage nextest-heavy nextest-heavy-with-coverage nextest-all nextest-all-with-coverage format lint generate-test-coverage-report

codecov.yml

@@ -0,0 +1,7 @@
+# FIXME: Remove this file to enable the `patch` coverage report and the corresponding PR check,
+#        once situation with the Rust's Doctests will be improved.
+#        - https://github.com/rust-lang/rust/issues/79417
+#        - https://github.com/nextest-rs/nextest/issues/16
+coverage:
+  status:
+    patch: off

mvpoly/src/utils.rs

@@ -219,3 +219,50 @@ 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();
+            indices
+        })
+        .collect()
+}

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);
+}