Simplify stwo proof (#2328)

Co-authored-by: ShuangWu121 <[email protected]>
powdr-labs · Jan 14, 2025 · 9f0285d · 9f0285d
1 parent 5a47b6e
commit 9f0285d
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 30 deletions.
diff --git a/backend/src/stwo/proof.rs b/backend/src/stwo/proof.rs
@@ -133,7 +133,5 @@ 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>,
+    pub machine_log_sizes: BTreeMap<String, u32>,
 }
diff --git a/backend/src/stwo/prover.rs b/backend/src/stwo/prover.rs
@@ -7,6 +7,7 @@ use powdr_number::FieldElement;
 use serde::de::DeserializeOwned;
 use serde::ser::Serialize;
 use std::collections::BTreeMap;
+use std::iter::repeat;
 use std::marker::PhantomData;
 use std::sync::Arc;
 use std::{fmt, io};
@@ -23,7 +24,7 @@ use stwo_prover::constraint_framework::{
 };
 
 use stwo_prover::core::air::{Component, ComponentProver};
-use stwo_prover::core::backend::{Backend, BackendForChannel, Column};
+use stwo_prover::core::backend::{Backend, BackendForChannel};
 use stwo_prover::core::channel::{Channel, MerkleChannel};
 use stwo_prover::core::fields::m31::{BaseField, M31};
 use stwo_prover::core::fields::qm31::SecureField;
@@ -245,7 +246,7 @@ where
 
         //The preprocessed columns needs to be indexed in the whole execution instead of each machine, so we need to keep track of the offset
         let mut constant_cols_offset_acc = 0;
-        let mut machine_log_sizes = Vec::new();
+        let mut machine_log_sizes = BTreeMap::new();
 
         let mut constant_cols = Vec::new();
 
@@ -293,7 +294,7 @@ where
                     );
                     components.push(component);
 
-                    machine_log_sizes.push(machine_length.ilog2());
+                    machine_log_sizes.insert(machine.clone(), machine_length.ilog2());
 
                     constant_cols_offset_acc +=
                         pil.constant_count() + get_constant_with_next_list(pil).len();
@@ -316,16 +317,6 @@ where
             .flatten()
             .collect();
 
-        let constant_col_log_sizes = constant_cols
-            .iter()
-            .map(|eval| eval.len().ilog2())
-            .collect::<Vec<_>>();
-
-        let witness_col_log_sizes = witness_by_machine
-            .iter()
-            .map(|eval| eval.len().ilog2())
-            .collect::<Vec<_>>();
-
         let twiddles_max_degree = B::precompute_twiddles(
             CanonicCoset::new(domain_degree_range.max.ilog2() + 1 + FRI_LOG_BLOWUP as u32)
                 .circle_domain()
@@ -366,8 +357,6 @@ where
 
         let proof: Proof<MC> = Proof {
             stark_proof,
-            constant_col_log_sizes,
-            witness_col_log_sizes,
             machine_log_sizes,
         };
         Ok(bincode::serialize(&proof).unwrap())
@@ -394,20 +383,35 @@ where
 
         let mut constant_cols_offset_acc = 0;
 
+        let mut constant_col_log_sizes = vec![];
+        let mut witness_col_log_sizes = vec![];
+
         let mut components = self
             .split
             .iter()
             .zip_eq(proof.machine_log_sizes.iter())
-            .map(|((_, pil), &machine_size)| {
-                constant_cols_offset_acc += pil.constant_count();
-
-                constant_cols_offset_acc += get_constant_with_next_list(pil).len();
-                PowdrComponent::new(
-                    tree_span_provider,
-                    PowdrEval::new((*pil).clone(), constant_cols_offset_acc, machine_size),
-                    (SecureField::zero(), None),
-                )
-            })
+            .map(
+                |((machine_name, pil), (proof_machine_name, &machine_log_size))| {
+                    assert_eq!(machine_name, proof_machine_name);
+                    let machine_component = PowdrComponent::new(
+                        tree_span_provider,
+                        PowdrEval::new((*pil).clone(), constant_cols_offset_acc, machine_log_size),
+                        (SecureField::zero(), None),
+                    );
+
+                    constant_cols_offset_acc += pil.constant_count();
+
+                    constant_cols_offset_acc += get_constant_with_next_list(pil).len();
+
+                    constant_col_log_sizes.extend(
+                        repeat(machine_log_size)
+                            .take(pil.constant_count() + get_constant_with_next_list(pil).len()),
+                    );
+                    witness_col_log_sizes
+                        .extend(repeat(machine_log_size).take(pil.commitment_count()));
+                    machine_component
+                },
+            )
             .collect::<Vec<_>>();
 
         let mut components_slice: Vec<&dyn Component> = components
@@ -419,12 +423,13 @@ where
 
         commitment_scheme.commit(
             proof.stark_proof.commitments[PREPROCESSED_TRACE_IDX],
-            &proof.constant_col_log_sizes,
+            &constant_col_log_sizes,
             verifier_channel,
         );
+
         commitment_scheme.commit(
             proof.stark_proof.commitments[ORIGINAL_TRACE_IDX],
-            &proof.witness_col_log_sizes,
+            &witness_col_log_sizes,
             verifier_channel,
         );