Uni-Stark Challenges

keccak-air/Cargo.toml

@@ -9,6 +9,7 @@ p3-air = { path = "../air" }
 p3-field = { path = "../field" }
 p3-matrix = { path = "../matrix" }
 p3-maybe-rayon = { path = "../maybe-rayon" }
+p3-uni-stark = { path = "../uni-stark" }
 p3-util = { path = "../util" }
 tracing = "0.1.37"
 
@@ -29,7 +30,6 @@ p3-poseidon = { path = "../poseidon" }
 p3-poseidon2 = { path = "../poseidon2" }
 p3-sha256 = { path = "../sha256", features = ["asm"] }
 p3-symmetric = { path = "../symmetric" }
-p3-uni-stark = { path = "../uni-stark" }
 rand = "0.8.5"
 tracing-subscriber = { version = "0.3.17", features = ["std", "env-filter"] }
 tracing-forest = { version = "0.1.6", features = ["ansi", "smallvec"] }

keccak-air/src/air.rs

@@ -3,6 +3,7 @@ use core::borrow::Borrow;
 use p3_air::{Air, AirBuilder, BaseAir};
 use p3_field::AbstractField;
 use p3_matrix::Matrix;
+use p3_uni_stark::MultiStageAir;
 
 use crate::columns::{KeccakCols, NUM_KECCAK_COLS};
 use crate::constants::rc_value_bit;
@@ -20,6 +21,8 @@ impl<F> BaseAir<F> for KeccakAir {
     }
 }
 
+impl<AB: AirBuilder> MultiStageAir<AB> for KeccakAir {}
+
 impl<AB: AirBuilder> Air<AB> for KeccakAir {
     #[inline]
     fn eval(&self, builder: &mut AB) {

uni-stark/src/check_constraints.rs

@@ -1,23 +1,28 @@
 use alloc::vec::Vec;
 
+use itertools::Itertools;
 use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, PairBuilder};
 use p3_field::Field;
 use p3_matrix::dense::{RowMajorMatrix, RowMajorMatrixView};
 use p3_matrix::stack::VerticalPair;
 use p3_matrix::Matrix;
 use tracing::instrument;
 
+use crate::traits::MultistageAirBuilder;
+
 #[instrument(name = "check constraints", skip_all)]
 pub(crate) fn check_constraints<F, A>(
     air: &A,
     preprocessed: &RowMajorMatrix<F>,
-    main: &RowMajorMatrix<F>,
-    public_values: &Vec<F>,
+    traces_by_stage: Vec<&RowMajorMatrix<F>>,
+    public_values_by_stage: &Vec<&Vec<F>>,
+    challenges: Vec<&Vec<F>>,
 ) where
     F: Field,
     A: for<'a> Air<DebugConstraintBuilder<'a, F>>,
 {
-    let height = main.height();
+    let num_stages = traces_by_stage.len();
+    let height = traces_by_stage[0].height();
 
     (0..height).for_each(|i| {
         let i_next = (i + 1) % height;
@@ -29,18 +34,30 @@ pub(crate) fn check_constraints<F, A>(
             RowMajorMatrixView::new_row(&*next_preprocessed),
         );
 
-        let local = main.row_slice(i);
-        let next = main.row_slice(i_next);
-        let main = VerticalPair::new(
-            RowMajorMatrixView::new_row(&*local),
-            RowMajorMatrixView::new_row(&*next),
-        );
+        let stages_local_next = traces_by_stage
+            .iter()
+            .map(|trace| {
+                let stage_local = trace.row_slice(i);
+                let stage_next = trace.row_slice(i_next);
+                (stage_local, stage_next)
+            })
+            .collect_vec();
+
+        let traces_by_stage = (0..num_stages)
+            .map(|stage| {
+                VerticalPair::new(
+                    RowMajorMatrixView::new_row(&*stages_local_next[stage].0),
+                    RowMajorMatrixView::new_row(&*stages_local_next[stage].1),
+                )
+            })
+            .collect();
 
         let mut builder = DebugConstraintBuilder {
             row_index: i,
+            challenges: challenges.clone(),
             preprocessed,
-            main,
-            public_values,
+            traces_by_stage,
+            public_values_by_stage,
             is_first_row: F::from_bool(i == 0),
             is_last_row: F::from_bool(i == height - 1),
             is_transition: F::from_bool(i != height - 1),
@@ -56,8 +73,9 @@ pub(crate) fn check_constraints<F, A>(
 pub struct DebugConstraintBuilder<'a, F: Field> {
     row_index: usize,
     preprocessed: VerticalPair<RowMajorMatrixView<'a, F>, RowMajorMatrixView<'a, F>>,
-    main: VerticalPair<RowMajorMatrixView<'a, F>, RowMajorMatrixView<'a, F>>,
-    public_values: &'a [F],
+    challenges: Vec<&'a Vec<F>>,
+    traces_by_stage: Vec<VerticalPair<RowMajorMatrixView<'a, F>, RowMajorMatrixView<'a, F>>>,
+    public_values_by_stage: &'a [&'a Vec<F>],
     is_first_row: F,
     is_last_row: F,
     is_transition: F,
@@ -89,7 +107,7 @@ where
     }
 
     fn main(&self) -> Self::M {
-        self.main
+        self.traces_by_stage[0]
     }
 
     fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
@@ -115,8 +133,8 @@ where
 impl<'a, F: Field> AirBuilderWithPublicValues for DebugConstraintBuilder<'a, F> {
     type PublicVar = Self::F;
 
-    fn public_values(&self) -> &[Self::F] {
-        self.public_values
+    fn public_values(&self) -> &[Self::PublicVar] {
+        self.stage_public_values(0)
     }
 }
 
@@ -125,3 +143,19 @@ impl<'a, F: Field> PairBuilder for DebugConstraintBuilder<'a, F> {
         self.preprocessed
     }
 }
+
+impl<'a, F: Field> MultistageAirBuilder for DebugConstraintBuilder<'a, F> {
+    type Challenge = Self::Expr;
+
+    fn stage_public_values(&self, stage: usize) -> &[Self::F] {
+        self.public_values_by_stage[stage]
+    }
+
+    fn stage_trace(&self, stage: usize) -> Self::M {
+        self.traces_by_stage[stage]
+    }
+
+    fn stage_challenges(&self, stage: usize) -> &[Self::Expr] {
+        self.challenges[stage]
+    }
+}

uni-stark/src/folder.rs

@@ -5,13 +5,15 @@ use p3_field::AbstractField;
 use p3_matrix::dense::{RowMajorMatrix, RowMajorMatrixView};
 use p3_matrix::stack::VerticalPair;
 
+use crate::traits::MultistageAirBuilder;
 use crate::{PackedChallenge, PackedVal, StarkGenericConfig, Val};
 
 #[derive(Debug)]
 pub struct ProverConstraintFolder<'a, SC: StarkGenericConfig> {
-    pub main: RowMajorMatrix<PackedVal<SC>>,
+    pub challenges: Vec<Vec<Val<SC>>>,
+    pub traces_by_stage: Vec<RowMajorMatrix<PackedVal<SC>>>,
     pub preprocessed: RowMajorMatrix<PackedVal<SC>>,
-    pub public_values: &'a Vec<Val<SC>>,
+    pub public_values_by_stage: &'a Vec<Vec<Val<SC>>>,
     pub is_first_row: PackedVal<SC>,
     pub is_last_row: PackedVal<SC>,
     pub is_transition: PackedVal<SC>,
@@ -23,9 +25,10 @@ type ViewPair<'a, T> = VerticalPair<RowMajorMatrixView<'a, T>, RowMajorMatrixVie
 
 #[derive(Debug)]
 pub struct VerifierConstraintFolder<'a, SC: StarkGenericConfig> {
-    pub main: ViewPair<'a, SC::Challenge>,
+    pub challenges: Vec<Vec<Val<SC>>>,
+    pub traces_by_stage: Vec<ViewPair<'a, SC::Challenge>>,
     pub preprocessed: ViewPair<'a, SC::Challenge>,
-    pub public_values: &'a Vec<Val<SC>>,
+    pub public_values_by_stage: Vec<&'a Vec<Val<SC>>>,
     pub is_first_row: SC::Challenge,
     pub is_last_row: SC::Challenge,
     pub is_transition: SC::Challenge,
@@ -40,7 +43,7 @@ impl<'a, SC: StarkGenericConfig> AirBuilder for ProverConstraintFolder<'a, SC> {
     type M = RowMajorMatrix<PackedVal<SC>>;
 
     fn main(&self) -> Self::M {
-        self.main.clone()
+        self.traces_by_stage[0].clone()
     }
 
     fn is_first_row(&self) -> Self::Expr {
@@ -67,10 +70,25 @@ impl<'a, SC: StarkGenericConfig> AirBuilder for ProverConstraintFolder<'a, SC> {
 }
 
 impl<'a, SC: StarkGenericConfig> AirBuilderWithPublicValues for ProverConstraintFolder<'a, SC> {
-    type PublicVar = Self::F;
+    type PublicVar = Val<SC>;
 
-    fn public_values(&self) -> &[Self::F] {
-        self.public_values
+    fn public_values(&self) -> &[Self::PublicVar] {
+        self.stage_public_values(0)
+    }
+}
+
+impl<'a, SC: StarkGenericConfig> MultistageAirBuilder for ProverConstraintFolder<'a, SC> {
+    type Challenge = Val<SC>;
+
+    fn stage_trace(&self, stage: usize) -> <Self as AirBuilder>::M {
+        self.traces_by_stage[stage].clone()
+    }
+
+    fn stage_challenges(&self, stage: usize) -> &[Self::Challenge] {
+        &self.challenges[stage]
+    }
+    fn stage_public_values(&self, stage: usize) -> &[Self::PublicVar] {
+        &self.public_values_by_stage[stage]
     }
 }
 
@@ -87,7 +105,7 @@ impl<'a, SC: StarkGenericConfig> AirBuilder for VerifierConstraintFolder<'a, SC>
     type M = ViewPair<'a, SC::Challenge>;
 
     fn main(&self) -> Self::M {
-        self.main
+        self.traces_by_stage[0]
     }
 
     fn is_first_row(&self) -> Self::Expr {
@@ -114,10 +132,25 @@ impl<'a, SC: StarkGenericConfig> AirBuilder for VerifierConstraintFolder<'a, SC>
 }
 
 impl<'a, SC: StarkGenericConfig> AirBuilderWithPublicValues for VerifierConstraintFolder<'a, SC> {
-    type PublicVar = Self::F;
+    type PublicVar = Val<SC>;
 
-    fn public_values(&self) -> &[Self::F] {
-        self.public_values
+    fn public_values(&self) -> &[Self::PublicVar] {
+        self.stage_public_values(0)
+    }
+}
+
+impl<'a, SC: StarkGenericConfig> MultistageAirBuilder for VerifierConstraintFolder<'a, SC> {
+    type Challenge = Val<SC>;
+
+    fn stage_trace(&self, stage: usize) -> <Self as AirBuilder>::M {
+        self.traces_by_stage[stage]
+    }
+
+    fn stage_challenges(&self, stage: usize) -> &[Self::Challenge] {
+        &self.challenges[stage]
+    }
+    fn stage_public_values(&self, stage: usize) -> &[Self::PublicVar] {
+        self.public_values_by_stage[stage]
     }
 }
 

uni-stark/src/lib.rs

@@ -11,6 +11,7 @@ mod prover;
 mod symbolic_builder;
 mod symbolic_expression;
 mod symbolic_variable;
+mod traits;
 mod verifier;
 mod zerofier_coset;
 
@@ -26,5 +27,6 @@ pub use prover::*;
 pub use symbolic_builder::*;
 pub use symbolic_expression::*;
 pub use symbolic_variable::*;
+pub use traits::*;
 pub use verifier::*;
 pub use zerofier_coset::*;

uni-stark/src/proof.rs

@@ -1,9 +1,10 @@
 use alloc::vec::Vec;
 
 use p3_commit::Pcs;
+use p3_matrix::dense::RowMajorMatrix;
 use serde::{Deserialize, Serialize};
 
-use crate::StarkGenericConfig;
+use crate::{StarkGenericConfig, Val};
 
 type Com<SC> = <<SC as StarkGenericConfig>::Pcs as Pcs<
     <SC as StarkGenericConfig>::Challenge,
@@ -29,16 +30,16 @@ pub struct Proof<SC: StarkGenericConfig> {
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct Commitments<Com> {
-    pub(crate) trace: Com,
+    pub(crate) traces_by_stage: Vec<Com>,
     pub(crate) quotient_chunks: Com,
 }
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct OpenedValues<Challenge> {
     pub(crate) preprocessed_local: Vec<Challenge>,
     pub(crate) preprocessed_next: Vec<Challenge>,
-    pub(crate) trace_local: Vec<Challenge>,
-    pub(crate) trace_next: Vec<Challenge>,
+    pub(crate) traces_by_stage_local: Vec<Vec<Challenge>>,
+    pub(crate) traces_by_stage_next: Vec<Vec<Challenge>>,
     pub(crate) quotient_chunks: Vec<Vec<Challenge>>,
 }
 
@@ -52,3 +53,12 @@ pub struct StarkProvingKey<SC: StarkGenericConfig> {
 pub struct StarkVerifyingKey<SC: StarkGenericConfig> {
     pub preprocessed_commit: Com<SC>,
 }
+
+pub struct ProcessedStage<SC: StarkGenericConfig> {
+    pub(crate) commitment: Com<SC>,
+    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>>,
+}