From 02fe2b1d59b8bdc8aa6884fafd122079bff3676f Mon Sep 17 00:00:00 2001 From: Andrew Milson Date: Mon, 23 Sep 2024 18:33:52 -1000 Subject: [PATCH] Make CommitmentSchemeProver::prove_values take ownership --- crates/prover/src/core/fri.rs | 2 +- crates/prover/src/core/pcs/prover.rs | 8 +- crates/prover/src/core/prover/mod.rs | 103 +++++++++--------- crates/prover/src/core/vcs/prover.rs | 11 +- crates/prover/src/core/vcs/test_utils.rs | 2 +- crates/prover/src/examples/blake/air.rs | 2 +- crates/prover/src/examples/plonk/mod.rs | 4 +- crates/prover/src/examples/poseidon/mod.rs | 4 +- .../prover/src/examples/state_machine/mod.rs | 4 +- .../prover/src/examples/wide_fibonacci/mod.rs | 12 +- .../src/examples/xor/gkr_lookups/mle_eval.rs | 8 +- 11 files changed, 74 insertions(+), 86 deletions(-) diff --git a/crates/prover/src/core/fri.rs b/crates/prover/src/core/fri.rs index ce5d4bddb..528d457dc 100644 --- a/crates/prover/src/core/fri.rs +++ b/crates/prover/src/core/fri.rs @@ -821,7 +821,7 @@ impl, H: MerkleHasher> FriLayerProver { let commitment = self.merkle_tree.root(); // TODO(andrew): Use _evals. let (_evals, decommitment) = self.merkle_tree.decommit( - [(self.evaluation.len().ilog2(), decommit_positions)] + &[(self.evaluation.len().ilog2(), decommit_positions)] .into_iter() .collect(), self.evaluation.values.columns.iter().collect_vec(), diff --git a/crates/prover/src/core/pcs/prover.rs b/crates/prover/src/core/pcs/prover.rs index 6ea42742a..ed2f67376 100644 --- a/crates/prover/src/core/pcs/prover.rs +++ b/crates/prover/src/core/pcs/prover.rs @@ -81,7 +81,7 @@ impl<'a, B: BackendForChannel, MC: MerkleChannel> CommitmentSchemeProver<'a, } pub fn prove_values( - &self, + self, sampled_points: TreeVec>>>, channel: &mut MC::C, ) -> CommitmentSchemeProof { @@ -133,13 +133,14 @@ impl<'a, B: BackendForChannel, MC: MerkleChannel> CommitmentSchemeProver<'a, .iter() .map(|(&log_size, domain)| (log_size, domain.flatten())) .collect(); - tree.decommit(queries) + tree.decommit(&queries) }); let queried_values = decommitment_results.as_ref().map(|(v, _)| v.clone()); let decommitments = decommitment_results.map(|(_, d)| d); CommitmentSchemeProof { + commitments: self.roots(), sampled_values, decommitments, queried_values, @@ -151,6 +152,7 @@ impl<'a, B: BackendForChannel, MC: MerkleChannel> CommitmentSchemeProver<'a, #[derive(Debug, Serialize, Deserialize)] pub struct CommitmentSchemeProof { + pub commitments: TreeVec, pub sampled_values: TreeVec>>, pub decommitments: TreeVec>, pub queried_values: TreeVec>>, @@ -231,7 +233,7 @@ impl, MC: MerkleChannel> CommitmentTreeProver { /// positions on each column of that size. fn decommit( &self, - queries: BTreeMap>, + queries: &BTreeMap>, ) -> (ColumnVec>, MerkleDecommitment) { let eval_vec = self .evaluations diff --git a/crates/prover/src/core/prover/mod.rs b/crates/prover/src/core/prover/mod.rs index a539f1a9b..30e7cbd54 100644 --- a/crates/prover/src/core/prover/mod.rs +++ b/crates/prover/src/core/prover/mod.rs @@ -1,3 +1,4 @@ +use std::ops::Deref; use std::{array, mem}; use serde::{Deserialize, Serialize}; @@ -9,7 +10,7 @@ use super::backend::BackendForChannel; use super::channel::MerkleChannel; use super::fields::secure_column::SECURE_EXTENSION_DEGREE; use super::fri::FriVerificationError; -use super::pcs::{CommitmentSchemeProof, TreeVec}; +use super::pcs::CommitmentSchemeProof; use super::vcs::ops::MerkleHasher; use crate::constraint_framework::PREPROCESSED_TRACE_IDX; use crate::core::channel::Channel; @@ -22,17 +23,11 @@ use crate::core::vcs::hash::Hash; use crate::core::vcs::prover::MerkleDecommitment; use crate::core::vcs::verifier::MerkleVerificationError; -#[derive(Debug, Serialize, Deserialize)] -pub struct StarkProof { - pub commitments: TreeVec, - pub commitment_scheme_proof: CommitmentSchemeProof, -} - #[instrument(skip_all)] pub fn prove, MC: MerkleChannel>( components: &[&dyn ComponentProver], channel: &mut MC::C, - commitment_scheme: &mut CommitmentSchemeProver<'_, B, MC>, + mut commitment_scheme: CommitmentSchemeProver<'_, B, MC>, ) -> Result, ProvingError> { let n_preprocessed_columns = commitment_scheme.trees[PREPROCESSED_TRACE_IDX] .polynomials @@ -67,25 +62,19 @@ pub fn prove, MC: MerkleChannel>( // Prove the trace and composition OODS values, and retrieve them. let commitment_scheme_proof = commitment_scheme.prove_values(sample_points, channel); - - let sampled_oods_values = &commitment_scheme_proof.sampled_values; - let composition_oods_eval = extract_composition_eval(sampled_oods_values).unwrap(); + let proof = StarkProof(commitment_scheme_proof); + info!(proof_size_estimate = proof.size_estimate()); // Evaluate composition polynomial at OODS point and check that it matches the trace OODS // values. This is a sanity check. - if composition_oods_eval + if proof.extract_composition_oods_eval().unwrap() != component_provers .components() - .eval_composition_polynomial_at_point(oods_point, sampled_oods_values, random_coeff) + .eval_composition_polynomial_at_point(oods_point, &proof.sampled_values, random_coeff) { return Err(ProvingError::ConstraintsNotSatisfied); } - let proof = StarkProof { - commitments: commitment_scheme.roots(), - commitment_scheme_proof, - }; - info!(proof_size_estimate = proof.size_estimate()); Ok(proof) } @@ -120,42 +109,21 @@ pub fn verify( // Add the composition polynomial mask points. sample_points.push(vec![vec![oods_point]; SECURE_EXTENSION_DEGREE]); - let sampled_oods_values = &proof.commitment_scheme_proof.sampled_values; - let composition_oods_eval = extract_composition_eval(sampled_oods_values).map_err(|_| { + let composition_oods_eval = proof.extract_composition_oods_eval().map_err(|_| { VerificationError::InvalidStructure("Unexpected sampled_values structure".to_string()) })?; if composition_oods_eval != components.eval_composition_polynomial_at_point( oods_point, - sampled_oods_values, + &proof.sampled_values, random_coeff, ) { return Err(VerificationError::OodsNotMatching); } - commitment_scheme.verify_values(sample_points, proof.commitment_scheme_proof, channel) -} - -/// Extracts the composition trace evaluation from the mask. -fn extract_composition_eval( - mask: &TreeVec>>, -) -> Result { - let mut composition_cols = mask.last().into_iter().flatten(); - - let coordinate_evals = array::try_from_fn(|_| { - let col = &**composition_cols.next().ok_or(InvalidOodsSampleStructure)?; - let [eval] = col.try_into().map_err(|_| InvalidOodsSampleStructure)?; - Ok(eval) - })?; - - // Too many columns. - if composition_cols.next().is_some() { - return Err(InvalidOodsSampleStructure); - } - - Ok(SecureField::from_partial_evals(coordinate_evals)) + commitment_scheme.verify_values(sample_points, proof.0, channel) } /// Error when the sampled values have an invalid structure. @@ -187,7 +155,33 @@ pub enum VerificationError { ProofOfWork, } +#[derive(Debug, Serialize, Deserialize)] +pub struct StarkProof(pub CommitmentSchemeProof); + impl StarkProof { + /// Extracts the composition trace Out-Of-Domain-Sample evaluation from the mask. + fn extract_composition_oods_eval(&self) -> Result { + // TODO(andrew): `[.., composition_mask, _quotients_mask]` when add quotients commitment. + let [.., composition_mask] = &**self.sampled_values else { + return Err(InvalidOodsSampleStructure); + }; + + let mut composition_cols = composition_mask.iter(); + + let coordinate_evals = array::try_from_fn(|_| { + let col = &**composition_cols.next().ok_or(InvalidOodsSampleStructure)?; + let [eval] = col.try_into().map_err(|_| InvalidOodsSampleStructure)?; + Ok(eval) + })?; + + // Too many columns. + if composition_cols.next().is_some() { + return Err(InvalidOodsSampleStructure); + } + + Ok(SecureField::from_partial_evals(coordinate_evals)) + } + /// Returns the estimate size (in bytes) of the proof. pub fn size_estimate(&self) -> usize { SizeEstimate::size_estimate(self) @@ -195,12 +189,10 @@ impl StarkProof { /// Returns size estimates (in bytes) for different parts of the proof. pub fn size_breakdown_estimate(&self) -> StarkProofSizeBreakdown { - let Self { - commitments, - commitment_scheme_proof, - } = self; + let Self(commitment_scheme_proof) = self; let CommitmentSchemeProof { + commitments, sampled_values, decommitments, queried_values, @@ -236,6 +228,14 @@ impl StarkProof { } } +impl Deref for StarkProof { + type Target = CommitmentSchemeProof; + + fn deref(&self) -> &CommitmentSchemeProof { + &self.0 + } +} + /// Size estimate (in bytes) for different parts of the proof. pub struct StarkProofSizeBreakdown { pub oods_samples: usize, @@ -313,13 +313,15 @@ impl SizeEstimate for FriProof { impl SizeEstimate for CommitmentSchemeProof { fn size_estimate(&self) -> usize { let Self { + commitments, sampled_values, decommitments, queried_values, proof_of_work, fri_proof, } = self; - sampled_values.size_estimate() + commitments.size_estimate() + + sampled_values.size_estimate() + decommitments.size_estimate() + queried_values.size_estimate() + mem::size_of_val(proof_of_work) @@ -329,11 +331,8 @@ impl SizeEstimate for CommitmentSchemeProof { impl SizeEstimate for StarkProof { fn size_estimate(&self) -> usize { - let Self { - commitments, - commitment_scheme_proof, - } = self; - commitments.size_estimate() + commitment_scheme_proof.size_estimate() + let Self(commitment_scheme_proof) = self; + commitment_scheme_proof.size_estimate() } } diff --git a/crates/prover/src/core/vcs/prover.rs b/crates/prover/src/core/vcs/prover.rs index de8afb9a3..bc788e51f 100644 --- a/crates/prover/src/core/vcs/prover.rs +++ b/crates/prover/src/core/vcs/prover.rs @@ -79,18 +79,9 @@ impl, H: MerkleHasher> MerkleProver { /// * A `MerkleDecommitment` containing the hash and column witnesses. pub fn decommit( &self, - queries_per_log_size: BTreeMap>, + queries_per_log_size: &BTreeMap>, columns: Vec<&Col>, ) -> (ColumnVec>, MerkleDecommitment) { - // Check that queries are sorted and deduped. - // TODO(andrew): Consider using a Queries struct to prevent this. - for queries in queries_per_log_size.values() { - assert!( - queries.windows(2).all(|w| w[0] < w[1]), - "Queries are not sorted." - ); - } - // Prepare output buffers. let mut queried_values_by_layer = vec![]; let mut decommitment = MerkleDecommitment::empty(); diff --git a/crates/prover/src/core/vcs/test_utils.rs b/crates/prover/src/core/vcs/test_utils.rs index 8fc535b05..b92f9e971 100644 --- a/crates/prover/src/core/vcs/test_utils.rs +++ b/crates/prover/src/core/vcs/test_utils.rs @@ -50,7 +50,7 @@ where queries.insert(log_size, layer_queries); } - let (values, decommitment) = merkle.decommit(queries.clone(), cols.iter().collect_vec()); + let (values, decommitment) = merkle.decommit(&queries, cols.iter().collect_vec()); let verifier = MerkleVerifier { root: merkle.root(), diff --git a/crates/prover/src/examples/blake/air.rs b/crates/prover/src/examples/blake/air.rs index f219f1f5f..2db5d95af 100644 --- a/crates/prover/src/examples/blake/air.rs +++ b/crates/prover/src/examples/blake/air.rs @@ -314,7 +314,7 @@ where // Setup protocol. let channel = &mut MC::C::default(); - let commitment_scheme = &mut CommitmentSchemeProver::new(config, &twiddles); + let mut commitment_scheme = CommitmentSchemeProver::new(config, &twiddles); // Preprocessed trace. // TODO(ShaharS): share is_first column between components when constant columns support this. diff --git a/crates/prover/src/examples/plonk/mod.rs b/crates/prover/src/examples/plonk/mod.rs index 8a6d0060e..d97cabd8e 100644 --- a/crates/prover/src/examples/plonk/mod.rs +++ b/crates/prover/src/examples/plonk/mod.rs @@ -183,8 +183,8 @@ pub fn prove_fibonacci_plonk( // Setup protocol. let channel = &mut Blake2sChannel::default(); - let commitment_scheme = - &mut CommitmentSchemeProver::<_, Blake2sMerkleChannel>::new(config, &twiddles); + let mut commitment_scheme = + CommitmentSchemeProver::<_, Blake2sMerkleChannel>::new(config, &twiddles); // Preprocessed trace. let span = span!(Level::INFO, "Constant").entered(); diff --git a/crates/prover/src/examples/poseidon/mod.rs b/crates/prover/src/examples/poseidon/mod.rs index 583cb82ba..56431d92f 100644 --- a/crates/prover/src/examples/poseidon/mod.rs +++ b/crates/prover/src/examples/poseidon/mod.rs @@ -341,8 +341,8 @@ pub fn prove_poseidon( // Setup protocol. let channel = &mut Blake2sChannel::default(); - let commitment_scheme = - &mut CommitmentSchemeProver::<_, Blake2sMerkleChannel>::new(config, &twiddles); + let mut commitment_scheme = + CommitmentSchemeProver::<_, Blake2sMerkleChannel>::new(config, &twiddles); // Preprocessed trace. let span = span!(Level::INFO, "Constant").entered(); diff --git a/crates/prover/src/examples/state_machine/mod.rs b/crates/prover/src/examples/state_machine/mod.rs index 4596eb974..30fcae4fc 100644 --- a/crates/prover/src/examples/state_machine/mod.rs +++ b/crates/prover/src/examples/state_machine/mod.rs @@ -51,8 +51,8 @@ pub fn prove_state_machine( ); // Setup protocol. - let commitment_scheme = - &mut CommitmentSchemeProver::<_, Blake2sMerkleChannel>::new(config, &twiddles); + let mut commitment_scheme = + CommitmentSchemeProver::<_, Blake2sMerkleChannel>::new(config, &twiddles); // Preprocessed trace. let mut tree_builder = commitment_scheme.tree_builder(); diff --git a/crates/prover/src/examples/wide_fibonacci/mod.rs b/crates/prover/src/examples/wide_fibonacci/mod.rs index 020a24c65..afdc64caa 100644 --- a/crates/prover/src/examples/wide_fibonacci/mod.rs +++ b/crates/prover/src/examples/wide_fibonacci/mod.rs @@ -183,10 +183,8 @@ mod tests { // Setup protocol. let prover_channel = &mut Blake2sChannel::default(); - let commitment_scheme = - &mut CommitmentSchemeProver::::new( - config, &twiddles, - ); + let mut commitment_scheme = + CommitmentSchemeProver::::new(config, &twiddles); // Preprocessed trace let mut tree_builder = commitment_scheme.tree_builder(); @@ -242,10 +240,8 @@ mod tests { // Setup protocol. let prover_channel = &mut Poseidon252Channel::default(); - let commitment_scheme = - &mut CommitmentSchemeProver::::new( - config, &twiddles, - ); + let mut commitment_scheme = + CommitmentSchemeProver::::new(config, &twiddles); // TODO(ilya): remove the following once preproccessed columns are not mandatory. // Preprocessed trace diff --git a/crates/prover/src/examples/xor/gkr_lookups/mle_eval.rs b/crates/prover/src/examples/xor/gkr_lookups/mle_eval.rs index 5b1b78aeb..69acedccb 100644 --- a/crates/prover/src/examples/xor/gkr_lookups/mle_eval.rs +++ b/crates/prover/src/examples/xor/gkr_lookups/mle_eval.rs @@ -791,8 +791,8 @@ mod tests { .half_coset, ); let config = PcsConfig::default(); - let commitment_scheme = - &mut CommitmentSchemeProver::<_, Blake2sMerkleChannel>::new(config, &twiddles); + let mut commitment_scheme = + CommitmentSchemeProver::<_, Blake2sMerkleChannel>::new(config, &twiddles); let channel = &mut Blake2sChannel::default(); // TODO(ilya): remove the following once preproccessed columns are not mandatory. // Preprocessed trace @@ -866,8 +866,8 @@ mod tests { .half_coset, ); let config = PcsConfig::default(); - let commitment_scheme = - &mut CommitmentSchemeProver::<_, Blake2sMerkleChannel>::new(config, &twiddles); + let mut commitment_scheme = + CommitmentSchemeProver::<_, Blake2sMerkleChannel>::new(config, &twiddles); let channel = &mut Blake2sChannel::default(); // TODO(ilya): remove the following once preproccessed columns are not mandatory.