[Feat] Calculate proof size and number of hashes (#712)

This PR is base on #692 and provided required features to calculate the
size of proofs and number of hashes

Now it has been implied in the e2e test `test_single_add_instance_e2e`
and `fibonacci` bench, obtained following output:

> encoded zkvm proof size: 2756371, hash_num: 4411

The encoding of proof is done by `bincode` to kept the serialization as
compact as possible

---------

Co-authored-by: Matthias Goergens <[email protected]>

ceno_zkvm/Cargo.toml

@@ -44,6 +44,7 @@ tempfile = "3.14"
 thread_local = "1.1"
 
 [dev-dependencies]
+bincode = "1"
 cfg-if.workspace = true
 criterion.workspace = true
 pprof2.workspace = true

ceno_zkvm/benches/fibonacci.rs

@@ -10,6 +10,7 @@ use ceno_zkvm::{
     e2e::{Checkpoint, Preset, run_e2e_with_checkpoint, setup_platform},
 };
 use criterion::*;
+use transcript::{BasicTranscriptWitStat, StatisticRecorder};
 
 use goldilocks::GoldilocksExt2;
 use mpcs::BasefoldDefault;
@@ -43,6 +44,36 @@ fn setup() -> (Program, Platform) {
 fn fibonacci_prove(c: &mut Criterion) {
     let (program, platform) = setup();
     for max_steps in [1usize << 20, 1usize << 21, 1usize << 22] {
+        // estimate proof size data first
+        let (sanity_check_state, _) = run_e2e_with_checkpoint::<E, Pcs>(
+            program.clone(),
+            platform.clone(),
+            vec![],
+            max_steps,
+            Checkpoint::PrepSanityCheck,
+        );
+        assert!(
+            sanity_check_state.is_some(),
+            "PrepSanityCheck do not provide proof and verifier"
+        );
+        if let Some((proof, verifier)) = sanity_check_state {
+            let serialize_size = bincode::serialize(&proof).unwrap().len();
+            let stat_recorder = StatisticRecorder::default();
+            let transcript = BasicTranscriptWitStat::new(&stat_recorder, b"riscv");
+            assert!(
+                verifier
+                    .verify_proof_halt(proof, transcript, false)
+                    .expect("verify proof return with error"),
+            );
+            println!();
+            println!(
+                "max_steps = {}, proof size = {}, hashes count = {}",
+                max_steps,
+                serialize_size,
+                stat_recorder.into_inner().field_appended_num
+            );
+        }
+
         // expand more input size once runtime is acceptable
         let mut group = c.benchmark_group(format!("fibonacci_max_steps_{}", max_steps));
         group.sample_size(NUM_SAMPLES);

ceno_zkvm/src/scheme.rs

@@ -16,7 +16,7 @@ pub mod mock_prover;
 #[cfg(test)]
 mod tests;
 
-#[derive(Clone)]
+#[derive(Clone, Serialize)]
 pub struct ZKVMOpcodeProof<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> {
     // TODO support >1 opcodes
     pub num_instances: usize,
@@ -115,7 +115,7 @@ impl PublicValues<u32> {
 /// Map circuit names to
 /// - an opcode or table proof,
 /// - an index unique across both types.
-#[derive(Clone)]
+#[derive(Clone, Serialize)]
 pub struct ZKVMProof<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> {
     // TODO preserve in serde only for auxiliary public input
     // other raw value can be construct by verifier directly.

ceno_zkvm/src/scheme/tests.rs

@@ -14,7 +14,7 @@ use mpcs::{Basefold, BasefoldDefault, BasefoldRSParams, PolynomialCommitmentSche
 use multilinear_extensions::{
     mle::IntoMLE, util::ceil_log2, virtual_poly_v2::ArcMultilinearExtension,
 };
-use transcript::{BasicTranscript, Transcript};
+use transcript::{BasicTranscript, BasicTranscriptWitStat, StatisticRecorder, Transcript};
 
 use crate::{
     circuit_builder::CircuitBuilder,
@@ -155,8 +155,9 @@ fn test_rw_lk_expression_combination() {
             .expect("create_proof failed");
 
         // verify proof
+        let stat_recorder = StatisticRecorder::default();
         let verifier = ZKVMVerifier::new(vk.clone());
-        let mut v_transcript = BasicTranscript::new(b"test");
+        let mut v_transcript = BasicTranscriptWitStat::new(&stat_recorder, b"test");
         // write commitment into transcript and derive challenges from it
         Pcs::write_commitment(&proof.wits_commit, &mut v_transcript).unwrap();
         let verifier_challenges = [
@@ -178,6 +179,10 @@ fn test_rw_lk_expression_combination() {
                 &verifier_challenges,
             )
             .expect("verifier failed");
+        println!(
+            "hashed fields {}",
+            stat_recorder.into_inner().field_appended_num
+        );
     }
 
     // <lookup count, rw count>
@@ -289,11 +294,21 @@ fn test_single_add_instance_e2e() {
         .create_proof(zkvm_witness, pi, transcript)
         .expect("create_proof failed");
 
-    let transcript = BasicTranscript::new(b"riscv");
-    assert!(
-        verifier
-            .verify_proof(zkvm_proof, transcript)
-            .expect("verify proof return with error"),
+    let encoded_bin = bincode::serialize(&zkvm_proof).unwrap();
+
+    let stat_recorder = StatisticRecorder::default();
+    {
+        let transcript = BasicTranscriptWitStat::new(&stat_recorder, b"riscv");
+        assert!(
+            verifier
+                .verify_proof(zkvm_proof, transcript)
+                .expect("verify proof return with error"),
+        );
+    }
+    println!(
+        "encoded zkvm proof size: {}, hash_num: {}",
+        encoded_bin.len(),
+        stat_recorder.into_inner().field_appended_num
     );
 }
 

mpcs/src/basefold/structure.rs

@@ -5,7 +5,7 @@ use crate::{
 use core::fmt::Debug;
 use ff_ext::ExtensionField;
 
-use serde::{Deserialize, Serialize, de::DeserializeOwned};
+use serde::{Deserialize, Serialize, Serializer, de::DeserializeOwned};
 
 use multilinear_extensions::mle::FieldType;
 
@@ -247,6 +247,15 @@ where
 #[derive(Debug)]
 pub struct Basefold<E: ExtensionField, Spec: BasefoldSpec<E>>(PhantomData<(E, Spec)>);
 
+impl<E: ExtensionField, Spec: BasefoldSpec<E>> Serialize for Basefold<E, Spec> {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: Serializer,
+    {
+        serializer.serialize_str("base_fold")
+    }
+}
+
 pub type BasefoldDefault<F> = Basefold<F, BasefoldRSParams>;
 
 impl<E: ExtensionField, Spec: BasefoldSpec<E>> Clone for Basefold<E, Spec> {

transcript/src/lib.rs

@@ -4,8 +4,10 @@
 #![feature(generic_arg_infer)]
 
 pub mod basic;
+mod statistics;
 pub mod syncronized;
 pub use basic::BasicTranscript;
+pub use statistics::{BasicTranscriptWitStat, StatisticRecorder};
 pub use syncronized::TranscriptSyncronized;
 
 #[derive(Default, Copy, Clone, Eq, PartialEq, Debug)]

transcript/src/statistics.rs

@@ -0,0 +1,59 @@
+use crate::{BasicTranscript, Challenge, ForkableTranscript, Transcript};
+use ff_ext::ExtensionField;
+use std::cell::RefCell;
+
+#[derive(Debug, Default)]
+pub struct Statistic {
+    pub field_appended_num: u32,
+}
+
+pub type StatisticRecorder = RefCell<Statistic>;
+
+#[derive(Clone)]
+pub struct BasicTranscriptWitStat<'a, E: ExtensionField> {
+    inner: BasicTranscript<E>,
+    stat: &'a StatisticRecorder,
+}
+
+impl<'a, E: ExtensionField> BasicTranscriptWitStat<'a, E> {
+    pub fn new(stat: &'a StatisticRecorder, label: &'static [u8]) -> Self {
+        Self {
+            inner: BasicTranscript::<_>::new(label),
+            stat,
+        }
+    }
+}
+
+impl<E: ExtensionField> Transcript<E> for BasicTranscriptWitStat<'_, E> {
+    fn append_field_elements(&mut self, elements: &[E::BaseField]) {
+        self.stat.borrow_mut().field_appended_num += 1;
+        self.inner.append_field_elements(elements)
+    }
+
+    fn append_field_element_ext(&mut self, element: &E) {
+        self.stat.borrow_mut().field_appended_num += E::DEGREE as u32;
+        self.inner.append_field_element_ext(element)
+    }
+
+    fn read_challenge(&mut self) -> Challenge<E> {
+        self.inner.read_challenge()
+    }
+
+    fn read_field_element_exts(&self) -> Vec<E> {
+        self.inner.read_field_element_exts()
+    }
+
+    fn read_field_element(&self) -> E::BaseField {
+        self.inner.read_field_element()
+    }
+
+    fn send_challenge(&self, challenge: E) {
+        self.inner.send_challenge(challenge)
+    }
+
+    fn commit_rolling(&mut self) {
+        self.inner.commit_rolling()
+    }
+}
+
+impl<E: ExtensionField> ForkableTranscript<E> for BasicTranscriptWitStat<'_, E> {}