diff --git a/uni-stark/src/check_constraints.rs b/uni-stark/src/check_constraints.rs index 6a7b48cd1..d033ef4d2 100644 --- a/uni-stark/src/check_constraints.rs +++ b/uni-stark/src/check_constraints.rs @@ -14,15 +14,15 @@ use crate::traits::MultistageAirBuilder; pub(crate) fn check_constraints( air: &A, preprocessed: &RowMajorMatrix, - stages: Vec<&RowMajorMatrix>, - public_values: &Vec<&Vec>, + traces_by_stage: Vec<&RowMajorMatrix>, + public_values_by_stage: &Vec<&Vec>, challenges: Vec<&Vec>, ) where F: Field, A: for<'a> Air>, { - let num_stages = stages.len(); - let height = stages[0].height(); + let num_stages = traces_by_stage.len(); + let height = traces_by_stage[0].height(); (0..height).for_each(|i| { let i_next = (i + 1) % height; @@ -34,7 +34,7 @@ pub(crate) fn check_constraints( RowMajorMatrixView::new_row(&*next_preprocessed), ); - let stages_local_next = stages + let stages_local_next = traces_by_stage .iter() .map(|trace| { let stage_local = trace.row_slice(i); @@ -43,7 +43,7 @@ pub(crate) fn check_constraints( }) .collect_vec(); - let stages = (0..num_stages) + let traces_by_stage = (0..num_stages) .map(|stage| { VerticalPair::new( RowMajorMatrixView::new_row(&*stages_local_next[stage].0), @@ -56,8 +56,8 @@ pub(crate) fn check_constraints( row_index: i, challenges: challenges.clone(), preprocessed, - stages, - public_values, + traces_by_stage, + public_values_by_stage, is_first_row: F::from_bool(i == 0), is_last_row: F::from_bool(i == height - 1), is_transition: F::from_bool(i != height - 1), @@ -74,8 +74,8 @@ pub struct DebugConstraintBuilder<'a, F: Field> { row_index: usize, preprocessed: VerticalPair, RowMajorMatrixView<'a, F>>, challenges: Vec<&'a Vec>, - stages: Vec, RowMajorMatrixView<'a, F>>>, - public_values: &'a [&'a Vec], + traces_by_stage: Vec, RowMajorMatrixView<'a, F>>>, + public_values_by_stage: &'a [&'a Vec], is_first_row: F, is_last_row: F, is_transition: F, @@ -107,7 +107,7 @@ where } fn main(&self) -> Self::M { - self.stages[0] + self.traces_by_stage[0] } fn assert_zero>(&mut self, x: I) { @@ -148,11 +148,11 @@ impl<'a, F: Field> MultistageAirBuilder for DebugConstraintBuilder<'a, F> { type Challenge = Self::Expr; fn stage_public_values(&self, stage: usize) -> &[Self::F] { - self.public_values[stage] + self.public_values_by_stage[stage] } fn stage_trace(&self, stage: usize) -> Self::M { - self.stages[stage] + self.traces_by_stage[stage] } fn stage_challenges(&self, stage: usize) -> &[Self::Expr] { diff --git a/uni-stark/src/folder.rs b/uni-stark/src/folder.rs index 307337922..fb280ee37 100644 --- a/uni-stark/src/folder.rs +++ b/uni-stark/src/folder.rs @@ -11,9 +11,9 @@ use crate::{PackedChallenge, PackedVal, StarkGenericConfig, Val}; #[derive(Debug)] pub struct ProverConstraintFolder<'a, SC: StarkGenericConfig> { pub challenges: Vec>>, - pub stages: Vec>>, + pub traces_by_stage: Vec>>, pub preprocessed: RowMajorMatrix>, - pub public_values: &'a Vec>>, + pub public_values_by_stage: &'a Vec>>, pub is_first_row: PackedVal, pub is_last_row: PackedVal, pub is_transition: PackedVal, @@ -26,9 +26,9 @@ type ViewPair<'a, T> = VerticalPair, RowMajorMatrixVie #[derive(Debug)] pub struct VerifierConstraintFolder<'a, SC: StarkGenericConfig> { pub challenges: Vec>>, - pub stages: Vec>, + pub traces_by_stage: Vec>, pub preprocessed: ViewPair<'a, SC::Challenge>, - pub public_values: Vec<&'a Vec>>, + pub public_values_by_stage: Vec<&'a Vec>>, pub is_first_row: SC::Challenge, pub is_last_row: SC::Challenge, pub is_transition: SC::Challenge, @@ -43,7 +43,7 @@ impl<'a, SC: StarkGenericConfig> AirBuilder for ProverConstraintFolder<'a, SC> { type M = RowMajorMatrix>; fn main(&self) -> Self::M { - self.stages[0].clone() + self.traces_by_stage[0].clone() } fn is_first_row(&self) -> Self::Expr { @@ -81,14 +81,14 @@ impl<'a, SC: StarkGenericConfig> MultistageAirBuilder for ProverConstraintFolder type Challenge = Val; fn stage_trace(&self, stage: usize) -> ::M { - self.stages[stage].clone() + self.traces_by_stage[stage].clone() } fn stage_challenges(&self, stage: usize) -> &[Self::Challenge] { &self.challenges[stage] } fn stage_public_values(&self, stage: usize) -> &[Self::PublicVar] { - &self.public_values[stage] + &self.public_values_by_stage[stage] } } @@ -105,7 +105,7 @@ impl<'a, SC: StarkGenericConfig> AirBuilder for VerifierConstraintFolder<'a, SC> type M = ViewPair<'a, SC::Challenge>; fn main(&self) -> Self::M { - self.stages[0] + self.traces_by_stage[0] } fn is_first_row(&self) -> Self::Expr { @@ -143,14 +143,14 @@ impl<'a, SC: StarkGenericConfig> MultistageAirBuilder for VerifierConstraintFold type Challenge = Val; fn stage_trace(&self, stage: usize) -> ::M { - self.stages[stage] + self.traces_by_stage[stage] } fn stage_challenges(&self, stage: usize) -> &[Self::Challenge] { &self.challenges[stage] } fn stage_public_values(&self, stage: usize) -> &[Self::PublicVar] { - self.public_values[stage] + self.public_values_by_stage[stage] } } diff --git a/uni-stark/src/prover.rs b/uni-stark/src/prover.rs index e69540a5a..06d5c6ca8 100644 --- a/uni-stark/src/prover.rs +++ b/uni-stark/src/prover.rs @@ -194,7 +194,7 @@ where .map(|stage| stage.challenge_values.clone()) .collect(); - let public_values = state + let public_values_by_stage = state .processed_stages .iter() .map(|stage| stage.public_values.clone()) @@ -202,7 +202,7 @@ where let quotient_values = quotient_values( air, - &public_values, + &public_values_by_stage, trace_domain, quotient_domain, preprocessed_on_quotient_domain, @@ -304,7 +304,7 @@ where #[instrument(name = "compute quotient polynomial", skip_all)] fn quotient_values<'a, SC, A, Mat>( air: &A, - public_values: &'a Vec>>, + public_values_by_stage: &'a Vec>>, trace_domain: Domain, quotient_domain: Domain, preprocessed_on_quotient_domain: Option, @@ -360,7 +360,7 @@ where preprocessed_width, ); - let stages = traces_on_quotient_domain + let traces_by_stage = traces_on_quotient_domain .iter() .map(|trace_on_quotient_domain| { RowMajorMatrix::new( @@ -378,9 +378,9 @@ where let accumulator = PackedChallenge::::zero(); let mut folder = ProverConstraintFolder { challenges: challenges.clone(), - stages, + traces_by_stage, preprocessed, - public_values, + public_values_by_stage, is_first_row, is_last_row, is_transition, diff --git a/uni-stark/src/symbolic_builder.rs b/uni-stark/src/symbolic_builder.rs index 767ad1961..6516b48e5 100644 --- a/uni-stark/src/symbolic_builder.rs +++ b/uni-stark/src/symbolic_builder.rs @@ -50,7 +50,7 @@ where A: MultiStageAir>, { let widths: Vec<_> = (0..air.stage_count()) - .map(|i| air.stage_width(i as u32)) + .map(|i| air.stage_trace_width(i as u32)) .collect(); let challenges: Vec<_> = (0..air.stage_count()) .map(|i| air.stage_challenge_count(i as u32)) @@ -70,8 +70,8 @@ where pub struct SymbolicAirBuilder { challenges: Vec>>, preprocessed: RowMajorMatrix>, - stages: Vec>>, - public_values: Vec>>, + traces_by_stage: Vec>>, + public_values_by_stage: Vec>>, constraints: Vec>, } @@ -89,7 +89,7 @@ impl SymbolicAirBuilder { .map(move |index| SymbolicVariable::new(Entry::Preprocessed { offset }, index)) }) .collect(); - let stages = stage_widths + let traces_by_stage = stage_widths .iter() .map(|width| { let values = [0, 1] @@ -116,7 +116,7 @@ impl SymbolicAirBuilder { }) .collect(); let mut public_value_index = 0; - let public_values = public_value_counts + let public_values_by_stage = public_value_counts .iter() .map(|count| { (0..*count) @@ -131,8 +131,8 @@ impl SymbolicAirBuilder { Self { challenges, preprocessed: RowMajorMatrix::new(prep_values, preprocessed_width), - stages, - public_values, + traces_by_stage, + public_values_by_stage, constraints: vec![], } } @@ -149,7 +149,7 @@ impl AirBuilder for SymbolicAirBuilder { type M = RowMajorMatrix; fn main(&self) -> Self::M { - self.stages[0].clone() + self.traces_by_stage[0].clone() } fn is_first_row(&self) -> Self::Expr { @@ -185,11 +185,11 @@ impl MultistageAirBuilder for SymbolicAirBuilder { type Challenge = Self::Var; fn stage_public_values(&self, stage: usize) -> &[Self::PublicVar] { - &self.public_values[stage] + &self.public_values_by_stage[stage] } fn stage_trace(&self, stage: usize) -> Self::M { - self.stages[stage].clone() + self.traces_by_stage[stage].clone() } fn stage_challenges(&self, stage: usize) -> &[Self::Challenge] { diff --git a/uni-stark/src/traits.rs b/uni-stark/src/traits.rs index 8314b0ceb..30d159481 100644 --- a/uni-stark/src/traits.rs +++ b/uni-stark/src/traits.rs @@ -9,6 +9,7 @@ pub trait MultistageAirBuilder: AirBuilderWithPublicValues { /// Challenges from each stage, drawn from the base field fn stage_challenges(&self, stage: usize) -> &[Self::Challenge]; + /// Public values for each stage fn stage_public_values(&self, stage: usize) -> &[Self::PublicVar] { match stage { 0 => self.public_values(), @@ -22,8 +23,8 @@ pub trait MultiStageAir: Air { 1 } - /// The number of columns in a given higher-stage trace. - fn stage_width(&self, stage: u32) -> usize { + /// The number of trace columns in this stage + fn stage_trace_width(&self, stage: u32) -> usize { match stage { 0 => self.width(), _ => unimplemented!(), diff --git a/uni-stark/src/verifier.rs b/uni-stark/src/verifier.rs index ecc2064cf..8dece8b68 100644 --- a/uni-stark/src/verifier.rs +++ b/uni-stark/src/verifier.rs @@ -40,7 +40,7 @@ pub fn verify_with_key( air: &A, challenger: &mut SC::Challenger, proof: &Proof, - public_values: Vec<&Vec>>, + public_values_by_stage: Vec<&Vec>>, ) -> Result<(), VerificationError>> where SC: StarkGenericConfig, @@ -57,7 +57,7 @@ where let degree = 1 << degree_bits; let log_quotient_degree = get_log_quotient_degree::, A>( air, - &public_values + &public_values_by_stage .iter() .map(|values| values.len()) .collect::>(), @@ -76,7 +76,7 @@ where let air_widths = (0..stage_count) .map(|stage| { - >>>::stage_width(air, stage as u32) + >>>::stage_trace_width(air, stage as u32) }) .collect::>(); let air_fixed_width = >>::preprocessed_width(air); @@ -97,8 +97,9 @@ where .quotient_chunks .iter() .all(|qc| qc.len() == >>::D) - && public_values.len() == stage_count + && public_values_by_stage.len() == stage_count && challenge_counts.len() == stage_count; + if !valid_shape { return Err(VerificationError::InvalidProofShape); } @@ -120,7 +121,7 @@ where commitments .traces_by_stage .iter() - .zip(&public_values) + .zip(&public_values_by_stage) .zip(challenge_counts) .for_each(|((commitment, public_values), challenge_count)| { challenger.observe(commitment.clone()); @@ -220,7 +221,7 @@ where RowMajorMatrixView::new_row(&opened_values.preprocessed_next), ); - let stages = opened_values + let traces_by_stage = opened_values .traces_by_stage_local .iter() .zip(opened_values.traces_by_stage_next.iter()) @@ -235,8 +236,8 @@ where let mut folder = VerifierConstraintFolder { challenges, preprocessed, - stages, - public_values, + traces_by_stage, + public_values_by_stage, is_first_row: sels.is_first_row, is_last_row: sels.is_last_row, is_transition: sels.is_transition,