Add support for constant column in STWO backend

backend/Cargo.toml

@@ -62,7 +62,8 @@ p3-commit = { git = "https://github.com/plonky3/Plonky3.git", rev = "2192432ddf2
 p3-matrix = { git = "https://github.com/plonky3/Plonky3.git", rev = "2192432ddf28e7359dd2c577447886463e6124f0", optional = true }
 p3-uni-stark = { git = "https://github.com/plonky3/Plonky3.git", rev = "2192432ddf28e7359dd2c577447886463e6124f0", optional = true }
 # TODO: Change this to main branch when the `andrew/dev/update-toolchain` branch is merged,the main branch is using "nightly-2024-01-04", not compatiable with plonky3
-stwo-prover = { git = "https://github.com/starkware-libs/stwo.git", optional = true, rev = "e6d10bc107c11cce54bb4aa152c3afa2e15e92c1" }
+# stwo-prover = { git = "https://github.com/starkware-libs/stwo.git", optional = true, rev = "e6d10bc107c11cce54bb4aa152c3afa2e15e92c1" }
+stwo-prover = { git = "https://github.com/ShuangWu121/stwo.git", optional = true, rev = "564a4ddcde376ba0ae78da4d86ea5ad7338ef6fe" }
 
 strum = { version = "0.24.1", features = ["derive"] }
 log = "0.4.17"

backend/src/stwo/circuit_builder.rs

@@ -13,7 +13,10 @@ use std::sync::Arc;
 use powdr_ast::analyzed::{
     AlgebraicUnaryOperation, AlgebraicUnaryOperator, PolyID, PolynomialType,
 };
-use stwo_prover::constraint_framework::{EvalAtRow, FrameworkComponent, FrameworkEval};
+use stwo_prover::constraint_framework::preprocessed_columns::{gen_is_first, PreprocessedColumn};
+use stwo_prover::constraint_framework::{
+    EvalAtRow, FrameworkComponent, FrameworkEval, ORIGINAL_TRACE_IDX, PREPROCESSED_TRACE_IDX,
+};
 use stwo_prover::core::backend::ColumnOps;
 use stwo_prover::core::fields::m31::{BaseField, M31};
 use stwo_prover::core::fields::{ExtensionOf, FieldExpOps, FieldOps};
@@ -53,6 +56,7 @@ where
 pub struct PowdrEval<T> {
     analyzed: Arc<Analyzed<T>>,
     witness_columns: BTreeMap<PolyID, usize>,
+    constant_columns: BTreeMap<PolyID, usize>,
 }
 
 impl<T: FieldElement> PowdrEval<T> {
@@ -63,10 +67,18 @@ impl<T: FieldElement> PowdrEval<T> {
             .enumerate()
             .map(|(index, (_, id))| (id, index))
             .collect();
+
+        let constant_columns: BTreeMap<PolyID, usize> = analyzed
+            .definitions_in_source_order(PolynomialType::Constant)
+            .flat_map(|(symbol, _)| symbol.array_elements())
+            .enumerate()
+            .map(|(index, (_, id))| (id, index))
+            .collect();
 
         Self {
             analyzed,
             witness_columns,
+            constant_columns,
         }
     }
 }
@@ -80,14 +92,29 @@ impl<T: FieldElement> FrameworkEval for PowdrEval<T> {
     }
     fn evaluate<E: EvalAtRow>(&self, mut eval: E) -> E {
         assert!(
-            self.analyzed.constant_count() == 0 && self.analyzed.publics_count() == 0,
-            "Error: Expected no fixed columns nor public inputs, as they are not supported yet.",
+            self.analyzed.publics_count() == 0,
+            "Error: Expected no public inputs, as they are not supported yet.",
         );
 
         let witness_eval: BTreeMap<PolyID, [<E as EvalAtRow>::F; 2]> = self
             .witness_columns
             .keys()
-            .map(|poly_id| (*poly_id, eval.next_interaction_mask(0, [0, 1])))
+            .map(|poly_id| {
+                (
+                    *poly_id,
+                    eval.next_interaction_mask(ORIGINAL_TRACE_IDX, [0, 1]),
+                )
+            })
+            .collect();
+        let constant_eval: BTreeMap<PolyID, <E as EvalAtRow>::F> = self
+            .constant_columns
+            .keys()
+            .map(|poly_id| {
+                (
+                    *poly_id,
+                    eval.get_preprocessed_column(PreprocessedColumn::Plonk(3)),
+                )
+            })
             .collect();
 
         for id in self
@@ -96,7 +123,8 @@ impl<T: FieldElement> FrameworkEval for PowdrEval<T> {
         {
             match id {
                 Identity::Polynomial(identity) => {
-                    let expr = to_stwo_expression(&identity.expression, &witness_eval);
+                    let expr =
+                        to_stwo_expression(&identity.expression, &witness_eval, &constant_eval);
                     eval.add_constraint(expr);
                 }
                 Identity::Connect(..) => {
@@ -119,6 +147,7 @@ impl<T: FieldElement> FrameworkEval for PowdrEval<T> {
 fn to_stwo_expression<T: FieldElement, F>(
     expr: &AlgebraicExpression<T>,
     witness_eval: &BTreeMap<PolyID, [F; 2]>,
+    constant_eval: &BTreeMap<PolyID, F>,
 ) -> F
 where
     F: FieldExpOps
@@ -145,7 +174,7 @@ where
                     true => witness_eval[&poly_id][1].clone(),
                 },
                 PolynomialType::Constant => {
-                    unimplemented!("Constant polynomials are not supported in stwo yet")
+                    constant_eval[&poly_id].clone()
                 }
                 PolynomialType::Intermediate => {
                     unimplemented!("Intermediate polynomials are not supported in stwo yet")
@@ -162,15 +191,15 @@ where
             right,
         }) => match **right {
             AlgebraicExpression::Number(n) => {
-                let left = to_stwo_expression(left, witness_eval);
+                let left = to_stwo_expression(left, witness_eval, constant_eval);
                 (0u32..n.to_integer().try_into_u32().unwrap())
                     .fold(F::one(), |acc, _| acc * left.clone())
             }
             _ => unimplemented!("pow with non-constant exponent"),
         },
         AlgebraicExpression::BinaryOperation(AlgebraicBinaryOperation { left, op, right }) => {
-            let left = to_stwo_expression(left, witness_eval);
-            let right = to_stwo_expression(right, witness_eval);
+            let left = to_stwo_expression(left, witness_eval, constant_eval);
+            let right = to_stwo_expression(right, witness_eval, constant_eval);
 
             match op {
                 Add => left + right,
@@ -180,7 +209,7 @@ where
             }
         }
         AlgebraicExpression::UnaryOperation(AlgebraicUnaryOperation { op, expr }) => {
-            let expr = to_stwo_expression(expr, witness_eval);
+            let expr = to_stwo_expression(expr, witness_eval, constant_eval);
 
             match op {
                 AlgebraicUnaryOperator::Minus => -expr,

backend/src/stwo/mod.rs

@@ -18,6 +18,7 @@ use stwo_prover::core::vcs::blake2_merkle::Blake2sMerkleChannel;
 
 mod circuit_builder;
 mod prover;
+//mod proof;
 #[allow(dead_code)]
 
 struct RestrictedFactory;
@@ -42,9 +43,6 @@ impl<F: FieldElement> BackendFactory<F> for RestrictedFactory {
         if pil.degrees().len() > 1 {
             return Err(Error::NoVariableDegreeAvailable);
         }
-        let fixed = Arc::new(
-            get_uniquely_sized_cloned(&fixed).map_err(|_| Error::NoVariableDegreeAvailable)?,
-        );
         let stwo: Box<StwoProver<F, SimdBackend, Blake2sMerkleChannel, Blake2sChannel>> =
             Box::new(StwoProver::new(pil, fixed)?);
         Ok(stwo)

backend/src/stwo/prover.rs

@@ -1,4 +1,6 @@
 use powdr_ast::analyzed::Analyzed;
+use powdr_backend_utils::machine_fixed_columns;
+use powdr_executor::constant_evaluator::VariablySizedColumn;
 use serde::de::DeserializeOwned;
 use serde::ser::Serialize;
 use std::io;
@@ -14,10 +16,12 @@ use powdr_number::FieldElement;
 use stwo_prover::core::air::{Component, ComponentProver};
 use stwo_prover::core::backend::{Backend, BackendForChannel};
 use stwo_prover::core::channel::{Channel, MerkleChannel};
-use stwo_prover::core::fields::m31::M31;
+use stwo_prover::core::fields::m31::{BaseField, M31};
 use stwo_prover::core::fri::FriConfig;
 use stwo_prover::core::pcs::{CommitmentSchemeProver, CommitmentSchemeVerifier, PcsConfig};
-use stwo_prover::core::poly::circle::CanonicCoset;
+use stwo_prover::core::poly::circle::{CanonicCoset, CircleEvaluation};
+use stwo_prover::core::poly::BitReversedOrder;
+use stwo_prover::core::ColumnVec;
 
 const FRI_LOG_BLOWUP: usize = 1;
 const FRI_NUM_QUERIES: usize = 100;
@@ -26,7 +30,7 @@ const LOG_LAST_LAYER_DEGREE_BOUND: usize = 0;
 
 pub struct StwoProver<T, B: Backend + Send, MC: MerkleChannel, C: Channel> {
     pub analyzed: Arc<Analyzed<T>>,
-    _fixed: Arc<Vec<(String, Vec<T>)>>,
+    fixed: Arc<Vec<(String, VariablySizedColumn<T>)>>,
     /// Proving key placeholder
     _proving_key: Option<()>,
     /// Verifying key placeholder
@@ -46,11 +50,11 @@ where
 {
     pub fn new(
         analyzed: Arc<Analyzed<F>>,
-        _fixed: Arc<Vec<(String, Vec<F>)>>,
+        fixed: Arc<Vec<(String, VariablySizedColumn<F>)>>,
     ) -> Result<Self, io::Error> {
         Ok(Self {
             analyzed,
-            _fixed,
+            fixed,
             _proving_key: None,
             _verifying_key: None,
             _channel_marker: PhantomData,
@@ -70,26 +74,69 @@ where
         );
 
         // Setup protocol.
-        let mut prover_channel = <MC as MerkleChannel>::C::default();
+        let prover_channel = &mut <MC as MerkleChannel>::C::default();
         let commitment_scheme = &mut CommitmentSchemeProver::<B, MC>::new(config, &twiddles);
 
+        // get fix_columns evaluations
+        let fixed_columns = machine_fixed_columns(&self.fixed, &self.analyzed);
+
+        let domain = CanonicCoset::new(
+            fixed_columns
+                .keys()
+                .next()
+                .map(|&first_key| first_key.ilog2())
+                .unwrap_or(0),
+        )
+        .circle_domain();
+
+        let constant_trace: ColumnVec<CircleEvaluation<B, BaseField, BitReversedOrder>> =
+            fixed_columns
+                .values()
+                .flat_map(|vec| {
+                    println!("vec is {:?}", vec);
+                    vec.iter().map(|(_name, values)| {
+                        let values = values
+                            .iter()
+                            .map(|v| v.try_into_i32().unwrap().into())
+                            .collect();
+                        CircleEvaluation::new(domain, values)
+                    })
+                })
+                .collect();
+
+        // Preprocessed trace
+        let mut tree_builder = commitment_scheme.tree_builder();
+        tree_builder.extend_evals(constant_trace);
+        tree_builder.commit(prover_channel);
+
+        // committed/witness trace
         let trace = gen_stwo_circuit_trace::<F, B, M31>(witness);
 
         let mut tree_builder = commitment_scheme.tree_builder();
         tree_builder.extend_evals(trace);
-        tree_builder.commit(&mut prover_channel);
+        tree_builder.commit(prover_channel);
 
         let component = PowdrComponent::new(
             &mut TraceLocationAllocator::default(),
             PowdrEval::new(self.analyzed.clone()),
         );
 
+        let n_preprocessed_columns = commitment_scheme.trees[0]
+        .polynomials
+        .len();
+        println!("n_preprocessed_columns is {}", n_preprocessed_columns);
+
+        let trace = commitment_scheme.trace();
+        println!("trace is {:?}", trace.evals);
+
         let proof = stwo_prover::core::prover::prove::<B, MC>(
             &[&component],
-            &mut prover_channel,
+            prover_channel,
             commitment_scheme,
         )
         .unwrap();
+
+
 
         Ok(bincode::serialize(&proof).unwrap())
     }
@@ -105,8 +152,8 @@ where
         let proof: StarkProof<MC::H> =
             bincode::deserialize(proof).map_err(|e| format!("Failed to deserialize proof: {e}"))?;
 
-        let mut verifier_channel = <MC as MerkleChannel>::C::default();
-        let mut commitment_scheme = CommitmentSchemeVerifier::<MC>::new(config);
+        let verifier_channel = &mut <MC as MerkleChannel>::C::default();
+        let commitment_scheme = &mut CommitmentSchemeVerifier::<MC>::new(config);
 
         //Constraints that are to be proved
         let component = PowdrComponent::new(
@@ -115,21 +162,13 @@ where
         );
 
         // Retrieve the expected column sizes in each commitment interaction, from the AIR.
-        // TODO: When constant columns are supported, there will be more than one sizes and proof.commitments
-        // size[0] is for constant columns, size[1] is for witness columns, size[2] is for lookup columns
-        // pass size[1] for witness columns now is not doable due to this branch is outdated for the new feature of constant columns
-        // it will throw errors.
         let sizes = component.trace_log_degree_bounds();
-        assert_eq!(sizes.len(), 1);
-        commitment_scheme.commit(proof.commitments[0], &sizes[0], &mut verifier_channel);
 
-        stwo_prover::core::prover::verify(
-            &[&component],
-            &mut verifier_channel,
-            &mut commitment_scheme,
-            proof,
-        )
-        .map_err(|e| e.to_string())
+        commitment_scheme.commit(proof.commitments[0], &sizes[0], verifier_channel);
+        commitment_scheme.commit(proof.commitments[1], &sizes[1], verifier_channel);
+
+        stwo_prover::core::prover::verify(&[&component], verifier_channel, commitment_scheme, proof)
+            .map_err(|e| e.to_string())
     }
 }
 

test_data/pil/fibo_no_public.pil

@@ -0,0 +1,13 @@
+let N = 4;
+
+// This uses the alternative nomenclature as well.
+
+namespace Fibonacci(N);
+    col fixed ISLAST(i) { if i == N - 1 { 1 } else { 0 } };
+    col witness x, y;
+
+    ISLAST * (y' - 1) = 0;
+    ISLAST * (x' - 1) = 0;
+
+    (1-ISLAST) * (x' - y) = 0;
+    (1-ISLAST) * (y' - (x + y)) = 0;