address review comments

uni-stark/src/proof.rs

@@ -26,21 +26,20 @@ pub struct Proof<SC: StarkGenericConfig> {
     pub(crate) opened_values: OpenedValues<SC::Challenge>,
     pub(crate) opening_proof: PcsProof<SC>,
     pub(crate) degree_bits: usize,
-    pub(crate) challenge_counts: Vec<usize>,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct Commitments<Com> {
-    pub(crate) stages: Vec<Com>,
+    pub(crate) traces_by_stage: Vec<Com>,
     pub(crate) quotient_chunks: Com,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct OpenedValues<Challenge> {
     pub(crate) preprocessed_local: Vec<Challenge>,
     pub(crate) preprocessed_next: Vec<Challenge>,
-    pub(crate) stages_local: Vec<Vec<Challenge>>,
-    pub(crate) stages_next: Vec<Vec<Challenge>>,
+    pub(crate) traces_by_stage_local: Vec<Vec<Challenge>>,
+    pub(crate) traces_by_stage_next: Vec<Vec<Challenge>>,
     pub(crate) quotient_chunks: Vec<Vec<Challenge>>,
 }
 

uni-stark/src/prover.rs

@@ -222,7 +222,7 @@ where
     challenger.observe(quotient_commit.clone());
 
     let commitments = Commitments {
-        stages: state
+        traces_by_stage: state
             .processed_stages
             .iter()
             .map(|s| s.commitment.clone())
@@ -272,7 +272,7 @@ where
     };
 
     // get values for the traces
-    let (stages_local, stages_next): (Vec<_>, Vec<_>) = state
+    let (traces_by_stage_local, traces_by_stage_next): (Vec<_>, Vec<_>) = state
         .processed_stages
         .iter()
         .map(|_| {
@@ -289,8 +289,8 @@ where
     let quotient_chunks = value.iter().map(|v| v[0].clone()).collect_vec();
 
     let opened_values = OpenedValues {
-        stages_local,
-        stages_next,
+        traces_by_stage_local,
+        traces_by_stage_next,
         preprocessed_local,
         preprocessed_next,
         quotient_chunks,
@@ -300,11 +300,6 @@ where
         opened_values,
         opening_proof,
         degree_bits: log_degree,
-        challenge_counts: state
-            .processed_stages
-            .iter()
-            .map(|s| s.challenge_values.len())
-            .collect(),
     }
 }
 
@@ -441,14 +436,13 @@ impl<'a, SC: StarkGenericConfig> ProverState<'a, SC> {
             .in_scope(|| self.pcs.commit(vec![(self.trace_domain, stage.trace)]));
 
         self.challenger.observe(commitment.clone());
+        // observe the public inputs for this stage
+        self.challenger.observe_slice(&stage.public_values);
 
         let challenge_values = (0..stage.challenge_count)
             .map(|_| self.challenger.sample())
             .collect();
 
-        // observe the public inputs for this stage
-        self.challenger.observe_slice(&stage.public_values);
-
         self.processed_stages.push(ProcessedStage {
             public_values: stage.public_values,
             prover_data,

uni-stark/src/verifier.rs

@@ -52,7 +52,6 @@ where
         opened_values,
         opening_proof,
         degree_bits,
-        challenge_counts,
     } = proof;
 
     let degree = 1 << degree_bits;
@@ -64,15 +63,18 @@ where
             .collect::<Vec<_>>(),
     );
     let quotient_degree = 1 << log_quotient_degree;
-    let stages = proof.commitments.stages.len();
+    let stage_count = proof.commitments.traces_by_stage.len();
+    let challenge_counts: Vec<usize> = (0..stage_count)
+        .map(|i| <A as MultiStageAir<SymbolicAirBuilder<_>>>::stage_challenge_count(air, i as u32))
+        .collect();
 
     let pcs = config.pcs();
     let trace_domain = pcs.natural_domain_for_degree(degree);
     let quotient_domain =
         trace_domain.create_disjoint_domain(1 << (degree_bits + log_quotient_degree));
     let quotient_chunks_domains = quotient_domain.split_domains(quotient_degree);
 
-    let air_widths = (0..stages)
+    let air_widths = (0..stage_count)
         .map(|stage| {
             <A as MultiStageAir<SymbolicAirBuilder<Val<SC>>>>::stage_width(air, stage as u32)
         })
@@ -81,12 +83,12 @@ where
     let valid_shape = opened_values.preprocessed_local.len() == air_fixed_width
         && opened_values.preprocessed_next.len() == air_fixed_width
         && opened_values
-            .stages_local
+            .traces_by_stage_local
             .iter()
             .zip(&air_widths)
             .all(|(stage, air_width)| stage.len() == *air_width)
         && opened_values
-            .stages_next
+            .traces_by_stage_next
             .iter()
             .zip(&air_widths)
             .all(|(stage, air_width)| stage.len() == *air_width)
@@ -95,8 +97,8 @@ where
             .quotient_chunks
             .iter()
             .all(|qc| qc.len() == <SC::Challenge as AbstractExtensionField<Val<SC>>>::D)
-        && public_values.len() == stages
-        && challenge_counts.len() == stages;
+        && public_values.len() == stage_count
+        && challenge_counts.len() == stage_count;
     if !valid_shape {
         return Err(VerificationError::InvalidProofShape);
     }
@@ -116,14 +118,14 @@ where
     let mut challenges = vec![];
 
     commitments
-        .stages
+        .traces_by_stage
         .iter()
         .zip(&public_values)
         .zip(challenge_counts)
         .for_each(|((commitment, public_values), challenge_count)| {
             challenger.observe(commitment.clone());
-            challenges.push((0..*challenge_count).map(|_| challenger.sample()).collect());
             challenger.observe_slice(public_values);
+            challenges.push((0..challenge_count).map(|_| challenger.sample()).collect());
         });
     let alpha: SC::Challenge = challenger.sample_ext_element();
     challenger.observe(commitments.quotient_chunks.clone());
@@ -151,9 +153,9 @@ where
             )
             .chain(
                 izip!(
-                    commitments.stages.iter(),
-                    opened_values.stages_local.iter(),
-                    opened_values.stages_next.iter()
+                    commitments.traces_by_stage.iter(),
+                    opened_values.traces_by_stage_local.iter(),
+                    opened_values.traces_by_stage_next.iter()
                 )
                 .map(|(trace_commit, opened_local, opened_next)| {
                     (
@@ -219,9 +221,9 @@ where
     );
 
     let stages = opened_values
-        .stages_local
+        .traces_by_stage_local
         .iter()
-        .zip(opened_values.stages_next.iter())
+        .zip(opened_values.traces_by_stage_next.iter())
         .map(|(trace_local, trace_next)| {
             VerticalPair::new(
                 RowMajorMatrixView::new_row(trace_local),