Merge pull request #550 from mir-protocol/stark_circuit_constraint_test

Add test to check that native and circuit STARK constraints are consistent

evm/src/cpu/cpu_stark.rs

@@ -44,3 +44,39 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for CpuStark<F, D
         vec![PermutationPair::singletons(8, 9)]
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+
+    use crate::cpu::cpu_stark::CpuStark;
+    use crate::stark_testing::{test_stark_circuit_constraints, test_stark_low_degree};
+
+    #[test]
+    #[ignore] // TODO: remove this when constraints are no longer all 0.
+    fn test_stark_degree() -> Result<()> {
+        const D: usize = 2;
+        type C = PoseidonGoldilocksConfig;
+        type F = <C as GenericConfig<D>>::F;
+        type S = CpuStark<F, D>;
+
+        let stark = S {
+            f: Default::default(),
+        };
+        test_stark_low_degree(stark)
+    }
+
+    #[test]
+    fn test_stark_circuit() -> Result<()> {
+        const D: usize = 2;
+        type C = PoseidonGoldilocksConfig;
+        type F = <C as GenericConfig<D>>::F;
+        type S = CpuStark<F, D>;
+
+        let stark = S {
+            f: Default::default(),
+        };
+        test_stark_circuit_constraints::<F, C, S, D>(stark)
+    }
+}

evm/src/keccak/keccak_stark.rs

@@ -44,3 +44,39 @@ impl<F: RichField + Extendable<D>, const D: usize> Stark<F, D> for KeccakStark<F
         vec![PermutationPair::singletons(0, 6)]
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use anyhow::Result;
+    use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+
+    use crate::keccak::keccak_stark::KeccakStark;
+    use crate::stark_testing::{test_stark_circuit_constraints, test_stark_low_degree};
+
+    #[test]
+    #[ignore] // TODO: remove this when constraints are no longer all 0.
+    fn test_stark_degree() -> Result<()> {
+        const D: usize = 2;
+        type C = PoseidonGoldilocksConfig;
+        type F = <C as GenericConfig<D>>::F;
+        type S = KeccakStark<F, D>;
+
+        let stark = S {
+            f: Default::default(),
+        };
+        test_stark_low_degree(stark)
+    }
+
+    #[test]
+    fn test_stark_circuit() -> Result<()> {
+        const D: usize = 2;
+        type C = PoseidonGoldilocksConfig;
+        type F = <C as GenericConfig<D>>::F;
+        type S = KeccakStark<F, D>;
+
+        let stark = S {
+            f: Default::default(),
+        };
+        test_stark_circuit_constraints::<F, C, S, D>(stark)
+    }
+}

evm/src/stark_testing.rs

@@ -1,14 +1,21 @@
 use anyhow::{ensure, Result};
 use plonky2::field::extension_field::Extendable;
+use plonky2::field::extension_field::FieldExtension;
 use plonky2::field::field_types::Field;
 use plonky2::field::polynomial::{PolynomialCoeffs, PolynomialValues};
 use plonky2::hash::hash_types::RichField;
+use plonky2::iop::witness::PartialWitness;
+use plonky2::iop::witness::Witness;
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2::plonk::circuit_data::CircuitConfig;
+use plonky2::plonk::config::GenericConfig;
+use plonky2::plonk::config::Hasher;
 use plonky2::util::transpose;
 use plonky2_util::{log2_ceil, log2_strict};
 
-use crate::constraint_consumer::ConstraintConsumer;
+use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
 use crate::stark::Stark;
-use crate::vars::StarkEvaluationVars;
+use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars};
 
 const WITNESS_SIZE: usize = 1 << 5;
 
@@ -71,6 +78,85 @@ where
     Ok(())
 }
 
+/// Tests that the circuit constraints imposed by the given STARK are coherent with the native constraints.
+pub fn test_stark_circuit_constraints<
+    F: RichField + Extendable<D>,
+    C: GenericConfig<D, F = F>,
+    S: Stark<F, D>,
+    const D: usize,
+>(
+    stark: S,
+) -> Result<()>
+where
+    [(); S::COLUMNS]:,
+    [(); S::PUBLIC_INPUTS]:,
+    [(); C::Hasher::HASH_SIZE]:,
+{
+    // Compute native constraint evaluation on random values.
+    let vars = StarkEvaluationVars {
+        local_values: &F::Extension::rand_arr::<{ S::COLUMNS }>(),
+        next_values: &F::Extension::rand_arr::<{ S::COLUMNS }>(),
+        public_inputs: &F::Extension::rand_arr::<{ S::PUBLIC_INPUTS }>(),
+    };
+    let alphas = F::rand_vec(1);
+    let z_last = F::Extension::rand();
+    let lagrange_first = F::Extension::rand();
+    let lagrange_last = F::Extension::rand();
+    let mut consumer = ConstraintConsumer::<F::Extension>::new(
+        alphas
+            .iter()
+            .copied()
+            .map(F::Extension::from_basefield)
+            .collect(),
+        z_last,
+        lagrange_first,
+        lagrange_last,
+    );
+    stark.eval_ext(vars, &mut consumer);
+    let native_eval = consumer.accumulators()[0];
+
+    // Compute circuit constraint evaluation on same random values.
+    let circuit_config = CircuitConfig::standard_recursion_config();
+    let mut builder = CircuitBuilder::<F, D>::new(circuit_config);
+    let mut pw = PartialWitness::<F>::new();
+
+    let locals_t = builder.add_virtual_extension_targets(S::COLUMNS);
+    pw.set_extension_targets(&locals_t, vars.local_values);
+    let nexts_t = builder.add_virtual_extension_targets(S::COLUMNS);
+    pw.set_extension_targets(&nexts_t, vars.next_values);
+    let pis_t = builder.add_virtual_extension_targets(S::PUBLIC_INPUTS);
+    pw.set_extension_targets(&pis_t, vars.public_inputs);
+    let alphas_t = builder.add_virtual_targets(1);
+    pw.set_target(alphas_t[0], alphas[0]);
+    let z_last_t = builder.add_virtual_extension_target();
+    pw.set_extension_target(z_last_t, z_last);
+    let lagrange_first_t = builder.add_virtual_extension_target();
+    pw.set_extension_target(lagrange_first_t, lagrange_first);
+    let lagrange_last_t = builder.add_virtual_extension_target();
+    pw.set_extension_target(lagrange_last_t, lagrange_last);
+
+    let vars = StarkEvaluationTargets::<D, { S::COLUMNS }, { S::PUBLIC_INPUTS }> {
+        local_values: &locals_t.try_into().unwrap(),
+        next_values: &nexts_t.try_into().unwrap(),
+        public_inputs: &pis_t.try_into().unwrap(),
+    };
+    let mut consumer = RecursiveConstraintConsumer::<F, D>::new(
+        builder.zero_extension(),
+        alphas_t,
+        z_last_t,
+        lagrange_first_t,
+        lagrange_last_t,
+    );
+    stark.eval_ext_circuit(&mut builder, vars, &mut consumer);
+    let circuit_eval = consumer.accumulators()[0];
+    let native_eval_t = builder.constant_extension(native_eval);
+    builder.connect_extension(circuit_eval, native_eval_t);
+
+    let data = builder.build::<C>();
+    let proof = data.prove(pw)?;
+    data.verify(proof)
+}
+
 fn random_low_degree_matrix<F: Field>(num_polys: usize, rate_bits: usize) -> Vec<Vec<F>> {
     let polys = (0..num_polys)
         .map(|_| random_low_degree_values(rate_bits))

starky/src/fibonacci_stark.rs

@@ -142,7 +142,7 @@ mod tests {
         verify_stark_proof_circuit,
     };
     use crate::stark::Stark;
-    use crate::stark_testing::test_stark_low_degree;
+    use crate::stark_testing::{test_stark_circuit_constraints, test_stark_low_degree};
     use crate::verifier::verify_stark_proof;
 
     fn fibonacci<F: Field>(n: usize, x0: F, x1: F) -> F {
@@ -184,6 +184,18 @@ mod tests {
         test_stark_low_degree(stark)
     }
 
+    #[test]
+    fn test_fibonacci_stark_circuit() -> Result<()> {
+        const D: usize = 2;
+        type C = PoseidonGoldilocksConfig;
+        type F = <C as GenericConfig<D>>::F;
+        type S = FibonacciStark<F, D>;
+
+        let num_rows = 1 << 5;
+        let stark = S::new(num_rows);
+        test_stark_circuit_constraints::<F, C, S, D>(stark)
+    }
+
     #[test]
     fn test_recursive_stark_verifier() -> Result<()> {
         init_logger();

starky/src/stark_testing.rs

@@ -1,14 +1,19 @@
 use anyhow::{ensure, Result};
-use plonky2::field::extension_field::Extendable;
+use plonky2::field::extension_field::{Extendable, FieldExtension};
 use plonky2::field::field_types::Field;
 use plonky2::field::polynomial::{PolynomialCoeffs, PolynomialValues};
 use plonky2::hash::hash_types::RichField;
+use plonky2::iop::witness::{PartialWitness, Witness};
+use plonky2::plonk::circuit_builder::CircuitBuilder;
+use plonky2::plonk::circuit_data::CircuitConfig;
+use plonky2::plonk::config::GenericConfig;
+use plonky2::plonk::config::Hasher;
 use plonky2::util::transpose;
 use plonky2_util::{log2_ceil, log2_strict};
 
-use crate::constraint_consumer::ConstraintConsumer;
+use crate::constraint_consumer::{ConstraintConsumer, RecursiveConstraintConsumer};
 use crate::stark::Stark;
-use crate::vars::StarkEvaluationVars;
+use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars};
 
 const WITNESS_SIZE: usize = 1 << 5;
 
@@ -71,6 +76,85 @@ where
     Ok(())
 }
 
+/// Tests that the circuit constraints imposed by the given STARK are coherent with the native constraints.
+pub fn test_stark_circuit_constraints<
+    F: RichField + Extendable<D>,
+    C: GenericConfig<D, F = F>,
+    S: Stark<F, D>,
+    const D: usize,
+>(
+    stark: S,
+) -> Result<()>
+where
+    [(); S::COLUMNS]:,
+    [(); S::PUBLIC_INPUTS]:,
+    [(); C::Hasher::HASH_SIZE]:,
+{
+    // Compute native constraint evaluation on random values.
+    let vars = StarkEvaluationVars {
+        local_values: &F::Extension::rand_arr::<{ S::COLUMNS }>(),
+        next_values: &F::Extension::rand_arr::<{ S::COLUMNS }>(),
+        public_inputs: &F::Extension::rand_arr::<{ S::PUBLIC_INPUTS }>(),
+    };
+    let alphas = F::rand_vec(1);
+    let z_last = F::Extension::rand();
+    let lagrange_first = F::Extension::rand();
+    let lagrange_last = F::Extension::rand();
+    let mut consumer = ConstraintConsumer::<F::Extension>::new(
+        alphas
+            .iter()
+            .copied()
+            .map(F::Extension::from_basefield)
+            .collect(),
+        z_last,
+        lagrange_first,
+        lagrange_last,
+    );
+    stark.eval_ext(vars, &mut consumer);
+    let native_eval = consumer.accumulators()[0];
+
+    // Compute circuit constraint evaluation on same random values.
+    let circuit_config = CircuitConfig::standard_recursion_config();
+    let mut builder = CircuitBuilder::<F, D>::new(circuit_config);
+    let mut pw = PartialWitness::<F>::new();
+
+    let locals_t = builder.add_virtual_extension_targets(S::COLUMNS);
+    pw.set_extension_targets(&locals_t, vars.local_values);
+    let nexts_t = builder.add_virtual_extension_targets(S::COLUMNS);
+    pw.set_extension_targets(&nexts_t, vars.next_values);
+    let pis_t = builder.add_virtual_extension_targets(S::PUBLIC_INPUTS);
+    pw.set_extension_targets(&pis_t, vars.public_inputs);
+    let alphas_t = builder.add_virtual_targets(1);
+    pw.set_target(alphas_t[0], alphas[0]);
+    let z_last_t = builder.add_virtual_extension_target();
+    pw.set_extension_target(z_last_t, z_last);
+    let lagrange_first_t = builder.add_virtual_extension_target();
+    pw.set_extension_target(lagrange_first_t, lagrange_first);
+    let lagrange_last_t = builder.add_virtual_extension_target();
+    pw.set_extension_target(lagrange_last_t, lagrange_last);
+
+    let vars = StarkEvaluationTargets::<D, { S::COLUMNS }, { S::PUBLIC_INPUTS }> {
+        local_values: &locals_t.try_into().unwrap(),
+        next_values: &nexts_t.try_into().unwrap(),
+        public_inputs: &pis_t.try_into().unwrap(),
+    };
+    let mut consumer = RecursiveConstraintConsumer::<F, D>::new(
+        builder.zero_extension(),
+        alphas_t,
+        z_last_t,
+        lagrange_first_t,
+        lagrange_last_t,
+    );
+    stark.eval_ext_circuit(&mut builder, vars, &mut consumer);
+    let circuit_eval = consumer.accumulators()[0];
+    let native_eval_t = builder.constant_extension(native_eval);
+    builder.connect_extension(circuit_eval, native_eval_t);
+
+    let data = builder.build::<C>();
+    let proof = data.prove(pw)?;
+    data.verify(proof)
+}
+
 fn random_low_degree_matrix<F: Field>(num_polys: usize, rate_bits: usize) -> Vec<Vec<F>> {
     let polys = (0..num_polys)
         .map(|_| random_low_degree_values(rate_bits))