From b74c3b4ab10ac4748411dd5762e7485ba16adf6f Mon Sep 17 00:00:00 2001
From: Ming <hero78119@gmail.com>
Date: Thu, 12 Sep 2024 22:59:54 +0800
Subject: [PATCH] fix add bench build error & example (#211)

---
 ceno_zkvm/benches/riscv_add.rs  | 25 +++++++++++++++++++------
 ceno_zkvm/examples/riscv_add.rs | 17 +++++++++--------
 ceno_zkvm/src/structs.rs        |  2 +-
 3 files changed, 29 insertions(+), 15 deletions(-)

diff --git a/ceno_zkvm/benches/riscv_add.rs b/ceno_zkvm/benches/riscv_add.rs
index aa05b218f..c1bb982f9 100644
--- a/ceno_zkvm/benches/riscv_add.rs
+++ b/ceno_zkvm/benches/riscv_add.rs
@@ -2,9 +2,10 @@ use std::time::{Duration, Instant};
 
 use ark_std::test_rng;
 use ceno_zkvm::{
-    circuit_builder::{CircuitBuilder, ConstraintSystem},
+    self,
     instructions::{riscv::addsub::AddInstruction, Instruction},
     scheme::prover::ZKVMProver,
+    structs::{ZKVMConstraintSystem, ZKVMFixedTraces},
 };
 use const_env::from_env;
 use criterion::*;
@@ -62,11 +63,22 @@ fn bench_add(c: &mut Criterion) {
             RAYON_NUM_THREADS
         }
     };
-    let mut cs = ConstraintSystem::new(|| "risv_add");
-    let mut circuit_builder = CircuitBuilder::<GoldilocksExt2>::new(&mut cs);
-    let _ = AddInstruction::construct_circuit(&mut circuit_builder);
-    let pk = cs.key_gen(None);
-    let num_witin = pk.get_cs().num_witin;
+    let mut zkvm_cs = ZKVMConstraintSystem::default();
+    let _ = zkvm_cs.register_opcode_circuit::<AddInstruction<E>>();
+    let mut zkvm_fixed_traces = ZKVMFixedTraces::default();
+    zkvm_fixed_traces.register_opcode_circuit::<AddInstruction<E>>(&zkvm_cs);
+
+    let pk = zkvm_cs
+        .clone()
+        .key_gen(zkvm_fixed_traces)
+        .expect("keygen failed");
+
+    let circuit_pk = pk
+        .circuit_pks
+        .get(&AddInstruction::<E>::name())
+        .unwrap()
+        .clone();
+    let num_witin = circuit_pk.get_cs().num_witin;
 
     let prover = ZKVMProver::new(pk);
     let mut transcript = Transcript::new(b"riscv");
@@ -101,6 +113,7 @@ fn bench_add(c: &mut Criterion) {
                         let timer = Instant::now();
                         let _ = prover
                             .create_opcode_proof(
+                                &circuit_pk,
                                 wits_in,
                                 num_instances,
                                 max_threads,
diff --git a/ceno_zkvm/examples/riscv_add.rs b/ceno_zkvm/examples/riscv_add.rs
index 5ff5df03a..773fbcf01 100644
--- a/ceno_zkvm/examples/riscv_add.rs
+++ b/ceno_zkvm/examples/riscv_add.rs
@@ -27,6 +27,7 @@ const RAYON_NUM_THREADS: usize = 8;
 //  - x2 is initialized to -1,
 //  - x3 is initialized to loop bound.
 // we use x4 to hold the acc_sum.
+#[allow(clippy::unusual_byte_groupings)]
 const PROGRAM_ADD_LOOP: [u32; 4] = [
     // func7   rs2   rs1   f3  rd    opcode
     0b_0000000_00100_00001_000_00100_0110011, // add x4, x4, x1 <=> addi x4, x4, 1
@@ -106,7 +107,7 @@ fn main() {
     let verifier = ZKVMVerifier::new(vk);
 
     for instance_num_vars in args.start..args.end {
-        let num_instances = 1 << instance_num_vars;
+        let step_loop = 1 << (instance_num_vars - 1); // 1 step in loop contribute to 2 add instance
         let mut vm = VMState::new(CENO_PLATFORM);
         let pc_start = ByteAddr(CENO_PLATFORM.pc_start()).waddr();
 
@@ -114,7 +115,7 @@ fn main() {
         // vm.x4 += vm.x1
         vm.init_register_unsafe(1usize, 1);
         vm.init_register_unsafe(2usize, u32::MAX); // -1 in two's complement
-        vm.init_register_unsafe(3usize, num_instances as u32);
+        vm.init_register_unsafe(3usize, step_loop as u32);
         for (i, inst) in PROGRAM_ADD_LOOP.iter().enumerate() {
             vm.init_memory(pc_start + i, *inst);
         }
@@ -148,17 +149,17 @@ fn main() {
             .create_proof(zkvm_witness, max_threads, &mut transcript, &real_challenges)
             .expect("create_proof failed");
 
+        println!(
+            "AddInstruction::create_proof, instance_num_vars = {}, time = {}",
+            instance_num_vars,
+            timer.elapsed().as_secs_f64()
+        );
+
         let mut transcript = Transcript::new(b"riscv");
         assert!(
             verifier
                 .verify_proof(zkvm_proof, &mut transcript, &real_challenges)
                 .expect("verify proof return with error"),
         );
-
-        println!(
-            "AddInstruction::create_proof, instance_num_vars = {}, time = {}",
-            instance_num_vars,
-            timer.elapsed().as_secs_f64()
-        );
     }
 }
diff --git a/ceno_zkvm/src/structs.rs b/ceno_zkvm/src/structs.rs
index c7c136d66..e2bbe0b56 100644
--- a/ceno_zkvm/src/structs.rs
+++ b/ceno_zkvm/src/structs.rs
@@ -245,7 +245,7 @@ impl<E: ExtensionField> ZKVMWitnesses<E> {
 #[derive(Default)]
 pub struct ZKVMProvingKey<E: ExtensionField> {
     // pk for opcode and table circuits
-    pub(crate) circuit_pks: BTreeMap<String, ProvingKey<E>>,
+    pub circuit_pks: BTreeMap<String, ProvingKey<E>>,
 }
 
 impl<E: ExtensionField> ZKVMProvingKey<E> {