Feat: add rw mismatch checkers and lookup checker in mock prover for …

…zkvm (#571)

This PR is extracted from #526 and it aims to improve current mock
prover to find out constraint errors in the zkvm setting.

- [x] lookup errors.
- [x] read / write not equal for global state 
- [x] read / write not equal for registers
- [x] read / write not equal for memory

ceno_zkvm/Cargo.toml

@@ -11,7 +11,9 @@ goldilocks.workspace = true
 rand_chacha.workspace = true
 rayon.workspace = true
 serde.workspace = true
+serde_json.workspace = true
 
+base64 = "0.22"
 ceno_emul = { path = "../ceno_emul" }
 ff_ext = { path = "../ff_ext" }
 mpcs = { path = "../mpcs" }
@@ -34,11 +36,9 @@ tempfile = "3.13"
 thread_local = "1.1"
 
 [dev-dependencies]
-base64 = "0.22"
 cfg-if.workspace = true
 criterion.workspace = true
 pprof.workspace = true
-serde_json.workspace = true
 
 [build-dependencies]
 glob = "0.3"

ceno_zkvm/examples/riscv_opcodes.rs

@@ -3,7 +3,7 @@ use std::{panic, time::Instant};
 use ceno_zkvm::{
     declare_program,
     instructions::riscv::{Rv32imConfig, constants::EXIT_PC},
-    scheme::prover::ZKVMProver,
+    scheme::{mock_prover::MockProver, prover::ZKVMProver},
     state::GlobalState,
     tables::{
         DynVolatileRamTable, MemFinalRecord, MemTable, ProgramTableCircuit, init_program_data,
@@ -134,7 +134,7 @@ fn main() {
 
     let pk = zkvm_cs
         .clone()
-        .key_gen::<Pcs>(pp.clone(), vp.clone(), zkvm_fixed_traces)
+        .key_gen::<Pcs>(pp.clone(), vp.clone(), zkvm_fixed_traces.clone())
         .expect("keygen failed");
     let vk = pk.get_vk();
 
@@ -274,6 +274,13 @@ fn main() {
             .assign_table_circuit::<ExampleProgramTableCircuit<E>>(&zkvm_cs, &prog_config, &program)
             .unwrap();
 
+        MockProver::assert_satisfied_full(
+            zkvm_cs.clone(),
+            zkvm_fixed_traces.clone(),
+            &zkvm_witness,
+            &pi,
+        );
+
         let timer = Instant::now();
 
         let transcript = Transcript::new(b"riscv");

ceno_zkvm/src/chip_handler/general.rs

@@ -93,27 +93,31 @@ impl<'a, E: ExtensionField> CircuitBuilder<'a, E> {
     pub fn r_table_record<NR, N>(
         &mut self,
         name_fn: N,
+        ram_type: RAMType,
         table_spec: SetTableSpec,
-        rlc_record: Expression<E>,
+        record: Vec<Expression<E>>,
     ) -> Result<(), ZKVMError>
     where
         NR: Into<String>,
         N: FnOnce() -> NR,
     {
-        self.cs.r_table_record(name_fn, table_spec, rlc_record)
+        self.cs
+            .r_table_record(name_fn, ram_type, table_spec, record)
     }
 
     pub fn w_table_record<NR, N>(
         &mut self,
         name_fn: N,
+        ram_type: RAMType,
         table_spec: SetTableSpec,
-        rlc_record: Expression<E>,
+        record: Vec<Expression<E>>,
     ) -> Result<(), ZKVMError>
     where
         NR: Into<String>,
         N: FnOnce() -> NR,
     {
-        self.cs.w_table_record(name_fn, table_spec, rlc_record)
+        self.cs
+            .w_table_record(name_fn, ram_type, table_spec, record)
     }
 
     /// Fetch an instruction at a given PC from the Program table.

ceno_zkvm/src/circuit_builder.rs

@@ -334,13 +334,15 @@ impl<E: ExtensionField> ConstraintSystem<E> {
     pub fn r_table_record<NR, N>(
         &mut self,
         name_fn: N,
+        ram_type: RAMType,
         table_spec: SetTableSpec,
-        rlc_record: Expression<E>,
+        record: Vec<Expression<E>>,
     ) -> Result<(), ZKVMError>
     where
         NR: Into<String>,
         N: FnOnce() -> NR,
     {
+        let rlc_record = self.rlc_chip_record(record.clone());
         assert_eq!(
             rlc_record.degree(),
             1,
@@ -353,20 +355,23 @@ impl<E: ExtensionField> ConstraintSystem<E> {
         });
         let path = self.ns.compute_path(name_fn().into());
         self.r_table_expressions_namespace_map.push(path);
+        self.r_ram_types.push((ram_type, record));
 
         Ok(())
     }
 
     pub fn w_table_record<NR, N>(
         &mut self,
         name_fn: N,
+        ram_type: RAMType,
         table_spec: SetTableSpec,
-        rlc_record: Expression<E>,
+        record: Vec<Expression<E>>,
     ) -> Result<(), ZKVMError>
     where
         NR: Into<String>,
         N: FnOnce() -> NR,
     {
+        let rlc_record = self.rlc_chip_record(record.clone());
         assert_eq!(
             rlc_record.degree(),
             1,
@@ -379,6 +384,7 @@ impl<E: ExtensionField> ConstraintSystem<E> {
         });
         let path = self.ns.compute_path(name_fn().into());
         self.w_table_expressions_namespace_map.push(path);
+        self.w_ram_types.push((ram_type, record));
 
         Ok(())
     }

ceno_zkvm/src/scheme.rs

@@ -12,7 +12,6 @@ pub mod prover;
 pub mod utils;
 pub mod verifier;
 
-#[cfg(test)]
 pub mod mock_prover;
 #[cfg(test)]
 mod tests;