check constraints, fix challenge values

air/src/air.rs

@@ -189,11 +189,13 @@ pub trait PermutationAirBuilder: ExtensionBuilder {
 }
 
 pub trait MultistageAirBuilder: AirBuilder {
+    type Challenge: Clone + Into<Self::Expr>;
+
     /// Traces from each stage.
     fn multi_stage(&self, stage: usize) -> Self::M;
 
     /// Challenges from each stage, drawn from the base field
-    fn challenges(&self, stage: usize) -> &[Self::Expr];
+    fn challenges(&self, stage: usize) -> &[Self::Challenge];
 }
 #[derive(Debug)]
 pub struct FilteredAirBuilder<'a, AB: AirBuilder> {

uni-stark/src/check_constraints.rs

@@ -143,6 +143,8 @@ impl<'a, F: Field> PairBuilder for DebugConstraintBuilder<'a, F> {
 }
 
 impl<'a, F: Field> MultistageAirBuilder for DebugConstraintBuilder<'a, F> {
+    type Challenge = Self::Expr;
+
     fn multi_stage(&self, stage: usize) -> Self::M {
         self.stages[stage]
     }

uni-stark/src/folder.rs

@@ -9,7 +9,7 @@ use crate::{PackedChallenge, PackedVal, StarkGenericConfig, Val};
 
 #[derive(Debug)]
 pub struct ProverConstraintFolder<'a, SC: StarkGenericConfig> {
-    pub challenges: Vec<Vec<PackedVal<SC>>>,
+    pub challenges: Vec<Vec<Val<SC>>>,
     pub stages: Vec<RowMajorMatrix<PackedVal<SC>>>,
     pub preprocessed: RowMajorMatrix<PackedVal<SC>>,
     pub public_values: &'a Vec<Vec<Val<SC>>>,
@@ -24,7 +24,7 @@ type ViewPair<'a, T> = VerticalPair<RowMajorMatrixView<'a, T>, RowMajorMatrixVie
 
 #[derive(Debug)]
 pub struct VerifierConstraintFolder<'a, SC: StarkGenericConfig> {
-    pub challenges: Vec<Vec<SC::Challenge>>,
+    pub challenges: Vec<Vec<Val<SC>>>,
     pub stages: Vec<ViewPair<'a, SC::Challenge>>,
     pub preprocessed: ViewPair<'a, SC::Challenge>,
     pub public_values: Vec<&'a Vec<Val<SC>>>,
@@ -83,11 +83,13 @@ impl<'a, SC: StarkGenericConfig> PairBuilder for ProverConstraintFolder<'a, SC>
 }
 
 impl<'a, SC: StarkGenericConfig> MultistageAirBuilder for ProverConstraintFolder<'a, SC> {
+    type Challenge = Val<SC>;
+
     fn multi_stage(&self, stage: usize) -> Self::M {
         self.stages[stage].clone()
     }
 
-    fn challenges(&self, stage: usize) -> &[Self::Expr] {
+    fn challenges(&self, stage: usize) -> &[Self::Challenge] {
         &self.challenges[stage]
     }
 }
@@ -140,11 +142,13 @@ impl<'a, SC: StarkGenericConfig> PairBuilder for VerifierConstraintFolder<'a, SC
 }
 
 impl<'a, SC: StarkGenericConfig> MultistageAirBuilder for VerifierConstraintFolder<'a, SC> {
+    type Challenge = Val<SC>;
+
     fn multi_stage(&self, stage: usize) -> Self::M {
         self.stages[stage]
     }
 
-    fn challenges(&self, stage: usize) -> &[Self::Expr] {
+    fn challenges(&self, stage: usize) -> &[Self::Challenge] {
         &self.challenges[stage]
     }
 }

uni-stark/src/proof.rs

@@ -65,6 +65,8 @@ pub struct ProcessedStage<SC: StarkGenericConfig> {
     pub(crate) prover_data: PcsProverData<SC>,
     pub(crate) challenge_values: Vec<Val<SC>>,
     pub(crate) public_values: Vec<Val<SC>>,
+    #[cfg(debug_assertions)]
+    pub(crate) trace: RowMajorMatrix<Val<SC>>,
 }
 
 /// Updating with each new trace in every stage
@@ -224,6 +226,10 @@ impl<'a, SC: StarkGenericConfig> State<'a, SC> {
     }
 
     pub(crate) fn run_stage(mut self, stage: Stage<SC>) -> Self {
+
+        #[cfg(debug_assertions)]
+        let trace = stage.trace.clone();
+
         // commit to the trace for this stage
         let (commitment, prover_data) = info_span!("commit to trace data")
             .in_scope(|| self.pcs.commit(vec![(self.trace_domain, stage.trace)]));
@@ -242,6 +248,8 @@ impl<'a, SC: StarkGenericConfig> State<'a, SC> {
             prover_data,
             commitment,
             challenge_values,
+            #[cfg(debug_assertions)]
+            trace,
         });
         self
     }

uni-stark/src/prover.rs

@@ -129,6 +129,16 @@ where
     // sanity check that we processed as many stages as expected
     assert_eq!(state.processed_stages.len(), stage_count);
 
+    #[cfg(debug_assertions)]
+    crate::check_constraints::check_constraints(
+        air,
+        &air.preprocessed_trace()
+            .unwrap_or(RowMajorMatrix::new(Default::default(), 0)),
+        state.processed_stages.iter().map(|s| &s.trace).collect(),
+        &state.processed_stages.iter().map(|s| &s.public_values).collect(),
+        state.processed_stages.iter().map(|s| &s.challenge_values).collect(),
+    );
+
     finish(proving_key, air, state)
 }
 
@@ -306,6 +316,7 @@ where
                 .collect::<Vec<RowMajorMatrix<PackedVal<SC>>>>();
 
             let accumulator = PackedChallenge::<SC>::zero();
+
             let mut folder = ProverConstraintFolder {
                 challenges: challenges.clone(),
                 stages,

uni-stark/src/symbolic_builder.rs

@@ -186,6 +186,8 @@ impl<F: Field> PairBuilder for SymbolicAirBuilder<F> {
 }
 
 impl<F: Field> MultistageAirBuilder for SymbolicAirBuilder<F> {
+    type Challenge = Self::Expr;
+
     fn multi_stage(&self, stage: usize) -> Self::M {
         self.stages[stage].clone()
     }