Feat: simplify basefold's verifier (#974)

## Rationale

The current API used by `mpcs` crate is tightly coupled with our zkvm's
design. This is not desireable.

## Summary

We decide to be align with the `Plonky3` style API. 

- [x] #979.
- [x] #981.
- [x] #982

Author: kunxian-xia

ceno_zkvm/benches/riscv_add.rs

@@ -1,4 +1,4 @@
-use std::{collections::BTreeMap, time::Duration};
+use std::time::Duration;
 
 use ceno_zkvm::{
     self,
@@ -77,10 +77,7 @@ fn bench_add(c: &mut Criterion) {
                     for _ in 0..iters {
                         // generate mock witness
                         let num_instances = 1 << instance_num_vars;
-                        let rmms = BTreeMap::from([(
-                            0,
-                            RowMajorMatrix::rand(&mut OsRng, num_instances, num_witin),
-                        )]);
+                        let rmms = vec![RowMajorMatrix::rand(&mut OsRng, num_instances, num_witin)];
 
                         let instant = std::time::Instant::now();
                         let num_instances = 1 << instance_num_vars;

ceno_zkvm/src/error.rs

@@ -13,6 +13,7 @@ pub enum ZKVMError {
     UtilError(UtilError),
     WitnessNotFound(String),
     InvalidWitness(String),
+    InvalidProof(String),
     VKNotFound(String),
     FixedTraceNotFound(String),
     VerifyError(String),

ceno_zkvm/src/scheme.rs

@@ -49,6 +49,7 @@ pub struct ZKVMChipProof<E: ExtensionField> {
 
     pub tower_proof: TowerProofs<E>,
 
+    pub num_instances: usize,
     pub fixed_in_evals: Vec<E>,
     pub wits_in_evals: Vec<E>,
 }
@@ -114,32 +115,25 @@ pub struct ZKVMProof<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> {
     pub raw_pi: Vec<Vec<E::BaseField>>,
     // the evaluation of raw_pi.
     pub pi_evals: Vec<E>,
-    // circuit size -> instance mapping
-    pub num_instances: Vec<(usize, usize)>,
-    opcode_proofs: BTreeMap<usize, ZKVMChipProof<E>>,
-    table_proofs: BTreeMap<usize, ZKVMChipProof<E>>,
+    chip_proofs: BTreeMap<usize, ZKVMChipProof<E>>,
     witin_commit: <PCS as PolynomialCommitmentScheme<E>>::Commitment,
-    pub fixed_witin_opening_proof: PCS::Proof,
+    pub opening_proof: PCS::Proof,
 }
 
 impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProof<E, PCS> {
     pub fn new(
         raw_pi: Vec<Vec<E::BaseField>>,
         pi_evals: Vec<E>,
-        opcode_proofs: BTreeMap<usize, ZKVMChipProof<E>>,
-        table_proofs: BTreeMap<usize, ZKVMChipProof<E>>,
+        chip_proofs: BTreeMap<usize, ZKVMChipProof<E>>,
         witin_commit: <PCS as PolynomialCommitmentScheme<E>>::Commitment,
-        fixed_witin_opening_proof: PCS::Proof,
-        num_instances: Vec<(usize, usize)>,
+        opening_proof: PCS::Proof,
     ) -> Self {
         Self {
             raw_pi,
             pi_evals,
-            opcode_proofs,
-            table_proofs,
+            chip_proofs,
             witin_commit,
-            fixed_witin_opening_proof,
-            num_instances,
+            opening_proof,
         }
     }
 
@@ -164,23 +158,19 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProof<E, PCS> {
     }
 
     pub fn num_circuits(&self) -> usize {
-        self.opcode_proofs.len() + self.table_proofs.len()
+        self.chip_proofs.len()
     }
 
     pub fn has_halt(&self, vk: &ZKVMVerifyingKey<E, PCS>) -> bool {
+        let halt_circuit_index = vk
+            .circuit_vks
+            .keys()
+            .position(|circuit_name| *circuit_name == HaltInstruction::<E>::name())
+            .expect("halt circuit not exist");
         let halt_instance_count = self
-            .num_instances
-            .iter()
-            .find_map(|(circuit_index, num_instances)| {
-                (*circuit_index
-                    == vk
-                        .circuit_vks
-                        .keys()
-                        .position(|circuit_name| *circuit_name == HaltInstruction::<E>::name())
-                        .expect("halt circuit not exist"))
-                .then_some(*num_instances)
-            })
-            .unwrap_or(0);
+            .chip_proofs
+            .get(&halt_circuit_index)
+            .map_or(0, |proof| proof.num_instances);
         if halt_instance_count > 0 {
             assert_eq!(
                 halt_instance_count, 1,
@@ -203,39 +193,11 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E> + Serialize> fmt::Dis
         let mpcs_opcode_commitment =
             bincode::serialized_size(&self.witin_commit).expect("serialization error");
         let mpcs_opcode_opening =
-            bincode::serialized_size(&self.fixed_witin_opening_proof).expect("serialization error");
+            bincode::serialized_size(&self.opening_proof).expect("serialization error");
 
-        // opcode circuit for tower proof size
-        let tower_proof_opcode = self
-            .opcode_proofs
-            .iter()
-            .map(|(circuit_index, proof)| {
-                let size = bincode::serialized_size(&proof.tower_proof);
-                size.inspect(|size| {
-                    *by_circuitname_stats.entry(circuit_index).or_insert(0) += size;
-                })
-            })
-            .collect::<Result<Vec<u64>, _>>()
-            .expect("serialization error")
-            .iter()
-            .sum::<u64>();
-        // opcode circuit main sumcheck
-        let main_sumcheck_opcode = self
-            .opcode_proofs
-            .iter()
-            .map(|(circuit_index, proof)| {
-                let size = bincode::serialized_size(&proof.main_sumcheck_proofs);
-                size.inspect(|size| {
-                    *by_circuitname_stats.entry(circuit_index).or_insert(0) += size;
-                })
-            })
-            .collect::<Result<Vec<u64>, _>>()
-            .expect("serialization error")
-            .iter()
-            .sum::<u64>();
-        // table circuit for tower proof size
-        let tower_proof_table = self
-            .table_proofs
+        // tower proof size
+        let tower_proof = self
+            .chip_proofs
             .iter()
             .map(|(circuit_index, proof)| {
                 let size = bincode::serialized_size(&proof.tower_proof);
@@ -247,9 +209,9 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E> + Serialize> fmt::Dis
             .expect("serialization error")
             .iter()
             .sum::<u64>();
-        // table circuit same r sumcheck
-        let same_r_sumcheck_table = self
-            .table_proofs
+        // main sumcheck
+        let main_sumcheck = self
+            .chip_proofs
             .iter()
             .map(|(circuit_index, proof)| {
                 let size = bincode::serialized_size(&proof.main_sumcheck_proofs);
@@ -286,20 +248,16 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E> + Serialize> fmt::Dis
             "overall_size {:.2}mb. \n\
             mpcs commitment {:?}% \n\
             mpcs opening {:?}% \n\
-            opcode tower proof {:?}% \n\
-            opcode main sumcheck proof {:?}% \n\
-            table tower proof {:?}% \n\
-            table same r sumcheck proof {:?}% \n\n\
+            tower proof {:?}% \n\
+            main sumcheck proof {:?}% \n\
             by circuit_name break down: \n\
             {}
             ",
             byte_to_mb(overall_size),
             (mpcs_opcode_commitment * 100).div(overall_size),
             (mpcs_opcode_opening * 100).div(overall_size),
-            (tower_proof_opcode * 100).div(overall_size),
-            (main_sumcheck_opcode * 100).div(overall_size),
-            (tower_proof_table * 100).div(overall_size),
-            (same_r_sumcheck_table * 100).div(overall_size),
+            (tower_proof * 100).div(overall_size),
+            (main_sumcheck * 100).div(overall_size),
             by_circuitname_stats,
         )
     }

ceno_zkvm/src/scheme/cpu/mod.rs

@@ -269,7 +269,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> TraceCommitter<CpuBa
             self.pp = Some(prover_param);
             self.pp.as_ref().unwrap()
         };
-        let pcs_data = PCS::batch_commit(prover_param, traces).unwrap();
+        let pcs_data = PCS::batch_commit(prover_param, traces.into_values().collect_vec()).unwrap();
         let commit = PCS::get_pure_commitment(&pcs_data);
         let mles = PCS::get_arc_mle_witness_from_commitment(&pcs_data)
             .into_par_iter()
@@ -752,22 +752,40 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> OpeningProver<CpuBac
         witness_data: PCS::CommitmentWithWitness,
         fixed_data: Option<Arc<PCS::CommitmentWithWitness>>,
         points: Vec<Point<E>>,
-        evals: Vec<Vec<E>>,
+        mut evals: Vec<Vec<E>>, // where each inner Vec<E> = wit_evals + fixed_evals
         circuit_num_polys: &[(usize, usize)],
         num_instances: &[(usize, usize)],
         transcript: &mut impl Transcript<E>,
     ) -> PCS::Proof {
-        PCS::batch_open(
-            self.pp.as_ref().unwrap(),
-            num_instances,
-            fixed_data.as_ref().map(|f| f.as_ref()),
+        let mut rounds = vec![];
+        rounds.push((
             &witness_data,
-            &points,
-            &evals,
-            circuit_num_polys,
-            transcript,
-        )
-        .unwrap()
+            points
+                .iter()
+                .zip_eq(evals.iter_mut())
+                .zip_eq(num_instances.iter())
+                .map(|((point, evals), (chip_idx, _))| {
+                    let (num_witin, _) = circuit_num_polys[*chip_idx];
+                    (point.clone(), evals.drain(..num_witin).collect_vec())
+                })
+                .collect_vec(),
+        ));
+        if let Some(fixed_data) = fixed_data.as_ref().map(|f| f.as_ref()) {
+            rounds.push((
+                fixed_data,
+                points
+                    .iter()
+                    .zip_eq(evals.iter_mut())
+                    .zip_eq(num_instances.iter())
+                    .filter(|(_, (chip_idx, _))| {
+                        let (_, num_fixed) = circuit_num_polys[*chip_idx];
+                        num_fixed > 0
+                    })
+                    .map(|((point, evals), _)| (point.clone(), evals.to_vec()))
+                    .collect_vec(),
+            ));
+        }
+        PCS::batch_open(self.pp.as_ref().unwrap(), rounds, transcript).unwrap()
     }
 }
 

ceno_zkvm/src/scheme/prover.rs

@@ -84,8 +84,7 @@ impl<
     ) -> Result<ZKVMProof<E, PCS>, ZKVMError> {
         let raw_pi = pi.to_vec::<E>();
         let mut pi_evals = ZKVMProof::<E, PCS>::pi_evals(&raw_pi);
-        let mut opcode_proofs: BTreeMap<usize, ZKVMChipProof<E>> = BTreeMap::new();
-        let mut table_proofs: BTreeMap<usize, ZKVMChipProof<E>> = BTreeMap::new();
+        let mut chip_proofs: BTreeMap<usize, ZKVMChipProof<E>> = BTreeMap::new();
 
         let span = entered_span!("commit_to_pi", profiling_1 = true);
         // including raw public input to transcript
@@ -233,23 +232,23 @@ impl<
                     );
                     points.push(input_opening_point);
                     evaluations.push(opcode_proof.wits_in_evals.clone());
-                    opcode_proofs.insert(index, opcode_proof);
+                    chip_proofs.insert(index, opcode_proof);
                 } else {
                     // FIXME: PROGRAM table circuit is not guaranteed to have 2^n instances
                     input.num_instances = 1 << input.log2_num_instances();
-                    let (table_proof, pi_in_evals, input_opening_point) = self.create_chip_proof(
-                        circuit_name,
-                        pk,
-                        input,
-                        &mut transcript,
-                        &challenges,
-                    )?;
+                    let (mut table_proof, pi_in_evals, input_opening_point) = self
+                        .create_chip_proof(circuit_name, pk, input, &mut transcript, &challenges)?;
                     points.push(input_opening_point);
-                    evaluations.push(table_proof.wits_in_evals.clone());
-                    if cs.num_fixed() > 0 {
-                        evaluations.push(table_proof.fixed_in_evals.clone());
-                    }
-                    table_proofs.insert(index, table_proof);
+                    evaluations.push(
+                        [
+                            table_proof.wits_in_evals.clone(),
+                            table_proof.fixed_in_evals.clone(),
+                        ]
+                        .concat(),
+                    );
+                    // FIXME: PROGRAM table circuit is not guaranteed to have 2^n instances
+                    table_proof.num_instances = num_instances;
+                    chip_proofs.insert(index, table_proof);
                     for (idx, eval) in pi_in_evals {
                         pi_evals[idx] = eval;
                     }
@@ -282,12 +281,9 @@ impl<
         let vm_proof = ZKVMProof::new(
             raw_pi,
             pi_evals,
-            opcode_proofs,
-            table_proofs,
+            chip_proofs,
             witin_commit,
             mpcs_opening_proof,
-            // verifier need this information from prover to achieve non-uniform design.
-            num_instances,
         );
         exit_span!(main_proofs_span);
 
@@ -363,6 +359,7 @@ impl<
                 tower_proof,
                 fixed_in_evals,
                 wits_in_evals,
+                num_instances: input.num_instances,
             },
             pi_in_evals,
             input_opening_point,

ceno_zkvm/src/scheme/tests.rs

@@ -1,4 +1,4 @@
-use std::{collections::BTreeMap, marker::PhantomData};
+use std::marker::PhantomData;
 
 use crate::{
     circuit_builder::CircuitBuilder,
@@ -138,12 +138,8 @@ fn test_rw_lk_expression_combination() {
         let rmm = zkvm_witness.into_iter_sorted().next().unwrap().1.remove(0);
         let wits_in = rmm.to_mles();
         // commit to main traces
-        let commit_with_witness = Pcs::batch_commit_and_write(
-            &prover.pk.pp,
-            vec![(0, rmm)].into_iter().collect::<BTreeMap<_, _>>(),
-            &mut transcript,
-        )
-        .unwrap();
+        let commit_with_witness =
+            Pcs::batch_commit_and_write(&prover.pk.pp, vec![rmm], &mut transcript).unwrap();
         let witin_commit = Pcs::get_pure_commitment(&commit_with_witness);
 
         let wits_in = wits_in.into_iter().map(|v| v.into()).collect_vec();
@@ -189,7 +185,6 @@ fn test_rw_lk_expression_combination() {
                 name.as_str(),
                 verifier.vk.circuit_vks.get(&name).unwrap(),
                 &proof,
-                num_instances,
                 &[],
                 &mut v_transcript,
                 NUM_FANIN,