Refactor RISC-V instruction types

Implement library of common gadgets for setting up register access and vm
state for RISC-V base instruction types.

ceno_zkvm/src/instructions/riscv.rs

@@ -6,13 +6,14 @@ pub mod branch;
 pub mod config;
 pub mod constants;
 pub mod divu;
-mod i_insn;
 pub mod logic;
+pub mod shift_imm;
 pub mod sltu;
 
 mod b_insn;
+mod i_insn;
+mod insn_base;
 mod r_insn;
-pub mod shift_imm;
 
 #[cfg(test)]
 mod test;

ceno_zkvm/src/instructions/riscv/b_insn.rs

@@ -5,13 +5,11 @@ use ff_ext::ExtensionField;
 
 use super::constants::PC_STEP_SIZE;
 use crate::{
-    chip_handler::{
-        GlobalStateRegisterMachineChipOperations, RegisterChipOperations, RegisterExpr,
-    },
+    chip_handler::RegisterExpr,
     circuit_builder::CircuitBuilder,
     error::ZKVMError,
     expression::{Expression, ToExpr, WitIn},
-    gadgets::IsLtConfig,
+    instructions::riscv::insn_base::{ReadRS1, ReadRS2, StateInOut},
     set_val,
     tables::InsnRecord,
     witness::LkMultiplicity,
@@ -36,108 +34,72 @@ use core::mem::MaybeUninit;
 /// It does _not_ range-check the `branch_taken_bit`.
 /// It does not witness of the register values, nor the actual function (e.g. BNE).
 #[derive(Debug)]
-pub struct BInstructionConfig {
-    pc: WitIn,
-    next_pc: WitIn,
-    ts: WitIn,
-    rs1_id: WitIn,
-    rs2_id: WitIn,
-    imm: WitIn,
-    prev_rs1_ts: WitIn,
-    prev_rs2_ts: WitIn,
-    lt_rs1_cfg: IsLtConfig,
-    lt_rs2_cfg: IsLtConfig,
+pub struct BInstructionConfig<E: ExtensionField> {
+    pub vm_state: StateInOut<E>,
+    pub rs1: ReadRS1<E>,
+    pub rs2: ReadRS2<E>,
+    pub imm: WitIn,
 }
 
-impl BInstructionConfig {
-    pub fn construct_circuit<E: ExtensionField>(
+impl<E: ExtensionField> BInstructionConfig<E> {
+    pub fn construct_circuit(
         circuit_builder: &mut CircuitBuilder<E>,
         insn_kind: InsnKind,
         rs1_read: RegisterExpr<E>,
         rs2_read: RegisterExpr<E>,
         branch_taken_bit: Expression<E>,
     ) -> Result<Self, ZKVMError> {
-        // State in.
-        let pc = circuit_builder.create_witin(|| "pc")?;
-        let cur_ts = circuit_builder.create_witin(|| "cur_ts")?;
-        circuit_builder.state_in(pc.expr(), cur_ts.expr())?;
+        // State in and out
+        let vm_state = StateInOut::construct_circuit(circuit_builder, true)?;
 
-        // Register indexes and immediate.
-        let rs1_id = circuit_builder.create_witin(|| "rs1_id")?;
-        let rs2_id = circuit_builder.create_witin(|| "rs2_id")?;
+        // Registers
+        let rs1 = ReadRS1::construct_circuit(circuit_builder, rs1_read, vm_state.ts)?;
+        let rs2 = ReadRS2::construct_circuit(circuit_builder, rs2_read, vm_state.ts)?;
+
+        // Immediate
         let imm = circuit_builder.create_witin(|| "imm")?;
 
-        // Fetch the instruction.
+        // Fetch instruction
         circuit_builder.lk_fetch(&InsnRecord::new(
-            pc.expr(),
+            vm_state.pc.expr(),
             (insn_kind.codes().opcode as usize).into(),
             0.into(),
             (insn_kind.codes().func3 as usize).into(),
-            rs1_id.expr(),
-            rs2_id.expr(),
+            rs1.id.expr(),
+            rs2.id.expr(),
             imm.expr(),
         ))?;
 
-        // Register state.
-        let prev_rs1_ts = circuit_builder.create_witin(|| "prev_rs1_ts")?;
-        let prev_rs2_ts = circuit_builder.create_witin(|| "prev_rs2_ts")?;
-
-        // Register reads.
-        let (ts, lt_rs1_cfg) = circuit_builder.register_read(
-            || "read_rs1",
-            &rs1_id,
-            prev_rs1_ts.expr(),
-            cur_ts.expr(),
-            rs1_read,
+        // Branch program counter
+        let pc_offset = branch_taken_bit.clone() * imm.expr()
+            - branch_taken_bit * PC_STEP_SIZE.into()
+            + PC_STEP_SIZE.into();
+        let next_pc = vm_state.next_pc.unwrap();
+        circuit_builder.require_equal(
+            || "pc_branch",
+            next_pc.expr(),
+            vm_state.pc.expr() + pc_offset,
         )?;
-        let (_ts, lt_rs2_cfg) = circuit_builder.register_read(
-            || "read_rs2",
-            &rs2_id,
-            prev_rs2_ts.expr(),
-            ts,
-            rs2_read,
-        )?;
-
-        // State out.
-        let next_pc = {
-            let pc_offset = branch_taken_bit.clone() * imm.expr()
-                - branch_taken_bit * PC_STEP_SIZE.into()
-                + PC_STEP_SIZE.into();
-            let next_pc = circuit_builder.create_witin(|| "next_pc")?;
-            circuit_builder.require_equal(|| "pc_branch", next_pc.expr(), pc.expr() + pc_offset)?;
-            next_pc
-        };
-        let next_ts = cur_ts.expr() + 4.into();
-        circuit_builder.state_out(next_pc.expr(), next_ts)?;
 
         Ok(BInstructionConfig {
-            pc,
-            next_pc,
-            ts: cur_ts,
-            rs1_id,
-            rs2_id,
+            vm_state,
+            rs1,
+            rs2,
             imm,
-            prev_rs1_ts,
-            prev_rs2_ts,
-            lt_rs1_cfg,
-            lt_rs2_cfg,
         })
     }
 
-    pub fn assign_instance<E: ExtensionField>(
+    pub fn assign_instance(
         &self,
         instance: &mut [MaybeUninit<<E as ExtensionField>::BaseField>],
         lk_multiplicity: &mut LkMultiplicity,
         step: &StepRecord,
     ) -> Result<(), ZKVMError> {
-        // State.
-        set_val!(instance, self.pc, step.pc().before.0 as u64);
-        set_val!(instance, self.next_pc, step.pc().after.0 as u64);
-        set_val!(instance, self.ts, step.cycle());
+        self.vm_state.assign_instance(instance, step)?;
+        self.rs1.assign_instance(instance, lk_multiplicity, step)?;
+        self.rs2.assign_instance(instance, lk_multiplicity, step)?;
 
-        // Register indexes and immediate.
-        set_val!(instance, self.rs1_id, step.insn().rs1() as u64);
-        set_val!(instance, self.rs2_id, step.insn().rs2() as u64);
+        // Immediate
         set_val!(
             instance,
             self.imm,
@@ -147,32 +109,6 @@ impl BInstructionConfig {
         // Fetch the instruction.
         lk_multiplicity.fetch(step.pc().before.0);
 
-        // Register state.
-        set_val!(
-            instance,
-            self.prev_rs1_ts,
-            step.rs1().unwrap().previous_cycle
-        );
-        set_val!(
-            instance,
-            self.prev_rs2_ts,
-            step.rs2().unwrap().previous_cycle
-        );
-
-        // Register read and write.
-        self.lt_rs1_cfg.assign_instance(
-            instance,
-            lk_multiplicity,
-            step.rs1().unwrap().previous_cycle,
-            step.cycle(),
-        )?;
-        self.lt_rs2_cfg.assign_instance(
-            instance,
-            lk_multiplicity,
-            step.rs2().unwrap().previous_cycle,
-            step.cycle() + 1,
-        )?;
-
         Ok(())
     }
 }

ceno_zkvm/src/instructions/riscv/branch/beq_circuit.rs

@@ -17,7 +17,7 @@ use crate::{
 };
 
 pub struct BeqConfig<E: ExtensionField> {
-    b_insn: BInstructionConfig,
+    b_insn: BInstructionConfig<E>,
 
     // TODO: Limb decomposition is not necessary. Replace with a single witness.
     rs1_read: UInt<E>,
@@ -78,7 +78,7 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for BeqCircuit<E, I> {
     ) -> Result<(), ZKVMError> {
         config
             .b_insn
-            .assign_instance::<E>(instance, lk_multiplicity, step)?;
+            .assign_instance(instance, lk_multiplicity, step)?;
 
         let rs1_read = step.rs1().unwrap().value;
         config

ceno_zkvm/src/instructions/riscv/branch/blt.rs

@@ -20,7 +20,7 @@ use ceno_emul::InsnKind;
 pub struct BltCircuit<I>(PhantomData<I>);
 
 pub struct InstructionConfig<E: ExtensionField> {
-    pub b_insn: BInstructionConfig,
+    pub b_insn: BInstructionConfig<E>,
     pub read_rs1: UInt<E>,
     pub read_rs2: UInt<E>,
     pub is_lt: UIntLtSignedConfig,
@@ -90,7 +90,7 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for BltCircuit<I> {
 
         config
             .b_insn
-            .assign_instance::<E>(instance, lk_multiplicity, step)?;
+            .assign_instance(instance, lk_multiplicity, step)?;
 
         Ok(())
     }

ceno_zkvm/src/instructions/riscv/branch/bltu.rs

@@ -23,7 +23,7 @@ use ceno_emul::InsnKind;
 pub struct BltuCircuit<I>(PhantomData<I>);
 
 pub struct InstructionConfig<E: ExtensionField> {
-    pub b_insn: BInstructionConfig,
+    pub b_insn: BInstructionConfig<E>,
     pub read_rs1: UInt<E>,
     pub read_rs2: UInt<E>,
     pub is_lt: IsLtConfig,
@@ -93,7 +93,7 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for BltuCircuit<I> {
 
         config
             .b_insn
-            .assign_instance::<E>(instance, lk_multiplicity, step)?;
+            .assign_instance(instance, lk_multiplicity, step)?;
 
         Ok(())
     }