Support single-shard e2e verification

Author: hero78119

ceno_recursion/src/bin/e2e_aggregate.rs

@@ -106,6 +106,10 @@ struct Args {
     // max cycle per shard
     #[arg(long, default_value = "536870912")] // 536870912 = 2^29
     max_cycle_per_shard: u64,
+
+    /// Restrict proving to a single shard id for debugging.
+    #[arg(long)]
+    shard_id: Option<u64>,
 }
 
 fn main() {
@@ -221,6 +225,7 @@ fn main() {
 
     let max_steps = args.max_steps.unwrap_or(usize::MAX);
     let public_io_digest = public_io_words_to_digest_words(&public_io);
+    let target_shard_id = args.shard_id.map(|v| v as usize);
     let multi_prover = MultiProver::new(
         args.prover_id as usize,
         args.num_provers as usize,
@@ -240,7 +245,7 @@ fn main() {
             public_io_digest,
             max_steps,
             Checkpoint::Complete,
-            None,
+            target_shard_id,
         );
 
     let zkvm_proofs = result

ceno_zkvm/src/bin/e2e.rs

@@ -350,7 +350,8 @@ fn run_inner<
     let vk_bytes = bincode::serialize(&vk).unwrap();
     fs::write(&vk_file, vk_bytes).unwrap();
 
-    if checkpoint > Checkpoint::PrepVerify && target_shard_id.is_none() {
+    if checkpoint > Checkpoint::PrepVerify && (target_shard_id.is_none() || zkvm_proofs.len() == 1)
+    {
         let verifier = ZKVMVerifier::new(vk);
         verify(zkvm_proofs.clone(), &verifier).expect("Verification failed");
         soundness_test(zkvm_proofs.first().cloned().unwrap(), &verifier);

ceno_zkvm/src/e2e.rs

@@ -1759,8 +1759,10 @@ pub fn run_e2e_with_checkpoint<
         &init_full_mem,
     );
 
-    if target_shard_id.is_some() {
-        // skip verify as the proof are in-completed
+    let can_verify_target_shard = target_shard_id.is_none() || zkvm_proofs.len() == 1;
+    if !can_verify_target_shard {
+        // Partial multi-shard subsets still skip verification because the
+        // continuation chain between omitted shards is unavailable.
         return E2ECheckpointResult {
             proofs: Some(zkvm_proofs),
             vk: Some(vk),
@@ -1775,13 +1777,25 @@ pub fn run_e2e_with_checkpoint<
             proofs: Some(zkvm_proofs.clone()),
             vk: Some(vk),
             next_step: Some(Box::new(move || {
-                run_e2e_verify(&verifier, zkvm_proofs, exit_code, max_steps)
+                run_e2e_verify(
+                    &verifier,
+                    zkvm_proofs,
+                    exit_code,
+                    max_steps,
+                    target_shard_id,
+                )
             })),
         };
     }
 
     let start = std::time::Instant::now();
-    run_e2e_verify(&verifier, zkvm_proofs.clone(), exit_code, max_steps);
+    run_e2e_verify(
+        &verifier,
+        zkvm_proofs.clone(),
+        exit_code,
+        max_steps,
+        target_shard_id,
+    );
     tracing::debug!("verified in {:?}", start.elapsed());
 
     E2ECheckpointResult {
@@ -2023,15 +2037,29 @@ pub fn run_e2e_verify<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>>(
     zkvm_proofs: Vec<ZKVMProof<E, PCS>>,
     exit_code: Option<u32>,
     max_steps: usize,
+    target_shard_id: Option<usize>,
 ) {
     let transcripts = (0..zkvm_proofs.len())
         .map(|_| Transcript::new(b"riscv"))
         .collect_vec();
-    assert!(
+    let expect_halt = zkvm_proofs
+        .last()
+        .map(|proof| proof.has_halt(&verifier.vk))
+        .unwrap_or(exit_code.is_some());
+    let verified = if target_shard_id.is_some() && zkvm_proofs.len() == 1 {
         verifier
-            .verify_proofs_halt(zkvm_proofs, transcripts, exit_code.is_some())
-            .expect("verify proof return with error"),
-    );
+            .verify_shard_proof_halt(
+                zkvm_proofs.into_iter().next().unwrap(),
+                transcripts.into_iter().next().unwrap(),
+                expect_halt,
+            )
+            .expect("verify proof return with error")
+    } else {
+        verifier
+            .verify_proofs_halt(zkvm_proofs, transcripts, expect_halt)
+            .expect("verify proof return with error")
+    };
+    assert!(verified);
     match exit_code {
         Some(0) => tracing::info!("exit code 0. Success."),
         Some(code) => tracing::error!("exit code {}. Failure.", code),
@@ -2099,11 +2127,19 @@ pub fn verify<E: ExtensionField, PCS: PolynomialCommitmentScheme<E> + serde::Ser
     {
         Instrumented::<<<E as ExtensionField>::BaseField as PoseidonField>::P>::clear_metrics();
     }
-    let transcripts = (0..zkvm_proofs.len())
-        .map(|_| Transcript::new(b"riscv"))
-        .collect_vec();
     let has_halt = zkvm_proofs.last().unwrap().has_halt(&verifier.vk);
-    verifier.verify_proofs_halt(zkvm_proofs, transcripts, has_halt)?;
+    if zkvm_proofs.len() == 1 {
+        verifier.verify_shard_proof_halt(
+            zkvm_proofs.into_iter().next().unwrap(),
+            Transcript::new(b"riscv"),
+            has_halt,
+        )?;
+    } else {
+        let transcripts = (0..zkvm_proofs.len())
+            .map(|_| Transcript::new(b"riscv"))
+            .collect_vec();
+        verifier.verify_proofs_halt(zkvm_proofs, transcripts, has_halt)?;
+    }
     // print verification statistics such as hash count
     #[cfg(debug_assertions)]
     {

ceno_zkvm/src/scheme/verifier.rs

@@ -121,6 +121,35 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
         self.verify_proofs_halt(vec![vm_proof], vec![transcript], expect_halt)
     }
 
+    /// Verify a single shard proof as a standalone segment.
+    ///
+    /// Unlike `verify_proof_halt`, this checks proof validity and halt/segment
+    /// invariants for one shard only and intentionally skips cross-shard
+    /// continuation checks such as init_pc/heap chaining.
+    pub fn verify_shard_proof_halt(
+        &self,
+        vm_proof: ZKVMProof<E, PCS>,
+        transcript: impl ForkableTranscript<E>,
+        expect_halt: bool,
+    ) -> Result<bool, ZKVMError> {
+        let has_halt = vm_proof.has_halt(&self.vk);
+        if has_halt != expect_halt {
+            return Err(ZKVMError::VerifyError(
+                format!("shard proof ecall/halt mismatch: expected {expect_halt} != {has_halt}",)
+                    .into(),
+            ));
+        }
+
+        assert_eq!(
+            vm_proof.public_values.query_by_index::<E>(INIT_CYCLE_IDX),
+            E::BaseField::from_canonical_u64(Tracer::SUBCYCLES_PER_INSN)
+        );
+
+        let shard_id = vm_proof.public_values.shard_id as usize;
+        self.verify_proof_validity(shard_id, vm_proof, transcript)?;
+        Ok(true)
+    }
+
     /// Verify a trace from start to optional halt.
     pub fn verify_proofs_halt(
         &self,
@@ -237,8 +266,15 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
             }
         }
 
-        // check shard id
-        assert_eq!(vm_proof.public_values.shard_id, shard_id as u32);
+        if vm_proof.public_values.shard_id != shard_id as u32 {
+            return Err(ZKVMError::VerifyError(
+                format!(
+                    "proof shard_id mismatch: expected {} != {}",
+                    shard_id, vm_proof.public_values.shard_id
+                )
+                .into(),
+            ));
+        }
 
         // write fixed commitment to transcript
         // TODO check soundness if there is no fixed_commit but got fixed proof?