From 9dd45dbf88fb11406b9c459b33cd9f441074cc8e Mon Sep 17 00:00:00 2001
From: kunxian xia <xiakunxian130@gmail.com>
Date: Wed, 30 Oct 2024 01:14:54 +0800
Subject: [PATCH] include all opcodes in rv32im config

---
 ceno_zkvm/examples/fibonacci_elf.rs           |   4 +-
 ceno_zkvm/src/instructions/riscv/branch.rs    |   8 +-
 .../src/instructions/riscv/branch/blt.rs      |   4 +-
 .../src/instructions/riscv/branch/bltu.rs     |   4 +-
 ceno_zkvm/src/instructions/riscv/jump/jalr.rs |   4 +
 .../src/instructions/riscv/memory/store.rs    |   3 -
 ceno_zkvm/src/instructions/riscv/rv32im.rs    | 251 ++++++++++++++----
 ceno_zkvm/src/instructions/riscv/shift.rs     |  10 +-
 ceno_zkvm/src/instructions/riscv/shift_imm.rs |   3 +
 ceno_zkvm/src/scheme/prover.rs                |  31 ++-
 ceno_zkvm/src/structs.rs                      |  14 +-
 11 files changed, 245 insertions(+), 91 deletions(-)

diff --git a/ceno_zkvm/examples/fibonacci_elf.rs b/ceno_zkvm/examples/fibonacci_elf.rs
index 041032f89..9f483a5e1 100644
--- a/ceno_zkvm/examples/fibonacci_elf.rs
+++ b/ceno_zkvm/examples/fibonacci_elf.rs
@@ -11,6 +11,7 @@ use ceno_zkvm::{
     structs::{ZKVMConstraintSystem, ZKVMFixedTraces, ZKVMWitnesses},
     tables::{MemFinalRecord, ProgramTableCircuit, initial_memory, initial_registers},
 };
+use ff_ext::ff::Field;
 use goldilocks::GoldilocksExt2;
 use itertools::Itertools;
 use mpcs::{Basefold, BasefoldRSParams, PolynomialCommitmentScheme};
@@ -19,7 +20,6 @@ use std::{panic, time::Instant};
 use tracing_flame::FlameLayer;
 use tracing_subscriber::{EnvFilter, Registry, fmt, layer::SubscriberExt};
 use transcript::Transcript;
-use ff_ext::ff::Field;
 
 fn main() {
     type E = GoldilocksExt2;
@@ -153,7 +153,7 @@ fn main() {
         .assign_table_circuit::<ExampleProgramTableCircuit<E>>(
             &zkvm_cs,
             &prog_config,
-            &vm.program(),
+            vm.program(),
         )
         .unwrap();
 
diff --git a/ceno_zkvm/src/instructions/riscv/branch.rs b/ceno_zkvm/src/instructions/riscv/branch.rs
index f9cb2c311..6fe929980 100644
--- a/ceno_zkvm/src/instructions/riscv/branch.rs
+++ b/ceno_zkvm/src/instructions/riscv/branch.rs
@@ -25,22 +25,22 @@ pub struct BltuOp;
 impl RIVInstruction for BltuOp {
     const INST_KIND: InsnKind = InsnKind::BLTU;
 }
-pub type BltuInstruction = bltu::BltuCircuit<BltuOp>;
+pub type BltuInstruction<E> = bltu::BltuCircuit<E, BltuOp>;
 
 pub struct BgeuOp;
 impl RIVInstruction for BgeuOp {
     const INST_KIND: InsnKind = InsnKind::BGEU;
 }
-pub type BgeuInstruction = bltu::BltuCircuit<BgeuOp>;
+pub type BgeuInstruction<E> = bltu::BltuCircuit<E, BgeuOp>;
 
 pub struct BltOp;
 impl RIVInstruction for BltOp {
     const INST_KIND: InsnKind = InsnKind::BLT;
 }
-pub type BltInstruction = blt::BltCircuit<BltOp>;
+pub type BltInstruction<E> = blt::BltCircuit<E, BltOp>;
 
 pub struct BgeOp;
 impl RIVInstruction for BgeOp {
     const INST_KIND: InsnKind = InsnKind::BGE;
 }
-pub type BgeInstruction = blt::BltCircuit<BgeOp>;
+pub type BgeInstruction<E> = blt::BltCircuit<E, BgeOp>;
diff --git a/ceno_zkvm/src/instructions/riscv/branch/blt.rs b/ceno_zkvm/src/instructions/riscv/branch/blt.rs
index 43f07bbfb..47caafc9e 100644
--- a/ceno_zkvm/src/instructions/riscv/branch/blt.rs
+++ b/ceno_zkvm/src/instructions/riscv/branch/blt.rs
@@ -16,7 +16,7 @@ use crate::{
 };
 use ceno_emul::{InsnKind, SWord};
 
-pub struct BltCircuit<I>(PhantomData<I>);
+pub struct BltCircuit<E, I>(PhantomData<(E, I)>);
 
 pub struct InstructionConfig<E: ExtensionField> {
     pub b_insn: BInstructionConfig<E>,
@@ -25,7 +25,7 @@ pub struct InstructionConfig<E: ExtensionField> {
     pub signed_lt: SignedLtConfig,
 }
 
-impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for BltCircuit<I> {
+impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for BltCircuit<E, I> {
     fn name() -> String {
         format!("{:?}", I::INST_KIND)
     }
diff --git a/ceno_zkvm/src/instructions/riscv/branch/bltu.rs b/ceno_zkvm/src/instructions/riscv/branch/bltu.rs
index 3043e7563..896bf19da 100644
--- a/ceno_zkvm/src/instructions/riscv/branch/bltu.rs
+++ b/ceno_zkvm/src/instructions/riscv/branch/bltu.rs
@@ -20,7 +20,7 @@ use crate::{
 };
 use ceno_emul::InsnKind;
 
-pub struct BltuCircuit<I>(PhantomData<I>);
+pub struct BltuCircuit<E, I>(PhantomData<(E, I)>);
 
 pub struct InstructionConfig<E: ExtensionField> {
     pub b_insn: BInstructionConfig<E>,
@@ -29,7 +29,7 @@ pub struct InstructionConfig<E: ExtensionField> {
     pub is_lt: IsLtConfig,
 }
 
-impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for BltuCircuit<I> {
+impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for BltuCircuit<E, I> {
     fn name() -> String {
         format!("{:?}", I::INST_KIND)
     }
diff --git a/ceno_zkvm/src/instructions/riscv/jump/jalr.rs b/ceno_zkvm/src/instructions/riscv/jump/jalr.rs
index a80e11c4a..9795f324b 100644
--- a/ceno_zkvm/src/instructions/riscv/jump/jalr.rs
+++ b/ceno_zkvm/src/instructions/riscv/jump/jalr.rs
@@ -109,6 +109,10 @@ impl<E: ExtensionField> Instruction<E> for JalrInstruction<E> {
 
         let rs1 = step.rs1().unwrap().value;
         let imm: i32 = insn.imm_or_funct7() as i32;
+        if step.rd().is_none() {
+            tracing::info!("step: {:?}", step.insn().kind());
+            tracing::info!("step: {:?}", step);
+        }
         let rd = step.rd().unwrap().value.after;
 
         let (sum, overflowing) = rs1.overflowing_add_signed(imm);
diff --git a/ceno_zkvm/src/instructions/riscv/memory/store.rs b/ceno_zkvm/src/instructions/riscv/memory/store.rs
index 878777de6..8d1fcfc13 100644
--- a/ceno_zkvm/src/instructions/riscv/memory/store.rs
+++ b/ceno_zkvm/src/instructions/riscv/memory/store.rs
@@ -38,7 +38,6 @@ impl RIVInstruction for SWOp {
     const INST_KIND: InsnKind = InsnKind::SW;
 }
 
-#[cfg(test)]
 pub type SwInstruction<E> = StoreInstruction<E, SWOp, 2>;
 
 pub struct SHOp;
@@ -47,7 +46,6 @@ impl RIVInstruction for SHOp {
     const INST_KIND: InsnKind = InsnKind::SH;
 }
 
-#[cfg(test)]
 pub type ShInstruction<E> = StoreInstruction<E, SHOp, 1>;
 
 pub struct SBOp;
@@ -56,7 +54,6 @@ impl RIVInstruction for SBOp {
     const INST_KIND: InsnKind = InsnKind::SB;
 }
 
-#[cfg(test)]
 pub type SbInstruction<E> = StoreInstruction<E, SBOp, 0>;
 
 impl<E: ExtensionField, I: RIVInstruction, const N_ZEROS: usize> Instruction<E>
diff --git a/ceno_zkvm/src/instructions/riscv/rv32im.rs b/ceno_zkvm/src/instructions/riscv/rv32im.rs
index 14a65ec35..ffdfddc61 100644
--- a/ceno_zkvm/src/instructions/riscv/rv32im.rs
+++ b/ceno_zkvm/src/instructions/riscv/rv32im.rs
@@ -1,6 +1,23 @@
 use crate::{
     error::ZKVMError,
-    instructions::{Instruction, riscv::*},
+    instructions::{
+        Instruction,
+        riscv::{
+            arith_imm::AddiInstruction,
+            branch::{
+                BeqInstruction, BgeInstruction, BgeuInstruction, BltInstruction, BneInstruction,
+            },
+            divu::DivUInstruction,
+            logic::{AndInstruction, OrInstruction, XorInstruction},
+            logic_imm::{AndiInstruction, OriInstruction, XoriInstruction},
+            mulh::MulhuInstruction,
+            shift::{SllInstruction, SrlInstruction},
+            shift_imm::{SlliInstruction, SraiInstruction, SrliInstruction},
+            slti::SltiInstruction,
+            sltu::SltuInstruction,
+            *,
+        },
+    },
     structs::{ZKVMConstraintSystem, ZKVMFixedTraces, ZKVMWitnesses},
     tables::{
         AndTableCircuit, LtuTableCircuit, MemFinalRecord, MemInitRecord, MemTableCircuit,
@@ -11,6 +28,7 @@ use ceno_emul::{CENO_PLATFORM, InsnKind, InsnKind::*, StepRecord};
 use ff_ext::ExtensionField;
 use itertools::Itertools;
 use num_traits::cast::ToPrimitive;
+use std::collections::BTreeMap;
 use strum::IntoEnumIterator;
 
 use super::{
@@ -25,17 +43,40 @@ pub struct Rv32imConfig<E: ExtensionField> {
     // ALU Opcodes.
     pub add_config: <AddInstruction<E> as Instruction<E>>::InstructionConfig,
     pub sub_config: <SubInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub and_config: <AndInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub or_config: <OrInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub xor_config: <XorInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub sll_config: <SllInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub srl_config: <SrlInstruction<E> as Instruction<E>>::InstructionConfig,
+    // TODO: sra / slt
+    pub sltu_config: <SltuInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub mul_config: <MulInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub mulhu_config: <MulhuInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub divu_config: <DivUInstruction<E> as Instruction<E>>::InstructionConfig,
 
-    // Branching Opcodes
-    pub bltu_config: <BltuInstruction as Instruction<E>>::InstructionConfig,
+    // ALU with imm
+    pub addi_config: <AddiInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub andi_config: <AndiInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub ori_config: <OriInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub xori_config: <XoriInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub slli_config: <SlliInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub srli_config: <SrliInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub srai_config: <SraiInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub slti_config: <SltiInstruction<E> as Instruction<E>>::InstructionConfig,
 
-    // Imm
-    pub lui_config: <LuiInstruction<E> as Instruction<E>>::InstructionConfig,
+    // Branching Opcodes
+    pub beq_config: <BeqInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub bne_config: <BneInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub blt_config: <BltInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub bltu_config: <BltuInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub bge_config: <BgeInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub bgeu_config: <BgeuInstruction<E> as Instruction<E>>::InstructionConfig,
 
     // Jump Opcodes
     pub jal_config: <JalInstruction<E> as Instruction<E>>::InstructionConfig,
     pub jalr_config: <JalrInstruction<E> as Instruction<E>>::InstructionConfig,
     pub auipc_config: <AuipcInstruction<E> as Instruction<E>>::InstructionConfig,
+    pub lui_config: <LuiInstruction<E> as Instruction<E>>::InstructionConfig,
 
     // Memory Opcodes
     pub lw_config: <LwInstruction<E> as Instruction<E>>::InstructionConfig,
@@ -52,7 +93,7 @@ pub struct Rv32imConfig<E: ExtensionField> {
     // Tables.
     pub u16_range_config: <U16TableCircuit<E> as TableCircuit<E>>::TableConfig,
     pub u14_range_config: <U14TableCircuit<E> as TableCircuit<E>>::TableConfig,
-    pub and_config: <AndTableCircuit<E> as TableCircuit<E>>::TableConfig,
+    pub and_table_config: <AndTableCircuit<E> as TableCircuit<E>>::TableConfig,
     pub ltu_config: <LtuTableCircuit<E> as TableCircuit<E>>::TableConfig,
 
     // RW tables.
@@ -66,9 +107,33 @@ impl<E: ExtensionField> Rv32imConfig<E> {
         // alu opcodes
         let add_config = cs.register_opcode_circuit::<AddInstruction<E>>();
         let sub_config = cs.register_opcode_circuit::<SubInstruction<E>>();
+        let and_config = cs.register_opcode_circuit::<AndInstruction<E>>();
+        let or_config = cs.register_opcode_circuit::<OrInstruction<E>>();
+        let xor_config = cs.register_opcode_circuit::<XorInstruction<E>>();
+        let sll_config = cs.register_opcode_circuit::<SllInstruction<E>>();
+        let srl_config = cs.register_opcode_circuit::<SrlInstruction<E>>();
+        let sltu_config = cs.register_opcode_circuit::<SltuInstruction<E>>();
+        let mul_config = cs.register_opcode_circuit::<MulInstruction<E>>();
+        let mulhu_config = cs.register_opcode_circuit::<MulhuInstruction<E>>();
+        let divu_config = cs.register_opcode_circuit::<DivUInstruction<E>>();
+
+        // alu with imm opcodes
+        let addi_config = cs.register_opcode_circuit::<AddiInstruction<E>>();
+        let andi_config = cs.register_opcode_circuit::<AndiInstruction<E>>();
+        let ori_config = cs.register_opcode_circuit::<OriInstruction<E>>();
+        let xori_config = cs.register_opcode_circuit::<XoriInstruction<E>>();
+        let slli_config = cs.register_opcode_circuit::<SlliInstruction<E>>();
+        let srli_config = cs.register_opcode_circuit::<SrliInstruction<E>>();
+        let srai_config = cs.register_opcode_circuit::<SraiInstruction<E>>();
+        let slti_config = cs.register_opcode_circuit::<SltiInstruction<E>>();
 
         // branching opcodes
-        let bltu_config = cs.register_opcode_circuit::<BltuInstruction>();
+        let beq_config = cs.register_opcode_circuit::<BeqInstruction<E>>();
+        let bne_config = cs.register_opcode_circuit::<BneInstruction<E>>();
+        let blt_config = cs.register_opcode_circuit::<BltInstruction<E>>();
+        let bltu_config = cs.register_opcode_circuit::<BltuInstruction<E>>();
+        let bge_config = cs.register_opcode_circuit::<BgeInstruction<E>>();
+        let bgeu_config = cs.register_opcode_circuit::<BgeuInstruction<E>>();
 
         // jump opcodes
         let lui_config = cs.register_opcode_circuit::<LuiInstruction<E>>();
@@ -91,7 +156,7 @@ impl<E: ExtensionField> Rv32imConfig<E> {
         // tables
         let u16_range_config = cs.register_table_circuit::<U16TableCircuit<E>>();
         let u14_range_config = cs.register_table_circuit::<U14TableCircuit<E>>();
-        let and_config = cs.register_table_circuit::<AndTableCircuit<E>>();
+        let and_table_config = cs.register_table_circuit::<AndTableCircuit<E>>();
         let ltu_config = cs.register_table_circuit::<LtuTableCircuit<E>>();
 
         // RW tables
@@ -102,8 +167,31 @@ impl<E: ExtensionField> Rv32imConfig<E> {
             // alu opcodes
             add_config,
             sub_config,
+            and_config,
+            or_config,
+            xor_config,
+            sll_config,
+            srl_config,
+            sltu_config,
+            mul_config,
+            mulhu_config,
+            divu_config,
+            // alu with imm
+            addi_config,
+            andi_config,
+            ori_config,
+            xori_config,
+            slli_config,
+            srli_config,
+            srai_config,
+            slti_config,
             // branching opcodes
+            beq_config,
+            bne_config,
+            blt_config,
             bltu_config,
+            bge_config,
+            bgeu_config,
             // jump opcodes
             lui_config,
             jal_config,
@@ -123,7 +211,7 @@ impl<E: ExtensionField> Rv32imConfig<E> {
             // tables
             u16_range_config,
             u14_range_config,
-            and_config,
+            and_table_config,
             ltu_config,
 
             reg_config,
@@ -138,16 +226,41 @@ impl<E: ExtensionField> Rv32imConfig<E> {
         reg_init: &[MemInitRecord],
         mem_init: &[MemInitRecord],
     ) {
+        // alu
         fixed.register_opcode_circuit::<AddInstruction<E>>(cs);
         fixed.register_opcode_circuit::<SubInstruction<E>>(cs);
-
-        fixed.register_opcode_circuit::<BltuInstruction>(cs);
-
+        fixed.register_opcode_circuit::<AndInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<OrInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<XorInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<SllInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<SrlInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<SltuInstruction<E>>(cs);
+        // TODO: add sra / slt
+        fixed.register_opcode_circuit::<MulInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<MulhuInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<DivUInstruction<E>>(cs);
+        // alu with imm
+        fixed.register_opcode_circuit::<AddiInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<AndiInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<OriInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<XoriInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<SlliInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<SrliInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<SraiInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<SltiInstruction<E>>(cs);
+        // branching
+        fixed.register_opcode_circuit::<BeqInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<BneInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<BltInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<BltuInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<BgeInstruction<E>>(cs);
+        fixed.register_opcode_circuit::<BgeuInstruction<E>>(cs);
+        // jump
         fixed.register_opcode_circuit::<JalInstruction<E>>(cs);
         fixed.register_opcode_circuit::<JalrInstruction<E>>(cs);
         fixed.register_opcode_circuit::<AuipcInstruction<E>>(cs);
         fixed.register_opcode_circuit::<LuiInstruction<E>>(cs);
-
+        // memory
         fixed.register_opcode_circuit::<SwInstruction<E>>(cs);
         fixed.register_opcode_circuit::<ShInstruction<E>>(cs);
         fixed.register_opcode_circuit::<SbInstruction<E>>(cs);
@@ -161,20 +274,23 @@ impl<E: ExtensionField> Rv32imConfig<E> {
 
         fixed.register_table_circuit::<U16TableCircuit<E>>(cs, self.u16_range_config.clone(), &());
         fixed.register_table_circuit::<U14TableCircuit<E>>(cs, self.u14_range_config.clone(), &());
-        fixed.register_table_circuit::<AndTableCircuit<E>>(cs, self.and_config.clone(), &());
+        fixed.register_table_circuit::<AndTableCircuit<E>>(cs, self.and_table_config.clone(), &());
         fixed.register_table_circuit::<LtuTableCircuit<E>>(cs, self.ltu_config.clone(), &());
 
         fixed.register_table_circuit::<RegTableCircuit<E>>(cs, self.reg_config.clone(), reg_init);
         fixed.register_table_circuit::<MemTableCircuit<E>>(cs, self.mem_config.clone(), mem_init);
     }
 
+
     pub fn assign_opcode_circuit(
         &self,
         cs: &ZKVMConstraintSystem<E>,
         witness: &mut ZKVMWitnesses<E>,
         steps: Vec<StepRecord>,
     ) -> Result<(), ZKVMError> {
-        let mut all_records = vec![Vec::new(); InsnKind::iter().count()];
+        let mut all_records: BTreeMap<usize, Vec<StepRecord>> = InsnKind::iter()
+            .map(|insn_kind| (insn_kind.to_usize().unwrap(), Vec::new()))
+            .collect();
         let mut halt_records = Vec::new();
         steps.into_iter().for_each(|record| {
             let insn_kind = record.insn().codes().kind;
@@ -183,56 +299,81 @@ impl<E: ExtensionField> Rv32imConfig<E> {
                 EANY if record.rs1().unwrap().value == CENO_PLATFORM.ecall_halt() => {
                     halt_records.push(record);
                 }
-                _ => all_records[insn_kind.to_usize().unwrap()].push(record),
+                _ => {
+                    let insn_kind = insn_kind.to_usize().unwrap();
+                    // it's safe to unwrap as all_records are initialized with Vec::new()
+                    all_records.get_mut(&insn_kind).unwrap().push(record);
+                }
             }
         });
 
-        for (insn_kind, records) in InsnKind::iter()
+        for (insn_kind, (_, records)) in InsnKind::iter()
             .zip(all_records.iter())
-            .sorted_by(|a, b| Ord::cmp(&a.1.len(), &b.1.len()))
+            .sorted_by(|a, b| Ord::cmp(&a.1.1.len(), &b.1.1.len()))
             .rev()
         {
-            if records.len() != 0 {
+            if !records.is_empty() {
                 tracing::info!("tracer generated {:?} {} records", insn_kind, records.len());
             }
         }
         assert_eq!(halt_records.len(), 1);
 
-        witness.assign_opcode_circuit::<AddInstruction<E>>(
-            cs,
-            &self.add_config,
-            all_records[ADD.to_usize().unwrap()].as_slice(),
-        )?;
-        witness.assign_opcode_circuit::<BltuInstruction>(
-            cs,
-            &self.bltu_config,
-            all_records[BLTU.to_usize().unwrap()].as_slice(),
-        )?;
-        witness.assign_opcode_circuit::<JalInstruction<E>>(
-            cs,
-            &self.jal_config,
-            all_records[JAL.to_usize().unwrap()].as_slice(),
-        )?;
-        witness.assign_opcode_circuit::<LuiInstruction<E>>(
-            cs,
-            &self.lui_config,
-            all_records[LUI.to_usize().unwrap()].as_slice(),
-        )?;
-        witness.assign_opcode_circuit::<LwInstruction<E>>(
-            cs,
-            &self.lw_config,
-            all_records[LW.to_usize().unwrap()].as_slice(),
-        )?;
-        witness.assign_opcode_circuit::<SwInstruction<E>>(
-            cs,
-            &self.sw_config,
-            all_records[SW.to_usize().unwrap()].as_slice(),
-        )?;
-        witness.assign_opcode_circuit::<HaltInstruction<E>>(
-            cs,
-            &self.halt_config,
-            &halt_records,
-        )?;
+        macro_rules! assign_opcode {
+            ($insn_kind:ident,$instruction:ty,$config:ident) => {
+                witness.assign_opcode_circuit::<$instruction>(
+                    cs,
+                    &self.$config,
+                    all_records.remove(&$insn_kind.to_usize().unwrap()).unwrap(),
+                )?;
+            }
+        }
+        // alu
+        assign_opcode!(ADD, AddInstruction<E>, add_config);
+        assign_opcode!(SUB, SubInstruction<E>, sub_config);
+        assign_opcode!(AND, AndInstruction<E>, and_config);
+        assign_opcode!(OR, OrInstruction<E>, or_config);
+        assign_opcode!(XOR, XorInstruction<E>, xor_config);
+        assign_opcode!(SLL, SllInstruction<E>, sll_config);
+        assign_opcode!(SRL, SrlInstruction<E>, srl_config);
+        assign_opcode!(SLTU, SltuInstruction<E>, sltu_config);
+        assign_opcode!(MUL, MulInstruction<E>, mul_config);
+        assign_opcode!(MULHU, MulhuInstruction<E>, mulhu_config);
+        assign_opcode!(DIVU, DivUInstruction<E>, divu_config);
+        // alu with imm
+        assign_opcode!(ADDI, AddiInstruction<E>, addi_config);
+        assign_opcode!(ANDI, AndiInstruction<E>, andi_config);
+        assign_opcode!(ORI, OriInstruction<E>, ori_config);
+        assign_opcode!(XORI, XoriInstruction<E>, xori_config);
+        assign_opcode!(SLLI, SlliInstruction<E>, slli_config);
+        assign_opcode!(SRLI, SrliInstruction<E>, srli_config);
+        assign_opcode!(SRAI, SraiInstruction<E>, srai_config);
+        assign_opcode!(SLTI, SltiInstruction<E>, slti_config);
+        // branching
+        assign_opcode!(BEQ, BeqInstruction<E>, beq_config);
+        assign_opcode!(BNE, BneInstruction<E>, bne_config);
+        assign_opcode!(BLT, BltInstruction<E>, blt_config);
+        assign_opcode!(BLTU, BltuInstruction<E>, bltu_config);
+        assign_opcode!(BGE, BgeInstruction<E>, bge_config);
+        assign_opcode!(BGEU, BgeuInstruction<E>, bgeu_config);
+        // jump
+        assign_opcode!(JAL, JalInstruction<E>, jal_config);
+        assign_opcode!(JALR, JalrInstruction<E>, jalr_config);
+        assign_opcode!(AUIPC, AuipcInstruction<E>, auipc_config);
+        assign_opcode!(LUI, LuiInstruction<E>, lui_config);
+        // memory
+        assign_opcode!(LW, LwInstruction<E>, lw_config);
+        assign_opcode!(LB, LbInstruction<E>, lb_config);
+        assign_opcode!(LBU, LbuInstruction<E>, lbu_config);
+        assign_opcode!(LH, LhInstruction<E>, lh_config);
+        assign_opcode!(LHU, LhuInstruction<E>, lhu_config);
+        assign_opcode!(SW, SwInstruction<E>, sw_config);
+        assign_opcode!(SH, ShInstruction<E>, sh_config);
+        assign_opcode!(SB, SbInstruction<E>, sb_config);
+
+        // ecall / halt
+        witness.assign_opcode_circuit::<HaltInstruction<E>>(cs, &self.halt_config, halt_records)?;
+
+        assert!(all_records.is_empty());
         Ok(())
     }
 
@@ -245,7 +386,7 @@ impl<E: ExtensionField> Rv32imConfig<E> {
     ) -> Result<(), ZKVMError> {
         witness.assign_table_circuit::<U16TableCircuit<E>>(cs, &self.u16_range_config, &())?;
         witness.assign_table_circuit::<U14TableCircuit<E>>(cs, &self.u14_range_config, &())?;
-        witness.assign_table_circuit::<AndTableCircuit<E>>(cs, &self.and_config, &())?;
+        witness.assign_table_circuit::<AndTableCircuit<E>>(cs, &self.and_table_config, &())?;
         witness.assign_table_circuit::<LtuTableCircuit<E>>(cs, &self.ltu_config, &())?;
 
         // assign register finalization.
diff --git a/ceno_zkvm/src/instructions/riscv/shift.rs b/ceno_zkvm/src/instructions/riscv/shift.rs
index f5d0f8e8d..a529b9e3c 100644
--- a/ceno_zkvm/src/instructions/riscv/shift.rs
+++ b/ceno_zkvm/src/instructions/riscv/shift.rs
@@ -31,19 +31,17 @@ pub struct ShiftConfig<E: ExtensionField> {
 
 pub struct ShiftLogicalInstruction<E, I>(PhantomData<(E, I)>);
 
-#[cfg(test)]
-struct SllOp;
-#[cfg(test)]
+pub struct SllOp;
 impl RIVInstruction for SllOp {
     const INST_KIND: InsnKind = InsnKind::SLL;
 }
+pub type SllInstruction<E> = ShiftLogicalInstruction<E, SllOp>;
 
-#[cfg(test)]
-struct SrlOp;
-#[cfg(test)]
+pub struct SrlOp;
 impl RIVInstruction for SrlOp {
     const INST_KIND: InsnKind = InsnKind::SRL;
 }
+pub type SrlInstruction<E> = ShiftLogicalInstruction<E, SrlOp>;
 
 impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for ShiftLogicalInstruction<E, I> {
     type InstructionConfig = ShiftConfig<E>;
diff --git a/ceno_zkvm/src/instructions/riscv/shift_imm.rs b/ceno_zkvm/src/instructions/riscv/shift_imm.rs
index 1f43da85c..04e141cec 100644
--- a/ceno_zkvm/src/instructions/riscv/shift_imm.rs
+++ b/ceno_zkvm/src/instructions/riscv/shift_imm.rs
@@ -35,16 +35,19 @@ pub struct SlliOp;
 impl RIVInstruction for SlliOp {
     const INST_KIND: ceno_emul::InsnKind = ceno_emul::InsnKind::SLLI;
 }
+pub type SlliInstruction<E> = ShiftImmInstruction<E, SlliOp>;
 
 pub struct SraiOp;
 impl RIVInstruction for SraiOp {
     const INST_KIND: ceno_emul::InsnKind = ceno_emul::InsnKind::SRAI;
 }
+pub type SraiInstruction<E> = ShiftImmInstruction<E, SraiOp>;
 
 pub struct SrliOp;
 impl RIVInstruction for SrliOp {
     const INST_KIND: ceno_emul::InsnKind = InsnKind::SRLI;
 }
+pub type SrliInstruction<E> = ShiftImmInstruction<E, SrliOp>;
 
 impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for ShiftImmInstruction<E, I> {
     type InstructionConfig = ShiftImmConfig<E>;
diff --git a/ceno_zkvm/src/scheme/prover.rs b/ceno_zkvm/src/scheme/prover.rs
index c87283c01..b83ca4689 100644
--- a/ceno_zkvm/src/scheme/prover.rs
+++ b/ceno_zkvm/src/scheme/prover.rs
@@ -76,17 +76,23 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
         for (circuit_name, witness) in witnesses.witnesses {
             let commit_dur = std::time::Instant::now();
             let num_instances = witness.num_instances();
-            let witness = witness.into_mles();
-            commitments.insert(
-                circuit_name.clone(),
-                PCS::batch_commit_and_write(&self.pk.pp, &witness, &mut transcript)
-                    .map_err(ZKVMError::PCSError)?,
-            );
-            tracing::info!(
-                "commit to {} traces took {:?}",
-                circuit_name,
-                commit_dur.elapsed()
-            );
+            let witness = match num_instances {
+                0 => vec![],
+                _ => {
+                    let witness = witness.into_mles();
+                    commitments.insert(
+                        circuit_name.clone(),
+                        PCS::batch_commit_and_write(&self.pk.pp, &witness, &mut transcript)
+                            .map_err(ZKVMError::PCSError)?,
+                    );
+                    tracing::info!(
+                        "commit to {} traces took {:?}",
+                        circuit_name,
+                        commit_dur.elapsed()
+                    );
+                    witness
+                }
+            };
             wits.insert(circuit_name, (witness, num_instances));
         }
 
@@ -107,6 +113,9 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
             let (witness, num_instances) = wits
                 .remove(circuit_name)
                 .ok_or(ZKVMError::WitnessNotFound(circuit_name.clone()))?;
+            if witness.is_empty() {
+                continue;
+            }
             let wits_commit = commitments.remove(circuit_name).unwrap();
             // TODO: add an enum for circuit type either in constraint_system or vk
             let cs = pk.get_cs();
diff --git a/ceno_zkvm/src/structs.rs b/ceno_zkvm/src/structs.rs
index 5ef37de55..e06c1d6dd 100644
--- a/ceno_zkvm/src/structs.rs
+++ b/ceno_zkvm/src/structs.rs
@@ -215,16 +215,18 @@ impl<E: ExtensionField> ZKVMWitnesses<E> {
         &mut self,
         cs: &ZKVMConstraintSystem<E>,
         config: &OC::InstructionConfig,
-        records: &[StepRecord],
+        records: Vec<StepRecord>,
     ) -> Result<(), ZKVMError> {
         assert!(self.combined_lk_mlt.is_none());
 
-        if records.len() == 0 {
-            return Ok(());
-        }
         let cs = cs.get_cs(&OC::name()).unwrap();
-        let (witness, logup_multiplicity) =
-            OC::assign_instances(config, cs.num_witin as usize, records.to_vec())?;
+        let (witness, logup_multiplicity) = match records.len() {
+            0 => (
+                RowMajorMatrix::new(0, cs.num_witin as usize),
+                LkMultiplicity::default(),
+            ),
+            _ => OC::assign_instances(config, cs.num_witin as usize, records)?,
+        };
         assert!(self.witnesses.insert(OC::name(), witness).is_none());
         assert!(
             self.lk_mlts