Stwo challenges sound (#2421)

# Make Challenge Function in `stwo` Sound

## Summary
This PR adds soundness to stwo challenge function

## Main Changes
- **Separation of Stage 0 and Stage 1 Witness Columns**  
separate stage 0 and stage 1 witness columns in circuit building phase
(PowdrEval) and proving phase
- **Make the challenge generation based on commitments of constant cols
and stage 0 witness cols**
- **Add more stwo tests to existing tests**

---------

Co-authored-by: Georg Wiese <[email protected]>
Co-authored-by: Thibaut Schaeffer <[email protected]>

ast/src/analyzed/mod.rs

@@ -103,6 +103,12 @@ impl<T> Analyzed<T> {
     pub fn commitment_count(&self) -> usize {
         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: u32) -> usize {
+        self.stage_declaration_type_count(PolynomialType::Committed, stage)
+    }
+
     /// @returns the number of intermediate polynomials (with multiplicities for arrays)
     pub fn intermediate_count(&self) -> usize {
         self.intermediate_columns
@@ -202,6 +208,19 @@ impl<T> Analyzed<T> {
             .sum()
     }
 
+    fn stage_declaration_type_count(&self, poly_type: PolynomialType, stage: u32) -> usize {
+        self.definitions
+            .values()
+            .filter(|(symbol, _)| symbol.stage.unwrap_or(0) == stage)
+            .filter_map(move |(symbol, _)| match symbol.kind {
+                SymbolKind::Poly(ptype) if ptype == poly_type => {
+                    Some(symbol.length.unwrap_or(1) as usize)
+                }
+                _ => None,
+            })
+            .sum()
+    }
+
     /// Returns the type (scheme) of a symbol with the given name.
     pub fn type_of_symbol(&self, name: &str) -> TypeScheme {
         let (sym, value) = &self.definitions[name];

backend/src/stwo/circuit_builder.rs

@@ -10,16 +10,18 @@ use powdr_ast::analyzed::{AlgebraicExpression, AlgebraicReference, Analyzed, Cha
 use powdr_ast::analyzed::{PolyID, PolynomialType};
 use powdr_number::Mersenne31Field as M31;
 use stwo_prover::constraint_framework::preprocessed_columns::PreprocessedColumn;
-use stwo_prover::constraint_framework::{
-    EvalAtRow, FrameworkComponent, FrameworkEval, ORIGINAL_TRACE_IDX,
-};
+use stwo_prover::constraint_framework::{EvalAtRow, FrameworkComponent, FrameworkEval};
 use stwo_prover::core::backend::{Column, ColumnOps};
 use stwo_prover::core::fields::m31::BaseField;
 use stwo_prover::core::fields::{ExtensionOf, FieldOps};
 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;
+
 pub type PowdrComponent = FrameworkComponent<PowdrEval>;
 
 pub fn gen_stwo_circle_column<B, F>(
@@ -56,11 +58,13 @@ pub struct PowdrEval {
     // the pre-processed are indexed in the whole proof, instead of in each component.
     // this offset represents the index of the first pre-processed column in this component
     preprocess_col_offset: usize,
-    witness_columns: BTreeMap<PolyID, usize>,
+    stage0_witness_columns: BTreeMap<PolyID, usize>,
+    stage1_witness_columns: BTreeMap<PolyID, usize>,
     constant_shifted: BTreeMap<PolyID, usize>,
     constant_columns: BTreeMap<PolyID, usize>,
     // stwo supports maximum 2 stages, challenges are only created after stage 0
     pub challenges: BTreeMap<u64, M31>,
+    poly_stage_map: BTreeMap<PolyID, usize>,
 }
 
 impl PowdrEval {
@@ -70,8 +74,17 @@ impl PowdrEval {
         log_degree: u32,
         challenges: BTreeMap<u64, M31>,
     ) -> Self {
-        let witness_columns: BTreeMap<PolyID, usize> = analyzed
+        let stage0_witness_columns: BTreeMap<PolyID, usize> = analyzed
             .definitions_in_source_order(PolynomialType::Committed)
+            .filter(|(symbol, _)| symbol.stage.unwrap_or(0) == 0)
+            .flat_map(|(symbol, _)| symbol.array_elements())
+            .enumerate()
+            .map(|(index, (_, id))| (id, index))
+            .collect();
+
+        let stage1_witness_columns: BTreeMap<PolyID, usize> = analyzed
+            .definitions_in_source_order(PolynomialType::Committed)
+            .filter(|(symbol, _)| symbol.stage.unwrap_or(0) == 1)
             .flat_map(|(symbol, _)| symbol.array_elements())
             .enumerate()
             .map(|(index, (_, id))| (id, index))
@@ -94,33 +107,48 @@ impl PowdrEval {
             .map(|(index, (_, id))| (id, index))
             .collect();
 
+        let poly_stage_map: BTreeMap<PolyID, usize> = stage0_witness_columns
+            .keys()
+            .map(|k| (*k, 0))
+            .chain(stage1_witness_columns.keys().map(|k| (*k, 1)))
+            .collect();
+
         Self {
             log_degree,
             analyzed,
             preprocess_col_offset,
-            witness_columns,
+            stage0_witness_columns,
+            stage1_witness_columns,
             constant_shifted,
             constant_columns,
             challenges,
+            poly_stage_map,
         }
     }
 }
 
 struct Data<'a, F> {
-    witness_eval: &'a BTreeMap<PolyID, [F; 2]>,
+    stage0_witness_eval: &'a BTreeMap<PolyID, [F; 2]>,
+    stage1_witness_eval: &'a BTreeMap<PolyID, [F; 2]>,
     constant_shifted_eval: &'a BTreeMap<PolyID, F>,
     constant_eval: &'a BTreeMap<PolyID, F>,
     // challenges for stage 1
     challenges: &'a BTreeMap<u64, F>,
+    poly_stage_map: &'a BTreeMap<PolyID, usize>,
 }
 
 impl<F: Clone> TerminalAccess<F> for &Data<'_, F> {
     fn get(&self, poly_ref: &AlgebraicReference) -> F {
         match poly_ref.poly_id.ptype {
-            PolynomialType::Committed => match poly_ref.next {
-                false => self.witness_eval[&poly_ref.poly_id][0].clone(),
-                true => self.witness_eval[&poly_ref.poly_id][1].clone(),
-            },
+            PolynomialType::Committed => {
+                match (self.poly_stage_map[&poly_ref.poly_id], poly_ref.next) {
+                    (0, false) => self.stage0_witness_eval[&poly_ref.poly_id][0].clone(),
+                    (0, true) => self.stage0_witness_eval[&poly_ref.poly_id][1].clone(),
+                    (1, false) => self.stage1_witness_eval[&poly_ref.poly_id][0].clone(),
+                    (1, true) => self.stage1_witness_eval[&poly_ref.poly_id][1].clone(),
+                    _ => unreachable!(),
+                }
+            }
             PolynomialType::Constant => match poly_ref.next {
                 false => self.constant_eval[&poly_ref.poly_id].clone(),
                 true => self.constant_shifted_eval[&poly_ref.poly_id].clone(),
@@ -151,13 +179,24 @@ impl FrameworkEval for PowdrEval {
             "Error: Expected no public inputs, as they are not supported yet.",
         );
 
-        let witness_eval: BTreeMap<PolyID, [<E as EvalAtRow>::F; 2]> = self
-            .witness_columns
+        let stage0_witness_eval: BTreeMap<PolyID, [<E as EvalAtRow>::F; 2]> = self
+            .stage0_witness_columns
+            .keys()
+            .map(|poly_id| {
+                (
+                    *poly_id,
+                    eval.next_interaction_mask(STAGE0_TRACE_IDX, [0, 1]),
+                )
+            })
+            .collect();
+
+        let stage1_witness_eval: BTreeMap<PolyID, [<E as EvalAtRow>::F; 2]> = self
+            .stage1_witness_columns
             .keys()
             .map(|poly_id| {
                 (
                     *poly_id,
-                    eval.next_interaction_mask(ORIGINAL_TRACE_IDX, [0, 1]),
+                    eval.next_interaction_mask(STAGE1_TRACE_IDX, [0, 1]),
                 )
             })
             .collect();
@@ -197,10 +236,12 @@ impl FrameworkEval for PowdrEval {
 
         let intermediate_definitions = self.analyzed.intermediate_definitions();
         let data = Data {
-            witness_eval: &witness_eval,
+            stage0_witness_eval: &stage0_witness_eval,
+            stage1_witness_eval: &stage1_witness_eval,
             constant_shifted_eval: &constant_shifted_eval,
             constant_eval: &constant_eval,
             challenges: &challenges,
+            poly_stage_map: &self.poly_stage_map,
         };
         let mut evaluator =
             ExpressionEvaluator::new_with_custom_expr(&data, &intermediate_definitions, |v| {