fix: public inputs in recursion program (#467)

core/src/air/public_values.rs

@@ -14,7 +14,7 @@ use super::Word;
 
 // TODO:  Create a config struct that will store the num_words setting and the hash function
 //        and initial entropy used.
-const PV_DIGEST_NUM_WORDS: usize = 8;
+pub const PV_DIGEST_NUM_WORDS: usize = 8;
 
 /// The PublicValuesDigest struct is used to represent the public values digest.  This is the hash of all the
 /// bytes that the guest program has written to public values.
@@ -81,6 +81,12 @@ impl<T: Debug + Copy> From<PublicValuesDigest<Word<T>>> for Vec<T> {
     }
 }
 
+impl<T> From<[T; PV_DIGEST_NUM_WORDS]> for PublicValuesDigest<T> {
+    fn from(words: [T; PV_DIGEST_NUM_WORDS]) -> Self {
+        PublicValuesDigest(words)
+    }
+}
+
 /// Implement the IndexMut trait for PublicValuesDigest to index specific words.
 impl<T> IndexMut<usize> for PublicValuesDigest<T> {
     fn index_mut(&mut self, index: usize) -> &mut Self::Output {

recursion/program/src/challenger.rs

@@ -7,6 +7,15 @@ use crate::types::Commitment;
 
 pub trait CanObserveVariable<C: Config, V> {
     fn observe(&mut self, builder: &mut Builder<C>, value: V);
+
+    fn observe_slice(&mut self, builder: &mut Builder<C>, values: &[V])
+    where
+        V: Copy,
+    {
+        for value in values {
+            self.observe(builder, *value);
+        }
+    }
 }
 
 pub trait CanSampleVariable<C: Config, V> {

recursion/program/src/constraints.rs

@@ -8,8 +8,11 @@ use p3_field::AbstractExtensionField;
 use p3_field::AbstractField;
 use p3_field::TwoAdicField;
 use sp1_core::air::MachineAir;
+use sp1_core::air::PublicValuesDigest;
+use sp1_core::air::Word;
 use sp1_core::stark::AirOpenedValues;
 use sp1_core::stark::{MachineChip, StarkGenericConfig};
+use sp1_recursion_compiler::ir::Felt;
 
 use crate::commit::PolynomialSpaceVariable;
 
@@ -30,6 +33,7 @@ where
         builder: &mut Builder<C>,
         chip: &MachineChip<SC, A>,
         opening: &ChipOpening<C>,
+        public_values_digest: PublicValuesDigest<Word<Felt<C::F>>>,
         selectors: &LagrangeSelectors<Ext<C::F, C::EF>>,
         alpha: Ext<C::F, C::EF>,
         permutation_challenges: &[C::EF],
@@ -56,13 +60,15 @@ where
         };
 
         let zero: Ext<SC::Val, SC::Challenge> = builder.eval(SC::Val::zero());
+        let public_values: Vec<Felt<C::F>> = public_values_digest.into();
         let mut folder = RecursiveVerifierConstraintFolder {
             builder,
             preprocessed: opening.preprocessed.view(),
             main: opening.main.view(),
             perm: perm_opening.view(),
             perm_challenges: permutation_challenges,
             cumulative_sum: opening.cumulative_sum,
+            public_values: &public_values,
             is_first_row: selectors.is_first_row,
             is_last_row: selectors.is_last_row,
             is_transition: selectors.is_transition,
@@ -124,6 +130,7 @@ where
         builder: &mut Builder<C>,
         chip: &MachineChip<SC, A>,
         opening: &ChipOpenedValuesVariable<C>,
+        public_values_digest: PublicValuesDigest<Word<Felt<C::F>>>,
         trace_domain: TwoAdicMultiplicativeCosetVariable<C>,
         qc_domains: Vec<TwoAdicMultiplicativeCosetVariable<C>>,
         zeta: Ext<C::F, C::EF>,
@@ -139,6 +146,7 @@ where
             builder,
             chip,
             &opening,
+            public_values_digest,
             &sels,
             alpha,
             permutation_challenges,
@@ -156,6 +164,7 @@ mod tests {
     use itertools::{izip, Itertools};
     use serde::{de::DeserializeOwned, Serialize};
     use sp1_core::{
+        air::{PublicValuesDigest, Word},
         runtime::Program,
         stark::{
             Chip, Com, Dom, MachineStark, OpeningProof, PcsProverData, RiscvAir, ShardCommitment,
@@ -272,7 +281,6 @@ mod tests {
     }
 
     #[test]
-    #[ignore]
     fn test_verify_constraints_parts() {
         type SC = BabyBearPoseidon2;
         type F = <SC as StarkGenericConfig>::Val;
@@ -287,21 +295,25 @@ mod tests {
         let machine = A::machine(SC::default());
         let (_, vk) = machine.setup(&Program::from(elf));
         let mut challenger = machine.config().challenger();
-        let proofs = SP1Prover::prove_with_config(elf, SP1Stdin::new(), machine.config().clone())
+        let proof = SP1Prover::prove_with_config(elf, SP1Stdin::new(), machine.config().clone())
             .unwrap()
-            .proof
-            .shard_proofs;
+            .proof;
         println!("Proof generated successfully");
 
         challenger.observe(vk.commit);
-        proofs.iter().for_each(|proof| {
+        proof.shard_proofs.iter().for_each(|proof| {
             challenger.observe(proof.commitment.main_commit);
         });
 
+        // Observe the public input digest
+        let pv_digest_field_elms: Vec<F> =
+            PublicValuesDigest::<Word<F>>::new(proof.public_values_digest).into();
+        challenger.observe_slice(&pv_digest_field_elms);
+
         // Run the verify inside the DSL and compare it to the calculated value.
         let mut builder = VmBuilder::<F, EF>::default();
 
-        for proof in proofs.into_iter().take(1) {
+        for proof in proof.shard_proofs.into_iter().take(1) {
             let (
                 chips,
                 trace_domains_vals,
@@ -311,6 +323,12 @@ mod tests {
                 zeta_val,
             ) = get_shard_data(&machine, &proof, &mut challenger);
 
+            // Set up the public values digest.
+            let public_values_digest = PublicValuesDigest::from(core::array::from_fn(|i| {
+                let word_val = proof.public_values_digest[i];
+                Word(core::array::from_fn(|j| builder.eval(word_val[j])))
+            }));
+
             for (chip, trace_domain_val, qc_domains_vals, values_vals) in izip!(
                 chips.iter(),
                 trace_domains_vals,
@@ -341,6 +359,7 @@ mod tests {
                     &mut builder,
                     chip,
                     &values,
+                    public_values_digest,
                     &sels,
                     alpha,
                     permutation_challenges.as_slice(),
@@ -385,7 +404,6 @@ mod tests {
     }
 
     #[test]
-    #[ignore]
     fn test_verify_constraints_whole() {
         type SC = BabyBearPoseidon2;
         type F = <SC as StarkGenericConfig>::Val;
@@ -404,19 +422,24 @@ mod tests {
             SP1Prover::prove_with_config(elf, SP1Stdin::new(), machine.config().clone()).unwrap();
         SP1Verifier::verify_with_config(elf, &proof, machine.config().clone()).unwrap();
 
-        let proofs = proof.proof.shard_proofs;
+        let proof = proof.proof;
         println!("Proof generated and verified successfully");
 
         challenger.observe(vk.commit);
 
-        proofs.iter().for_each(|proof| {
+        proof.shard_proofs.iter().for_each(|proof| {
             challenger.observe(proof.commitment.main_commit);
         });
 
+        // Observe the public input digest
+        let pv_digest_field_elms: Vec<F> =
+            PublicValuesDigest::<Word<F>>::new(proof.public_values_digest).into();
+        challenger.observe_slice(&pv_digest_field_elms);
+
         // Run the verify inside the DSL and compare it to the calculated value.
         let mut builder = VmBuilder::<F, EF>::default();
 
-        for proof in proofs.into_iter().take(1) {
+        for proof in proof.shard_proofs.into_iter().take(1) {
             let (
                 chips,
                 trace_domains_vals,
@@ -441,10 +464,17 @@ mod tests {
                     .map(|domain| builder.eval_const(*domain))
                     .collect::<Vec<_>>();
 
+                // Set up the public values digest.
+                let public_values_digest = PublicValuesDigest::from(core::array::from_fn(|i| {
+                    let word_val = proof.public_values_digest[i];
+                    Word(core::array::from_fn(|j| builder.eval(word_val[j])))
+                }));
+
                 StarkVerifier::<_, SC>::verify_constraints::<A>(
                     &mut builder,
                     chip,
                     &opening,
+                    public_values_digest,
                     trace_domain,
                     qc_domains,
                     zeta,

recursion/program/src/folder.rs

@@ -5,7 +5,7 @@ use p3_air::{
 use sp1_core::air::{EmptyMessageBuilder, MultiTableAirBuilder, PublicValuesBuilder};
 
 use sp1_recursion_compiler::{
-    ir::{Builder, Config, Ext},
+    ir::{Builder, Config, Ext, Felt},
     prelude::SymbolicExt,
 };
 
@@ -15,6 +15,7 @@ pub struct RecursiveVerifierConstraintFolder<'a, C: Config> {
     pub main: TwoRowMatrixView<'a, Ext<C::F, C::EF>>,
     pub perm: TwoRowMatrixView<'a, Ext<C::F, C::EF>>,
     pub perm_challenges: &'a [C::EF],
+    pub public_values: &'a [Felt<C::F>],
     pub cumulative_sum: Ext<C::F, C::EF>,
     pub is_first_row: Ext<C::F, C::EF>,
     pub is_last_row: Ext<C::F, C::EF>,
@@ -101,9 +102,9 @@ impl<'a, C: Config> EmptyMessageBuilder for RecursiveVerifierConstraintFolder<'a
 impl<'a, C: Config> PublicValuesBuilder for RecursiveVerifierConstraintFolder<'a, C> {}
 
 impl<'a, C: Config> AirBuilderWithPublicValues for RecursiveVerifierConstraintFolder<'a, C> {
-    type PublicVar = C::F;
+    type PublicVar = Felt<C::F>;
 
     fn public_values(&self) -> &[Self::PublicVar] {
-        &[]
+        self.public_values
     }
 }