pi logic

Author: darth-cy

ceno_recursion/src/aggregation/mod.rs

@@ -47,6 +47,16 @@ use openvm_circuit::{
     },
     system::{memory::CHUNK, program::trace::compute_exe_commit},
 };
+use openvm_continuations::{
+    verifier::{
+        common::{
+            assert_or_assign_connector_pvs, assert_or_assign_memory_pvs,
+            assert_required_air_for_app_vm_present, get_connector_pvs, get_memory_pvs,
+            get_program_commit,
+        },
+        leaf::types::UserPublicValuesRootProof,
+    },
+};
 
 const LEAF_LOG_BLOWUP: usize = 1;
 const INTERNAL_LOG_BLOWUP: usize = 2;
@@ -63,18 +73,76 @@ impl CenoLeafVmVerifierConfig {
     pub fn build_program(&self) -> Program<F> {
         let mut builder = Builder::<C>::default();
 
+        let pvs = VmVerifierPvs::<Felt<F>>::uninit(&mut builder);
+
         {
             builder.cycle_tracker_start("VerifyCenoProof");
 
             let zkvm_proof_input_variables = ZKVMProofInput::read(&mut builder);
             verify_zkvm_proof(&mut builder, zkvm_proof_input_variables, &self.vk);
 
+            {
+                let commit = get_program_commit(builder, &proof);
+                builder.if_eq(i, RVar::zero()).then_or_else(
+                    |builder| {
+                        builder.assign(&pvs.app_commit, commit);
+                    },
+                    |builder| builder.assert_eq::<[_; DIGEST_SIZE]>(pvs.app_commit, commit),
+                );
+            }
+
+            let proof_connector_pvs = get_connector_pvs(builder, &proof);
+            assert_or_assign_connector_pvs(builder, &pvs.connector, i, &proof_connector_pvs);
+
+            let proof_memory_pvs = get_memory_pvs(builder, &proof);
+            assert_or_assign_memory_pvs(builder, &pvs.memory, i, &proof_memory_pvs);
+
+            let is_terminate = builder.cast_felt_to_var(pvs.connector.is_terminate);
+            builder.if_eq(is_terminate, F::ONE).then(|builder| {
+                let (pv_commit, expected_memory_root) =
+                    self.verify_user_public_values_root(builder);
+                builder.assert_eq::<[_; DIGEST_SIZE]>(pvs.memory.final_root, expected_memory_root);
+                builder.assign(&pvs.public_values_commit, pv_commit);
+            });
+            for pv in pvs.flatten() {
+                builder.commit_public_value(pv);
+            }
+
             builder.cycle_tracker_end("VerifyCenoProof");
             builder.halt();
         }
 
         builder.compile_isa_with_options(self.compiler_options)
     }
+
+    fn verify_user_public_values_root(
+        &self,
+        builder: &mut Builder<C>,
+    ) -> ([Felt<F>; DIGEST_SIZE], [Felt<F>; DIGEST_SIZE]) {
+        let memory_dimensions = self.app_system_config.memory_config.memory_dimensions();
+        let pv_as = PUBLIC_VALUES_ADDRESS_SPACE_OFFSET + memory_dimensions.as_offset;
+        let pv_start_idx = memory_dimensions.label_to_index((pv_as, 0));
+        let pv_height = log2_strict_usize(self.app_system_config.num_public_values / DIGEST_SIZE);
+        let proof_len = memory_dimensions.overall_height() - pv_height;
+        let idx_prefix = pv_start_idx >> pv_height;
+
+        // Read the public values root proof from the input stream.
+        let root_proof = UserPublicValuesRootProof::<F>::read(builder);
+        builder.assert_eq::<Usize<_>>(root_proof.sibling_hashes.len(), Usize::from(proof_len));
+        let mut curr_commit = root_proof.public_values_commit;
+        // Share the same state array to avoid unnecessary allocations.
+        let compressor = VariableP2Compressor::new(builder);
+        for i in 0..proof_len {
+            let sibling_hash = builder.get(&root_proof.sibling_hashes, i);
+            let (l_hash, r_hash) = if idx_prefix & (1 << i) != 0 {
+                (sibling_hash, curr_commit)
+            } else {
+                (curr_commit, sibling_hash)
+            };
+            curr_commit = compressor.compress(builder, &l_hash, &r_hash);
+        }
+        (root_proof.public_values_commit, curr_commit)
+    }
 }
 
 #[derive(Clone, Serialize, Deserialize)]
@@ -96,6 +164,10 @@ pub fn compress_to_root_proof(
         .enumerate()
         .map(|(shard_id, p)| ZKVMProofInput::from((shard_id, p)))
         .collect();
+    let user_public_values = zkvm_proof_inputs
+            .iter()
+            .flat_map(|p| p.raw_pi.iter().flat_map(|v| v.clone()).collect::<Vec<F>>())
+            .collect();
 
     let aggregation_start_timestamp = Instant::now();
 
@@ -273,10 +345,7 @@ pub fn compress_to_root_proof(
     // Export e2e stark proof (used in verify_e2e_stark_proof)
     let root_stark_proof = VmStarkProof {
         proof: proofs.pop().unwrap(),
-        user_public_values: zkvm_proof_inputs
-            .iter()
-            .flat_map(|p| p.raw_pi.iter().flat_map(|v| v.clone()).collect::<Vec<F>>())
-            .collect(),
+        user_public_values,
     };
     let file = File::create("root_stark_proof.bin").expect("Create export proof file");
     bincode::serialize_into(file, &root_stark_proof).expect("failed to serialize internal proof");