Extend stwo to Support Machines with Varying Degrees (#2280)

### Summary
This PR introduces support for proving machines with differing degrees
in the `stwo` prover. It enables handling proving columns of varying
sizes.

---

### Details
1. **Support for Proving Columns of Different Sizes**:
- Machines with different column sizes require their own distinct
components.
- All components are stored in a single vector, from which a slice is
created and passed into the `prove` function to handle the proving
process.
- Constant columns are indexed globally across all components to ensure
consistency and correctness.

2. **Verification**:
- The verification function requires the sizes of all columns (including
constant columns and witness columns) as input.

---

### Changes
- Updated the logic for managing machine-specific components and
assembling them for the proving process.
- Implemented indexing of constant columns across components.

---

### Testing
- Added tests to validate proving functionality with machines of
differing degrees for stwo prover.


---

---------

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

backend/src/stwo/circuit_builder.rs

@@ -2,7 +2,6 @@ use core::unreachable;
 use powdr_ast::parsed::visitor::AllChildren;
 use powdr_executor_utils::expression_evaluator::{ExpressionEvaluator, TerminalAccess};
 use std::collections::HashSet;
-use std::sync::Arc;
 
 extern crate alloc;
 use alloc::collections::btree_map::BTreeMap;
@@ -53,14 +52,18 @@ where
 }
 
 pub struct PowdrEval<T> {
-    analyzed: Arc<Analyzed<T>>,
+    log_degree: u32,
+    analyzed: Analyzed<T>,
+    // 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>,
     constant_shifted: BTreeMap<PolyID, usize>,
     constant_columns: BTreeMap<PolyID, usize>,
 }
 
 impl<T: FieldElement> PowdrEval<T> {
-    pub fn new(analyzed: Arc<Analyzed<T>>) -> Self {
+    pub fn new(analyzed: Analyzed<T>, preprocess_col_offset: usize, log_degree: u32) -> Self {
         let witness_columns: BTreeMap<PolyID, usize> = analyzed
             .definitions_in_source_order(PolynomialType::Committed)
             .flat_map(|(symbol, _)| symbol.array_elements())
@@ -86,7 +89,9 @@ impl<T: FieldElement> PowdrEval<T> {
             .collect();
 
         Self {
+            log_degree,
             analyzed,
+            preprocess_col_offset,
             witness_columns,
             constant_shifted,
             constant_columns,
@@ -126,10 +131,10 @@ impl<F: Clone> TerminalAccess<F> for &Data<'_, F> {
 
 impl<T: FieldElement> FrameworkEval for PowdrEval<T> {
     fn log_size(&self) -> u32 {
-        self.analyzed.degree().ilog2()
+        self.log_degree
     }
     fn max_constraint_log_degree_bound(&self) -> u32 {
-        self.analyzed.degree().ilog2() + 1
+        self.log_degree + 1
     }
     fn evaluate<E: EvalAtRow>(&self, mut eval: E) -> E {
         assert!(
@@ -155,8 +160,9 @@ impl<T: FieldElement> FrameworkEval for PowdrEval<T> {
             .map(|(i, poly_id)| {
                 (
                     *poly_id,
-                    // PreprocessedColumn::Plonk(i) is unused argument in get_preprocessed_column
-                    eval.get_preprocessed_column(PreprocessedColumn::Plonk(i)),
+                    eval.get_preprocessed_column(PreprocessedColumn::Plonk(
+                        i + self.preprocess_col_offset,
+                    )),
                 )
             })
             .collect();
@@ -169,7 +175,7 @@ impl<T: FieldElement> FrameworkEval for PowdrEval<T> {
                 (
                     *poly_id,
                     eval.get_preprocessed_column(PreprocessedColumn::Plonk(
-                        i + constant_eval.len(),
+                        i + constant_eval.len() + self.preprocess_col_offset,
                     )),
                 )
             })

backend/src/stwo/mod.rs

@@ -39,10 +39,6 @@ impl<F: FieldElement> BackendFactory<F> for RestrictedFactory {
             return Err(Error::BackendError("Proving key unused".to_string()));
         }
 
-        if pil.degrees().len() > 1 {
-            return Err(Error::NoVariableDegreeAvailable);
-        }
-
         let mut stwo: Box<StwoProver<F, SimdBackend, Blake2sMerkleChannel, Blake2sChannel>> =
             Box::new(StwoProver::new(pil, fixed)?);
 

backend/src/stwo/proof.rs

@@ -1,13 +1,15 @@
-use serde::Deserialize;
-use serde::Serialize;
+use serde::de::DeserializeOwned;
+use serde::{Deserialize, Serialize};
 use std::collections::BTreeMap;
 use stwo_prover::core::backend::Backend;
 use stwo_prover::core::backend::Column;
 use stwo_prover::core::backend::ColumnOps;
+use stwo_prover::core::channel::MerkleChannel;
 use stwo_prover::core::fields::m31::BaseField;
 use stwo_prover::core::fields::m31::M31;
 use stwo_prover::core::poly::circle::{CanonicCoset, CircleEvaluation};
 use stwo_prover::core::poly::BitReversedOrder;
+use stwo_prover::core::prover::StarkProof;
 use stwo_prover::core::ColumnVec;
 
 /// For each possible size, the commitment and prover data
@@ -124,3 +126,14 @@ impl<B: Backend> From<StarkProvingKey<B>> for SerializableStarkProvingKey {
         Self { preprocessed }
     }
 }
+
+#[derive(Serialize, Deserialize)]
+pub struct Proof<MC: MerkleChannel>
+where
+    MC::H: DeserializeOwned + Serialize,
+{
+    pub stark_proof: StarkProof<MC::H>,
+    pub constant_col_log_sizes: Vec<u32>,
+    pub witness_col_log_sizes: Vec<u32>,
+    pub machine_log_sizes: Vec<u32>,
+}