polish, add more tests

Author: ShuangWu121

ast/src/analyzed/mod.rs

@@ -102,6 +102,7 @@ impl<T> Analyzed<T> {
         self.declaration_type_count(PolynomialType::Committed)
     }
 
+    /// @returns the number of committed polynomials (with multiplicities for arrays) in a specific stage
     pub fn stage_commitment_count(&self, stage: Option<u32>) -> usize {
         self.stage_declaration_type_count(PolynomialType::Committed, stage)
     }

backend/src/stwo/circuit_builder.rs

@@ -18,6 +18,7 @@ use stwo_prover::core::poly::circle::{CircleDomain, CircleEvaluation};
 use stwo_prover::core::poly::BitReversedOrder;
 use stwo_prover::core::utils::{bit_reverse_index, coset_index_to_circle_domain_index};
 
+pub const PREPROCESSED_TRACE_IDX: usize = 0;
 pub const STAGE0_TRACE_IDX: usize = 1;
 pub const STAGE1_TRACE_IDX: usize = 2;
 
@@ -85,6 +86,10 @@ impl PowdrEval {
             .filter(|(symbol, _)| symbol.stage == Some(1))
             .flat_map(|(symbol, _)| symbol.array_elements())
             .enumerate()
+            // During the circuit-building phase, stage1 witness columns are indexed after stage0 witness columns,
+            // resulting in a concatenation of witness columns from both stages.
+            // In the proving phase, all stage1 witness columns are appended after stage0 witness columns
+            // to maintain consistency with the circuit-building phase.
             .map(|(index, (_, id))| (id, index + witness_columns_stage0.len()))
             .collect();
 
@@ -217,7 +222,9 @@ impl FrameworkEval for PowdrEval {
             .iter()
             .map(|(k, v)| (*k, E::F::from(into_stwo_field(v))))
             .collect();
+        // concatenate stage0 and stage1 witness columns
         witness_eval_stage0.extend(witness_eval_stage1);
+
         let intermediate_definitions = self.analyzed.intermediate_definitions();
         let data = Data {
             witness_eval: &witness_eval_stage0,

backend/src/stwo/prover.rs

@@ -20,13 +20,13 @@ use std::{fmt, io};
 
 use crate::stwo::circuit_builder::{
     gen_stwo_circle_column, get_constant_with_next_list, PowdrComponent, PowdrEval,
-    STAGE0_TRACE_IDX, STAGE1_TRACE_IDX,
+    PREPROCESSED_TRACE_IDX, STAGE0_TRACE_IDX, STAGE1_TRACE_IDX,
 };
 use crate::stwo::proof::{
     Proof, SerializableStarkProvingKey, StarkProvingKey, TableProvingKey, TableProvingKeyCollection,
 };
 
-use stwo_prover::constraint_framework::{TraceLocationAllocator, PREPROCESSED_TRACE_IDX};
+use stwo_prover::constraint_framework::TraceLocationAllocator;
 
 use stwo_prover::core::air::{Component, ComponentProver};
 use stwo_prover::core::backend::{Backend, BackendForChannel};
@@ -459,6 +459,7 @@ where
         let mut stage0_witness_col_log_sizes = vec![];
         let mut stage1_witness_col_log_sizes = vec![];
 
+        // get the column sizes for stage 0 and stage 1, used by verifier to read the commitments from the proof
         self.split
             .iter()
             .zip_eq(proof.machine_log_sizes.iter())
@@ -477,6 +478,7 @@ where
                 },
             );
 
+        // Verifier gets the commitments of the constant columns and stage 0 witness columns (two Merkle tree roots)
         commitment_scheme.commit(
             proof.stark_proof.commitments[PREPROCESSED_TRACE_IDX],
             &constant_col_log_sizes,
@@ -489,10 +491,10 @@ where
             verifier_channel,
         );
 
-        // TODO: make the challenge sound, now the challenge is built the same way in prover.
+        // Get challenges based on the commitments of constant columns and stage 0 witness columns
         let stage0_challenges = get_challenges::<MC>(&self.analyzed, verifier_channel);
 
-        let mut components = self
+        let components = self
             .split
             .iter()
             .zip_eq(proof.machine_log_sizes.iter())
@@ -519,13 +521,11 @@ where
             )
             .collect::<Vec<_>>();
 
-        let mut components_slice: Vec<&dyn Component> = components
-            .iter_mut()
+        let components_slice: Vec<&dyn Component> = components
+            .iter()
             .map(|component| component as &dyn Component)
             .collect();
 
-        let components_slice = components_slice.as_mut_slice();
-
         if self.analyzed.stage_count() > 1 {
             commitment_scheme.commit(
                 proof.stark_proof.commitments[STAGE1_TRACE_IDX],
@@ -535,7 +535,7 @@ where
         }
 
         stwo_prover::core::prover::verify(
-            components_slice,
+            &components_slice,
             verifier_channel,
             commitment_scheme,
             proof.stark_proof,

pipeline/tests/asm.rs

@@ -73,6 +73,8 @@ fn second_phase_hint() {
     let pipeline = make_simple_prepared_pipeline::<GoldilocksField>(f, LinkerMode::Bus);
     test_mock_backend(pipeline.clone());
     test_plonky3_pipeline(pipeline);
+    let pipeline = make_simple_prepared_pipeline::<Mersenne31Field>(f, LinkerMode::Bus);
+    test_stwo_pipeline(pipeline);
 }
 
 #[test]