Fix bench_recursion and a bit of cleanup

0xPolygonZero · Sep 7, 2024 · b79391b · b79391b
1 parent 1650fee
commit b79391b
Show file tree

Hide file tree

Showing 12 changed files with 546 additions and 180 deletions.
diff --git a/plonky2/src/fri/witness_util.rs b/plonky2/src/fri/witness_util.rs
@@ -80,7 +80,6 @@ where
             tmp.push((vec![], dummy_proof));
             tmp
         } else {
-            println!("we're good");
             q.initial_trees_proof.evals_proofs.clone()
         };
         for (at, a) in qt

diff --git a/plonky2/src/iop/witness.rs b/plonky2/src/iop/witness.rs
@@ -2,7 +2,7 @@
 use alloc::{vec, vec::Vec};
 use core::iter::zip;
 
-use anyhow::{anyhow, Result};
+use anyhow::{anyhow, ensure, Result};
 use hashbrown::HashMap;
 use itertools::{zip_eq, Itertools};
 
@@ -45,6 +45,24 @@ pub trait WitnessWrite<F: Field> {
         Ok(())
     }
 
+    fn set_opt_cap_target<H: AlgebraicHasher<F>>(
+        &mut self,
+        opt_ct: &Option<MerkleCapTarget>,
+        opt_value: &Option<MerkleCap<F, H>>,
+    ) -> Result<()>
+    where
+        F: RichField,
+    {
+        ensure!(opt_ct.is_some() == opt_value.is_some());
+        if let (Some(ct), Some(value)) = (opt_ct, opt_value) {
+            for (ht, h) in ct.0.iter().zip(&value.0) {
+                self.set_hash_target(*ht, *h)?;
+            }
+        }
+
+        Ok(())
+    }
+
     fn set_extension_target<const D: usize>(
         &mut self,
         et: ExtensionTarget<D>,
@@ -128,7 +146,7 @@ pub trait WitnessWrite<F: Field> {
             &proof.plonk_zs_partial_products_cap,
         )?;
         self.set_cap_target(&proof_target.quotient_polys_cap, &proof.quotient_polys_cap)?;
-        self.set_cap_target(&proof_target.random_r, &proof.random_r)?;
+        self.set_opt_cap_target(&proof_target.opt_random_r, &proof.opt_random_r)?;
 
         self.set_fri_openings(
             &proof_target.openings.to_fri_openings(),

diff --git a/plonky2/src/plonk/circuit_data.rs b/plonky2/src/plonk/circuit_data.rs
@@ -576,28 +576,49 @@ impl<F: RichField + Extendable<D>, const D: usize> CommonCircuitData<F, D> {
     }
 
     fn fri_oracles(&self) -> Vec<FriOracleInfo> {
-        vec![
-            FriOracleInfo {
-                num_polys: self.num_preprocessed_polys(),
-                blinding: PlonkOracle::CONSTANTS_SIGMAS.blinding,
-            },
-            FriOracleInfo {
-                num_polys: self.config.num_wires,
-                blinding: PlonkOracle::WIRES.blinding,
-            },
-            FriOracleInfo {
-                num_polys: self.num_zs_partial_products_polys() + self.num_all_lookup_polys(),
-                blinding: PlonkOracle::ZS_PARTIAL_PRODUCTS.blinding,
-            },
-            FriOracleInfo {
-                num_polys: self.num_quotient_polys(),
-                blinding: PlonkOracle::QUOTIENT.blinding,
-            },
-            FriOracleInfo {
-                num_polys: self.num_r_polys(),
-                blinding: PlonkOracle::R.blinding,
-            },
-        ]
+        if self.num_r_polys() > 0 {
+            vec![
+                FriOracleInfo {
+                    num_polys: self.num_preprocessed_polys(),
+                    blinding: PlonkOracle::CONSTANTS_SIGMAS.blinding,
+                },
+                FriOracleInfo {
+                    num_polys: self.config.num_wires,
+                    blinding: PlonkOracle::WIRES.blinding,
+                },
+                FriOracleInfo {
+                    num_polys: self.num_zs_partial_products_polys() + self.num_all_lookup_polys(),
+                    blinding: PlonkOracle::ZS_PARTIAL_PRODUCTS.blinding,
+                },
+                FriOracleInfo {
+                    num_polys: self.num_quotient_polys(),
+                    blinding: PlonkOracle::QUOTIENT.blinding,
+                },
+                FriOracleInfo {
+                    num_polys: self.num_r_polys(),
+                    blinding: PlonkOracle::R.blinding,
+                },
+            ]
+        } else {
+            vec![
+                FriOracleInfo {
+                    num_polys: self.num_preprocessed_polys(),
+                    blinding: PlonkOracle::CONSTANTS_SIGMAS.blinding,
+                },
+                FriOracleInfo {
+                    num_polys: self.config.num_wires,
+                    blinding: PlonkOracle::WIRES.blinding,
+                },
+                FriOracleInfo {
+                    num_polys: self.num_zs_partial_products_polys() + self.num_all_lookup_polys(),
+                    blinding: PlonkOracle::ZS_PARTIAL_PRODUCTS.blinding,
+                },
+                FriOracleInfo {
+                    num_polys: self.num_quotient_polys(),
+                    blinding: PlonkOracle::QUOTIENT.blinding,
+                },
+            ]
+        }
     }
 
     fn fri_preprocessed_polys(&self) -> Vec<FriPolynomialInfo> {
@@ -667,15 +688,27 @@ impl<F: RichField + Extendable<D>, const D: usize> CommonCircuitData<F, D> {
     }
 
     fn fri_all_polys(&self) -> Vec<FriPolynomialInfo> {
-        [
-            self.fri_preprocessed_polys(),
-            self.fri_wire_polys(),
-            self.fri_zs_partial_products_polys(),
-            self.fri_quotient_polys(),
-            self.fri_lookup_polys(),
-            self.fri_r_polys(),
-        ]
-        .concat()
+        let num_r_polys = self.num_r_polys();
+        if num_r_polys > 0 {
+            [
+                self.fri_preprocessed_polys(),
+                self.fri_wire_polys(),
+                self.fri_zs_partial_products_polys(),
+                self.fri_quotient_polys(),
+                self.fri_lookup_polys(),
+                self.fri_r_polys(),
+            ]
+            .concat()
+        } else {
+            [
+                self.fri_preprocessed_polys(),
+                self.fri_wire_polys(),
+                self.fri_zs_partial_products_polys(),
+                self.fri_quotient_polys(),
+                self.fri_lookup_polys(),
+            ]
+            .concat()
+        }
     }
 }
 

diff --git a/plonky2/src/plonk/get_challenges.rs b/plonky2/src/plonk/get_challenges.rs
@@ -114,7 +114,7 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
             wires_cap,
             plonk_zs_partial_products_cap,
             quotient_polys_cap,
-            random_r,
+            opt_random_r,
             openings,
             opening_proof:
                 FriProof {
@@ -154,7 +154,7 @@ impl<F: RichField + Extendable<D>, C: GenericConfig<D, F = F>, const D: usize>
             wires_cap,
             plonk_zs_partial_products_cap,
             quotient_polys_cap,
-            random_r,
+            opt_random_r,
             openings,
             opening_proof:
                 CompressedFriProof {
@@ -344,7 +344,7 @@ impl<const D: usize> ProofWithPublicInputsTarget<D> {
             wires_cap,
             plonk_zs_partial_products_cap,
             quotient_polys_cap,
-            random_r,
+            opt_random_r,
             openings,
             opening_proof:
                 FriProofTarget {