Adds 'AirBuilderWithPublicValues' trait and support for using public …

…values when proving/verifying
valida-xyz · May 15, 2024 · 317df5a · 317df5a
1 parent bdd338d
commit 317df5a
Show file tree

Hide file tree

Showing 11 changed files with 260 additions and 38 deletions.
diff --git a/air/src/air.rs b/air/src/air.rs
@@ -108,6 +108,10 @@ pub trait AirBuilder: Sized {
     }
 }
 
+pub trait AirBuilderWithPublicValues: AirBuilder {
+    fn public_values(&self) -> &[Self::F];
+}
+
 pub trait PairBuilder: AirBuilder {
     fn preprocessed(&self) -> Self::M;
 }

diff --git a/keccak-air/examples/prove_baby_bear_keccak.rs b/keccak-air/examples/prove_baby_bear_keccak.rs
@@ -69,8 +69,8 @@ fn main() -> Result<(), VerificationError> {
 
     let inputs = (0..NUM_HASHES).map(|_| random()).collect::<Vec<_>>();
     let trace = generate_trace_rows::<Val>(inputs);
-    let proof = prove::<MyConfig, _>(&config, &KeccakAir {}, &mut challenger, trace);
+    let proof = prove::<MyConfig, _>(&config, &KeccakAir {}, &mut challenger, trace, &vec![]);
 
     let mut challenger = Challenger::new(perm);
-    verify(&config, &KeccakAir {}, &mut challenger, &proof)
+    verify(&config, &KeccakAir {}, &mut challenger, &proof, &vec![])
 }
diff --git a/keccak-air/examples/prove_baby_bear_poseidon2.rs b/keccak-air/examples/prove_baby_bear_poseidon2.rs
@@ -69,8 +69,8 @@ fn main() -> Result<(), VerificationError> {
 
     let inputs = (0..NUM_HASHES).map(|_| random()).collect::<Vec<_>>();
     let trace = generate_trace_rows::<Val>(inputs);
-    let proof = prove::<MyConfig, _>(&config, &KeccakAir {}, &mut challenger, trace);
+    let proof = prove::<MyConfig, _>(&config, &KeccakAir {}, &mut challenger, trace, &vec![]);
 
     let mut challenger = Challenger::new(perm);
-    verify(&config, &KeccakAir {}, &mut challenger, &proof)
+    verify(&config, &KeccakAir {}, &mut challenger, &proof, &vec![])
 }
diff --git a/keccak-air/examples/prove_goldilocks_keccak.rs b/keccak-air/examples/prove_goldilocks_keccak.rs
@@ -68,8 +68,8 @@ fn main() -> Result<(), VerificationError> {
 
     let inputs = (0..NUM_HASHES).map(|_| random()).collect::<Vec<_>>();
     let trace = generate_trace_rows::<Val>(inputs);
-    let proof = prove::<MyConfig, _>(&config, &KeccakAir {}, &mut challenger, trace);
+    let proof = prove::<MyConfig, _>(&config, &KeccakAir {}, &mut challenger, trace, &vec![]);
 
     let mut challenger = Challenger::new(perm);
-    verify(&config, &KeccakAir {}, &mut challenger, &proof)
+    verify(&config, &KeccakAir {}, &mut challenger, &proof, &vec![])
 }
diff --git a/uni-stark/src/check_constraints.rs b/uni-stark/src/check_constraints.rs
@@ -1,11 +1,13 @@
-use p3_air::{Air, AirBuilder, TwoRowMatrixView};
+use alloc::vec::Vec;
+
+use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, TwoRowMatrixView};
 use p3_field::Field;
 use p3_matrix::dense::RowMajorMatrix;
 use p3_matrix::{Matrix, MatrixRowSlices};
 use tracing::instrument;
 
 #[instrument(name = "check constraints", skip_all)]
-pub(crate) fn check_constraints<F, A>(air: &A, main: &RowMajorMatrix<F>)
+pub(crate) fn check_constraints<F, A>(air: &A, main: &RowMajorMatrix<F>, public_values: &Vec<F>)
 where
     F: Field,
     A: for<'a> Air<DebugConstraintBuilder<'a, F>>,
@@ -25,6 +27,7 @@ where
         let mut builder = DebugConstraintBuilder {
             row_index: i,
             main,
+            public_values,
             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),
@@ -39,6 +42,7 @@ where
 pub struct DebugConstraintBuilder<'a, F: Field> {
     row_index: usize,
     main: TwoRowMatrixView<'a, F>,
+    public_values: &'a [F],
     is_first_row: F,
     is_last_row: F,
     is_transition: F,
@@ -92,3 +96,9 @@ where
         );
     }
 }
+
+impl<'a, F: Field> AirBuilderWithPublicValues for DebugConstraintBuilder<'a, F> {
+    fn public_values(&self) -> &[Self::F] {
+        self.public_values
+    }
+}
diff --git a/uni-stark/src/folder.rs b/uni-stark/src/folder.rs
@@ -1,24 +1,27 @@
-use p3_air::{AirBuilder, TwoRowMatrixView};
-use p3_field::{AbstractField, Field};
+use alloc::vec::Vec;
+use p3_air::{AirBuilder, AirBuilderWithPublicValues, TwoRowMatrixView};
+use p3_field::AbstractField;
 
 use crate::{PackedChallenge, PackedVal, StarkGenericConfig};
 
 pub struct ProverConstraintFolder<'a, SC: StarkGenericConfig> {
     pub main: TwoRowMatrixView<'a, PackedVal<SC>>,
+    pub public_values: &'a Vec<SC::Val>,
     pub is_first_row: PackedVal<SC>,
     pub is_last_row: PackedVal<SC>,
     pub is_transition: PackedVal<SC>,
     pub alpha: SC::Challenge,
     pub accumulator: PackedChallenge<SC>,
 }
 
-pub struct VerifierConstraintFolder<'a, Challenge> {
-    pub main: TwoRowMatrixView<'a, Challenge>,
-    pub is_first_row: Challenge,
-    pub is_last_row: Challenge,
-    pub is_transition: Challenge,
-    pub alpha: Challenge,
-    pub accumulator: Challenge,
+pub struct VerifierConstraintFolder<'a, SC: StarkGenericConfig> {
+    pub main: TwoRowMatrixView<'a, SC::Challenge>,
+    pub public_values: &'a Vec<SC::Val>,
+    pub is_first_row: SC::Challenge,
+    pub is_last_row: SC::Challenge,
+    pub is_transition: SC::Challenge,
+    pub alpha: SC::Challenge,
+    pub accumulator: SC::Challenge,
 }
 
 impl<'a, SC: StarkGenericConfig> AirBuilder for ProverConstraintFolder<'a, SC> {
@@ -54,11 +57,17 @@ impl<'a, SC: StarkGenericConfig> AirBuilder for ProverConstraintFolder<'a, SC> {
     }
 }
 
-impl<'a, Challenge: Field> AirBuilder for VerifierConstraintFolder<'a, Challenge> {
-    type F = Challenge;
-    type Expr = Challenge;
-    type Var = Challenge;
-    type M = TwoRowMatrixView<'a, Challenge>;
+impl<'a, SC: StarkGenericConfig> AirBuilderWithPublicValues for ProverConstraintFolder<'a, SC> {
+    fn public_values(&self) -> &[Self::F] {
+        self.public_values
+    }
+}
+
+impl<'a, SC: StarkGenericConfig> AirBuilder for VerifierConstraintFolder<'a, SC> {
+    type F = SC::Val;
+    type Expr = SC::Challenge;
+    type Var = SC::Challenge;
+    type M = TwoRowMatrixView<'a, SC::Challenge>;
 
     fn main(&self) -> Self::M {
         self.main
@@ -81,8 +90,13 @@ impl<'a, Challenge: Field> AirBuilder for VerifierConstraintFolder<'a, Challenge
     }
 
     fn assert_zero<I: Into<Self::Expr>>(&mut self, x: I) {
-        let x: Challenge = x.into();
+        let x: SC::Challenge = x.into();
         self.accumulator *= self.alpha;
         self.accumulator += x;
     }
 }
+impl<'a, SC: StarkGenericConfig> AirBuilderWithPublicValues for VerifierConstraintFolder<'a, SC> {
+    fn public_values(&self) -> &[Self::F] {
+        self.public_values
+    }
+}
diff --git a/uni-stark/src/prover.rs b/uni-stark/src/prover.rs
@@ -31,18 +31,19 @@ pub fn prove<
     air: &A,
     challenger: &mut SC::Challenger,
     trace: RowMajorMatrix<SC::Val>,
+    public_values: &Vec<SC::Val>,
 ) -> Proof<SC>
 where
     SC: StarkGenericConfig,
     A: Air<SymbolicAirBuilder<SC::Val>> + for<'a> Air<ProverConstraintFolder<'a, SC>>,
 {
     #[cfg(debug_assertions)]
-    crate::check_constraints::check_constraints(air, &trace);
+    crate::check_constraints::check_constraints(air, &trace, public_values);
 
     let degree = trace.height();
     let log_degree = log2_strict_usize(degree);
 
-    let log_quotient_degree = get_log_quotient_degree::<SC::Val, A>(air);
+    let log_quotient_degree = get_log_quotient_degree::<SC::Val, A>(air, public_values.len());
 
     let g_subgroup = SC::Val::two_adic_generator(log_degree);
 
@@ -63,6 +64,7 @@ where
     let quotient_values = quotient_values(
         config,
         air,
+        public_values,
         log_degree,
         log_quotient_degree,
         trace_lde_for_quotient,
@@ -118,6 +120,7 @@ where
 fn quotient_values<SC, A, Mat>(
     config: &SC,
     air: &A,
+    public_values: &Vec<SC::Val>,
     degree_bits: usize,
     quotient_degree_bits: usize,
     trace_lde: Mat,
@@ -192,6 +195,7 @@ where
                     local: &local,
                     next: &next,
                 },
+                public_values,
                 is_first_row,
                 is_last_row,
                 is_transition,

diff --git a/uni-stark/src/symbolic_builder.rs b/uni-stark/src/symbolic_builder.rs
@@ -2,7 +2,7 @@ use alloc::vec;
 use alloc::vec::Vec;
 use core::marker::PhantomData;
 
-use p3_air::{Air, AirBuilder};
+use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues};
 use p3_field::Field;
 use p3_matrix::dense::RowMajorMatrix;
 use p3_util::log2_ceil_usize;
@@ -12,13 +12,13 @@ use crate::symbolic_expression::SymbolicExpression;
 use crate::symbolic_variable::SymbolicVariable;
 
 #[instrument(name = "infer log of constraint degree", skip_all)]
-pub fn get_log_quotient_degree<F, A>(air: &A) -> usize
+pub fn get_log_quotient_degree<F, A>(air: &A, num_public_values: usize) -> usize
 where
     F: Field,
     A: Air<SymbolicAirBuilder<F>>,
 {
     // We pad to at least degree 2, since a quotient argument doesn't make sense with smaller degrees.
-    let constraint_degree = get_max_constraint_degree(air).max(2);
+    let constraint_degree = get_max_constraint_degree(air, num_public_values).max(2);
 
     // The quotient's actual degree is approximately (max_constraint_degree - 1) n,
     // where subtracting 1 comes from division by the zerofier.
@@ -27,37 +27,41 @@ where
 }
 
 #[instrument(name = "infer constraint degree", skip_all, level = "debug")]
-pub fn get_max_constraint_degree<F, A>(air: &A) -> usize
+pub fn get_max_constraint_degree<F, A>(air: &A, num_public_values: usize) -> usize
 where
     F: Field,
     A: Air<SymbolicAirBuilder<F>>,
 {
-    get_symbolic_constraints(air)
+    get_symbolic_constraints(air, num_public_values)
         .iter()
         .map(|c| c.degree_multiple())
         .max()
         .unwrap_or(0)
 }
 
-#[instrument(name = "evalute constraints symbolically", skip_all, level = "debug")]
-pub fn get_symbolic_constraints<F, A>(air: &A) -> Vec<SymbolicExpression<F>>
+#[instrument(name = "evaluate constraints symbolically", skip_all, level = "debug")]
+pub fn get_symbolic_constraints<F, A>(
+    air: &A,
+    num_public_values: usize,
+) -> Vec<SymbolicExpression<F>>
 where
     F: Field,
     A: Air<SymbolicAirBuilder<F>>,
 {
-    let mut builder = SymbolicAirBuilder::new(air.width());
+    let mut builder = SymbolicAirBuilder::new(air.width(), num_public_values);
     air.eval(&mut builder);
     builder.constraints()
 }
 
 /// An `AirBuilder` for evaluating constraints symbolically, and recording them for later use.
 pub struct SymbolicAirBuilder<F: Field> {
     main: RowMajorMatrix<SymbolicVariable<F>>,
+    public_values: Vec<F>,
     constraints: Vec<SymbolicExpression<F>>,
 }
 
 impl<F: Field> SymbolicAirBuilder<F> {
-    pub(crate) fn new(width: usize) -> Self {
+    pub(crate) fn new(width: usize, num_public_values: usize) -> Self {
         let values = [false, true]
             .into_iter()
             .flat_map(|is_next| {
@@ -70,6 +74,8 @@ impl<F: Field> SymbolicAirBuilder<F> {
             .collect();
         Self {
             main: RowMajorMatrix::new(values, width),
+            // TODO replace zeros once we have SymbolicExpression::PublicValue
+            public_values: vec![F::zero(); num_public_values],
             constraints: vec![],
         }
     }
@@ -109,3 +115,9 @@ impl<F: Field> AirBuilder for SymbolicAirBuilder<F> {
         self.constraints.push(x.into());
     }
 }
+
+impl<F: Field> AirBuilderWithPublicValues for SymbolicAirBuilder<F> {
+    fn public_values(&self) -> &[Self::F] {
+        self.public_values.as_slice()
+    }
+}
diff --git a/uni-stark/src/verifier.rs b/uni-stark/src/verifier.rs
@@ -18,12 +18,13 @@ pub fn verify<SC, A>(
     air: &A,
     challenger: &mut SC::Challenger,
     proof: &Proof<SC>,
+    public_values: &Vec<SC::Val>,
 ) -> Result<(), VerificationError>
 where
     SC: StarkGenericConfig,
-    A: Air<SymbolicAirBuilder<SC::Val>> + for<'a> Air<VerifierConstraintFolder<'a, SC::Challenge>>,
+    A: Air<SymbolicAirBuilder<SC::Val>> + for<'a> Air<VerifierConstraintFolder<'a, SC>>,
 {
-    let log_quotient_degree = get_log_quotient_degree::<SC::Val, A>(air);
+    let log_quotient_degree = get_log_quotient_degree::<SC::Val, A>(air, public_values.len());
     let quotient_degree = 1 << log_quotient_degree;
 
     let Proof {
@@ -110,6 +111,7 @@ where
             local: &opened_values.trace_local,
             next: &opened_values.trace_next,
         },
+        public_values,
         is_first_row,
         is_last_row,
         is_transition,