From 7ed82fc658545c2ecc2fa369ec67dca0c0e5ce14 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Nicolas?= <info@nau.re>
Date: Wed, 11 Sep 2024 15:34:23 +0200
Subject: [PATCH 01/12] program-table: initial program table

---
 ceno_zkvm/.gitignore            |   1 +
 ceno_zkvm/examples/riscv_add.rs |  25 +++++-
 ceno_zkvm/src/structs.rs        |  12 +--
 ceno_zkvm/src/tables/mod.rs     |   6 +-
 ceno_zkvm/src/tables/program.rs | 134 ++++++++++++++++++++++++++++++++
 ceno_zkvm/src/tables/range.rs   |   3 +-
 6 files changed, 170 insertions(+), 11 deletions(-)
 create mode 100644 ceno_zkvm/.gitignore
 create mode 100644 ceno_zkvm/src/tables/program.rs

diff --git a/ceno_zkvm/.gitignore b/ceno_zkvm/.gitignore
new file mode 100644
index 000000000..8db038a03
--- /dev/null
+++ b/ceno_zkvm/.gitignore
@@ -0,0 +1 @@
+tracing.folded
diff --git a/ceno_zkvm/examples/riscv_add.rs b/ceno_zkvm/examples/riscv_add.rs
index ec8732085..e02577166 100644
--- a/ceno_zkvm/examples/riscv_add.rs
+++ b/ceno_zkvm/examples/riscv_add.rs
@@ -1,7 +1,10 @@
 use std::time::Instant;
 
 use ark_std::test_rng;
-use ceno_zkvm::{instructions::riscv::addsub::AddInstruction, scheme::prover::ZKVMProver};
+use ceno_zkvm::{
+    instructions::riscv::addsub::AddInstruction, scheme::prover::ZKVMProver,
+    tables::ProgramTableCircuit,
+};
 use const_env::from_env;
 
 use ceno_emul::{ByteAddr, InsnKind::ADD, StepRecord, VMState, CENO_PLATFORM};
@@ -63,6 +66,7 @@ fn main() {
         .with(
             fmt::layer()
                 .compact()
+                .without_time()
                 .with_thread_ids(false)
                 .with_thread_names(false),
         )
@@ -74,11 +78,20 @@ fn main() {
     let mut zkvm_cs = ZKVMConstraintSystem::default();
     let add_config = zkvm_cs.register_opcode_circuit::<AddInstruction<E>>();
     let range_config = zkvm_cs.register_table_circuit::<RangeTableCircuit<E>>();
+    let prog_config = zkvm_cs.register_table_circuit::<ProgramTableCircuit<E>>();
 
     let mut zkvm_fixed_traces = ZKVMFixedTraces::default();
     zkvm_fixed_traces.register_opcode_circuit::<AddInstruction<E>>(&zkvm_cs);
-    zkvm_fixed_traces
-        .register_table_circuit::<RangeTableCircuit<E>>(&zkvm_cs, range_config.clone());
+    zkvm_fixed_traces.register_table_circuit::<RangeTableCircuit<E>>(
+        &zkvm_cs,
+        range_config.clone(),
+        &(),
+    );
+    zkvm_fixed_traces.register_table_circuit::<ProgramTableCircuit<E>>(
+        &zkvm_cs,
+        prog_config.clone(),
+        &PROGRAM_ADD_LOOP,
+    );
 
     let pk = zkvm_cs
         .clone()
@@ -90,7 +103,8 @@ fn main() {
     let prover = ZKVMProver::new(pk);
     let verifier = ZKVMVerifier::new(vk);
 
-    for instance_num_vars in 8..22 {
+    // for instance_num_vars in 8..22 {
+    for instance_num_vars in 2..3 {
         let num_instances = 1 << instance_num_vars;
         let mut vm = VMState::new(CENO_PLATFORM);
         let pc_start = ByteAddr(CENO_PLATFORM.pc_start()).waddr();
@@ -122,6 +136,9 @@ fn main() {
         zkvm_witness
             .assign_table_circuit::<RangeTableCircuit<E>>(&zkvm_cs, &range_config)
             .unwrap();
+        zkvm_witness
+            .assign_table_circuit::<ProgramTableCircuit<E>>(&zkvm_cs, &prog_config)
+            .unwrap();
 
         let timer = Instant::now();
 
diff --git a/ceno_zkvm/src/structs.rs b/ceno_zkvm/src/structs.rs
index 16af00f7c..639c04ca0 100644
--- a/ceno_zkvm/src/structs.rs
+++ b/ceno_zkvm/src/structs.rs
@@ -40,10 +40,11 @@ pub type WitnessId = u16;
 pub type ChallengeId = u16;
 
 pub enum ROMType {
-    U5 = 0, // 2^5 = 32
-    U16,    // 2^16 = 65,536
-    And,    // a ^ b where a, b are bytes
-    Ltu,    // a <(usign) b where a, b are bytes
+    U5 = 0,      // 2^5 = 32
+    U16,         // 2^16 = 65,536
+    And,         // a ^ b where a, b are bytes
+    Ltu,         // a <(usign) b where a, b are bytes
+    Instruction, // Decoded instruction from the fixed program.
 }
 
 #[derive(Clone, Debug, Copy)]
@@ -154,13 +155,14 @@ impl<E: ExtensionField> ZKVMFixedTraces<E> {
         &mut self,
         cs: &ZKVMConstraintSystem<E>,
         config: TC::TableConfig,
+        input: &TC::Input,
     ) {
         let cs = cs.get_cs(&TC::name()).expect("cs not found");
         assert!(
             self.circuit_fixed_traces
                 .insert(
                     TC::name(),
-                    Some(TC::generate_fixed_traces(&config, cs.num_fixed,)),
+                    Some(TC::generate_fixed_traces(&config, cs.num_fixed, input)),
                 )
                 .is_none()
         );
diff --git a/ceno_zkvm/src/tables/mod.rs b/ceno_zkvm/src/tables/mod.rs
index a66094952..fdd74a4b0 100644
--- a/ceno_zkvm/src/tables/mod.rs
+++ b/ceno_zkvm/src/tables/mod.rs
@@ -5,9 +5,12 @@ use std::collections::HashMap;
 mod range;
 pub use range::RangeTableCircuit;
 
+mod program;
+pub use program::ProgramTableCircuit;
+
 pub trait TableCircuit<E: ExtensionField> {
     type TableConfig: Send + Sync;
-    type Input: Send + Sync;
+    type Input: Send + Sync + ?Sized;
 
     fn name() -> String;
 
@@ -18,6 +21,7 @@ pub trait TableCircuit<E: ExtensionField> {
     fn generate_fixed_traces(
         config: &Self::TableConfig,
         num_fixed: usize,
+        input: &Self::Input,
     ) -> RowMajorMatrix<E::BaseField>;
 
     fn assign_instances(
diff --git a/ceno_zkvm/src/tables/program.rs b/ceno_zkvm/src/tables/program.rs
new file mode 100644
index 000000000..23529bd62
--- /dev/null
+++ b/ceno_zkvm/src/tables/program.rs
@@ -0,0 +1,134 @@
+use std::{collections::HashMap, marker::PhantomData, mem::MaybeUninit};
+
+use crate::{
+    circuit_builder::CircuitBuilder,
+    error::ZKVMError,
+    expression::{Expression, Fixed, ToExpr, WitIn},
+    scheme::constants::MIN_PAR_SIZE,
+    set_fixed_val, set_val,
+    structs::ROMType,
+    tables::TableCircuit,
+    uint::constants::RANGE_CHIP_BIT_WIDTH,
+    witness::RowMajorMatrix,
+};
+use ff_ext::ExtensionField;
+use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
+
+#[derive(Clone, Debug)]
+pub struct ProgramTableConfig {
+    pc: Fixed,
+    // Decoded instruction fields.
+    opcode: Fixed,
+    rd: Fixed,
+    funct3: Fixed,
+    rs1: Fixed,
+    rs2: Fixed,
+    /// The complete immediate value, for instruction types I/S/B/U/J.
+    /// Otherwise, the field funct7 of R-Type instructions.
+    imm_or_funct7: Fixed,
+    /// Multiplicity of the record - how many times an instruction is visited.
+    mlt: WitIn,
+}
+
+pub struct ProgramTableCircuit<E>(PhantomData<E>);
+
+impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
+    type TableConfig = ProgramTableConfig;
+    type Input = [u32];
+
+    fn name() -> String {
+        "PROGRAM".into()
+    }
+
+    fn construct_circuit(cb: &mut CircuitBuilder<E>) -> Result<ProgramTableConfig, ZKVMError> {
+        let pc = cb.create_fixed(|| "pc")?;
+        let opcode = cb.create_fixed(|| "opcode")?;
+        let rd = cb.create_fixed(|| "rd")?;
+        let funct3 = cb.create_fixed(|| "funct3")?;
+        let rs1 = cb.create_fixed(|| "rs1")?;
+        let rs2 = cb.create_fixed(|| "rs2")?;
+        let imm_or_funct7 = cb.create_fixed(|| "imm_or_funct7")?;
+
+        let mlt = cb.create_witin(|| "mlt")?;
+
+        let record_exprs = cb.rlc_chip_record(vec![
+            Expression::Constant(E::BaseField::from(ROMType::Instruction as u64)),
+            Expression::Fixed(pc.clone()),
+            Expression::Fixed(opcode.clone()),
+            Expression::Fixed(rd.clone()),
+            Expression::Fixed(funct3.clone()),
+            Expression::Fixed(rs1.clone()),
+            Expression::Fixed(rs2.clone()),
+            Expression::Fixed(imm_or_funct7.clone()),
+        ]);
+
+        cb.lk_table_record(|| "prog table", record_exprs, mlt.expr())?;
+
+        Ok(ProgramTableConfig {
+            pc,
+            opcode,
+            rd,
+            funct3,
+            rs1,
+            rs2,
+            imm_or_funct7,
+            mlt,
+        })
+    }
+
+    fn generate_fixed_traces(
+        config: &ProgramTableConfig,
+        num_fixed: usize,
+        _input: &[u32],
+    ) -> RowMajorMatrix<E::BaseField> {
+        // TODO: get bytecode of the program.
+        let num_instructions = 1 << RANGE_CHIP_BIT_WIDTH;
+
+        let mut fixed = RowMajorMatrix::<E::BaseField>::new(num_instructions, num_fixed);
+
+        fixed
+            .par_iter_mut()
+            .with_min_len(MIN_PAR_SIZE)
+            .zip((0..num_instructions).into_par_iter())
+            .for_each(|(row, i)| {
+                set_fixed_val!(row, config.pc.0, E::BaseField::from(0 as u64));
+                set_fixed_val!(row, config.opcode.0, E::BaseField::from(0 as u64));
+                set_fixed_val!(row, config.rd.0, E::BaseField::from(0 as u64));
+                set_fixed_val!(row, config.funct3.0, E::BaseField::from(0 as u64));
+                set_fixed_val!(row, config.rs1.0, E::BaseField::from(0 as u64));
+                set_fixed_val!(row, config.rs2.0, E::BaseField::from(0 as u64));
+                set_fixed_val!(row, config.imm_or_funct7.0, E::BaseField::from(0 as u64));
+            });
+
+        fixed
+    }
+
+    fn assign_instances(
+        config: &Self::TableConfig,
+        num_witin: usize,
+        multiplicity: &[HashMap<u64, usize>],
+    ) -> Result<RowMajorMatrix<E::BaseField>, ZKVMError> {
+        // TODO: get instruction count.
+        tracing::debug!("num_witin: {}", num_witin);
+        let num_instructions = 1 << RANGE_CHIP_BIT_WIDTH;
+
+        let multiplicity = &multiplicity[ROMType::Instruction as usize];
+        tracing::debug!("multiplicity: {:?}", multiplicity);
+
+        let mut prog_mlt = vec![0_usize; num_instructions];
+        for (limb, mlt) in multiplicity {
+            prog_mlt[*limb as usize] = *mlt;
+        }
+
+        let mut witness = RowMajorMatrix::<E::BaseField>::new(prog_mlt.len(), num_witin);
+        witness
+            .par_iter_mut()
+            .with_min_len(MIN_PAR_SIZE)
+            .zip(prog_mlt.into_par_iter())
+            .for_each(|(row, mlt)| {
+                set_val!(row, config.mlt, E::BaseField::from(0 as u64)); // TODO
+            });
+
+        Ok(witness)
+    }
+}
diff --git a/ceno_zkvm/src/tables/range.rs b/ceno_zkvm/src/tables/range.rs
index 293077f46..0ae419aa4 100644
--- a/ceno_zkvm/src/tables/range.rs
+++ b/ceno_zkvm/src/tables/range.rs
@@ -24,7 +24,7 @@ pub struct RangeTableCircuit<E>(PhantomData<E>);
 
 impl<E: ExtensionField> TableCircuit<E> for RangeTableCircuit<E> {
     type TableConfig = RangeTableConfig;
-    type Input = u64;
+    type Input = ();
 
     fn name() -> String {
         "RANGE".into()
@@ -47,6 +47,7 @@ impl<E: ExtensionField> TableCircuit<E> for RangeTableCircuit<E> {
     fn generate_fixed_traces(
         config: &RangeTableConfig,
         num_fixed: usize,
+        _input: &(),
     ) -> RowMajorMatrix<E::BaseField> {
         let num_u16s = 1 << 16;
         let mut fixed = RowMajorMatrix::<E::BaseField>::new(num_u16s, num_fixed);

From adcfa594c0f41ddd5ec7f9497e76705f8ab9e0f2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Nicolas?= <info@nau.re>
Date: Wed, 11 Sep 2024 17:39:37 +0200
Subject: [PATCH 02/12] program-table: fetch instruction from ADD circuit

---
 ceno_zkvm/examples/riscv_add.rs            |  8 +++-
 ceno_zkvm/src/chip_handler/general.rs      | 16 ++++++++
 ceno_zkvm/src/instructions/riscv/addsub.rs |  4 ++
 ceno_zkvm/src/structs.rs                   |  4 +-
 ceno_zkvm/src/tables/mod.rs                |  6 ++-
 ceno_zkvm/src/tables/program.rs            | 45 +++++++++++++---------
 ceno_zkvm/src/tables/range.rs              |  6 ++-
 ceno_zkvm/src/witness.rs                   | 11 +++++-
 8 files changed, 73 insertions(+), 27 deletions(-)

diff --git a/ceno_zkvm/examples/riscv_add.rs b/ceno_zkvm/examples/riscv_add.rs
index e02577166..f06c88b55 100644
--- a/ceno_zkvm/examples/riscv_add.rs
+++ b/ceno_zkvm/examples/riscv_add.rs
@@ -134,10 +134,14 @@ fn main() {
         zkvm_witness.finalize_lk_multiplicities();
         // assign table circuits
         zkvm_witness
-            .assign_table_circuit::<RangeTableCircuit<E>>(&zkvm_cs, &range_config)
+            .assign_table_circuit::<RangeTableCircuit<E>>(&zkvm_cs, &range_config, &())
             .unwrap();
         zkvm_witness
-            .assign_table_circuit::<ProgramTableCircuit<E>>(&zkvm_cs, &prog_config)
+            .assign_table_circuit::<ProgramTableCircuit<E>>(
+                &zkvm_cs,
+                &prog_config,
+                &PROGRAM_ADD_LOOP.len(),
+            )
             .unwrap();
 
         let timer = Instant::now();
diff --git a/ceno_zkvm/src/chip_handler/general.rs b/ceno_zkvm/src/chip_handler/general.rs
index 0b4a672a8..413ee417c 100644
--- a/ceno_zkvm/src/chip_handler/general.rs
+++ b/ceno_zkvm/src/chip_handler/general.rs
@@ -55,6 +55,22 @@ impl<'a, E: ExtensionField> CircuitBuilder<'a, E> {
         self.cs.lk_table_record(name_fn, rlc_record, multiplicity)
     }
 
+    /// Fetch an instruction at a given PC from the Program table.
+    pub fn lk_fetch(&mut self, pc: Expression<E>) -> Result<(), ZKVMError> {
+        let insn_record = self.rlc_chip_record(vec![
+            E::BaseField::from(ROMType::Instruction as u64).expr(),
+            pc,
+            // TODO: instruction fields.
+            0_usize.into(),
+            0_usize.into(),
+            0_usize.into(),
+            0_usize.into(),
+            0_usize.into(),
+            0_usize.into(),
+        ]);
+        self.cs.lk_record(|| "fetch", insn_record)
+    }
+
     pub fn read_record<NR, N>(
         &mut self,
         name_fn: N,
diff --git a/ceno_zkvm/src/instructions/riscv/addsub.rs b/ceno_zkvm/src/instructions/riscv/addsub.rs
index f563ba68b..bc92c510f 100644
--- a/ceno_zkvm/src/instructions/riscv/addsub.rs
+++ b/ceno_zkvm/src/instructions/riscv/addsub.rs
@@ -59,6 +59,9 @@ fn add_sub_gadget<E: ExtensionField, const IS_ADD: bool>(
 
     let next_pc = pc.expr() + PC_STEP_SIZE.into();
 
+    // Instruction check.
+    circuit_builder.lk_fetch(pc.expr())?;
+
     // Execution result = addend0 + addend1, with carry.
     let prev_rd_value = RegUInt::new_unchecked(|| "prev_rd_value", circuit_builder)?;
 
@@ -157,6 +160,7 @@ impl<E: ExtensionField> Instruction<E> for AddInstruction<E> {
         lk_multiplicity: &mut LkMultiplicity,
         step: &StepRecord,
     ) -> Result<(), ZKVMError> {
+        lk_multiplicity.fetch(step.pc().before.0);
         // TODO use fields from step
         set_val!(instance, config.pc, 1);
         set_val!(instance, config.ts, 2);
diff --git a/ceno_zkvm/src/structs.rs b/ceno_zkvm/src/structs.rs
index 639c04ca0..e29a8afcd 100644
--- a/ceno_zkvm/src/structs.rs
+++ b/ceno_zkvm/src/structs.rs
@@ -155,7 +155,7 @@ impl<E: ExtensionField> ZKVMFixedTraces<E> {
         &mut self,
         cs: &ZKVMConstraintSystem<E>,
         config: TC::TableConfig,
-        input: &TC::Input,
+        input: &TC::FixedInput,
     ) {
         let cs = cs.get_cs(&TC::name()).expect("cs not found");
         assert!(
@@ -228,6 +228,7 @@ impl<E: ExtensionField> ZKVMWitnesses<E> {
         &mut self,
         cs: &ZKVMConstraintSystem<E>,
         config: &TC::TableConfig,
+        input: &TC::WitnessInput,
     ) -> Result<(), ZKVMError> {
         assert!(self.combined_lk_mlt.is_some());
 
@@ -236,6 +237,7 @@ impl<E: ExtensionField> ZKVMWitnesses<E> {
             config,
             cs.num_witin as usize,
             self.combined_lk_mlt.as_ref().unwrap(),
+            input,
         )?;
         assert!(self.witnesses.insert(TC::name(), witness).is_none());
 
diff --git a/ceno_zkvm/src/tables/mod.rs b/ceno_zkvm/src/tables/mod.rs
index fdd74a4b0..89046b5f2 100644
--- a/ceno_zkvm/src/tables/mod.rs
+++ b/ceno_zkvm/src/tables/mod.rs
@@ -10,7 +10,8 @@ pub use program::ProgramTableCircuit;
 
 pub trait TableCircuit<E: ExtensionField> {
     type TableConfig: Send + Sync;
-    type Input: Send + Sync + ?Sized;
+    type FixedInput: Send + Sync + ?Sized;
+    type WitnessInput: Send + Sync + ?Sized;
 
     fn name() -> String;
 
@@ -21,12 +22,13 @@ pub trait TableCircuit<E: ExtensionField> {
     fn generate_fixed_traces(
         config: &Self::TableConfig,
         num_fixed: usize,
-        input: &Self::Input,
+        input: &Self::FixedInput,
     ) -> RowMajorMatrix<E::BaseField>;
 
     fn assign_instances(
         config: &Self::TableConfig,
         num_witin: usize,
         multiplicity: &[HashMap<u64, usize>],
+        input: &Self::WitnessInput,
     ) -> Result<RowMajorMatrix<E::BaseField>, ZKVMError>;
 }
diff --git a/ceno_zkvm/src/tables/program.rs b/ceno_zkvm/src/tables/program.rs
index 23529bd62..2ce3ac468 100644
--- a/ceno_zkvm/src/tables/program.rs
+++ b/ceno_zkvm/src/tables/program.rs
@@ -8,9 +8,9 @@ use crate::{
     set_fixed_val, set_val,
     structs::ROMType,
     tables::TableCircuit,
-    uint::constants::RANGE_CHIP_BIT_WIDTH,
     witness::RowMajorMatrix,
 };
+use ceno_emul::{CENO_PLATFORM, WORD_SIZE};
 use ff_ext::ExtensionField;
 use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
 
@@ -34,7 +34,8 @@ pub struct ProgramTableCircuit<E>(PhantomData<E>);
 
 impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
     type TableConfig = ProgramTableConfig;
-    type Input = [u32];
+    type FixedInput = [u32];
+    type WitnessInput = usize;
 
     fn name() -> String {
         "PROGRAM".into()
@@ -79,10 +80,11 @@ impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
     fn generate_fixed_traces(
         config: &ProgramTableConfig,
         num_fixed: usize,
-        _input: &[u32],
+        input: &[u32],
     ) -> RowMajorMatrix<E::BaseField> {
         // TODO: get bytecode of the program.
-        let num_instructions = 1 << RANGE_CHIP_BIT_WIDTH;
+        let num_instructions = input.len();
+        let pc_start = CENO_PLATFORM.pc_start() as u64;
 
         let mut fixed = RowMajorMatrix::<E::BaseField>::new(num_instructions, num_fixed);
 
@@ -91,13 +93,14 @@ impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
             .with_min_len(MIN_PAR_SIZE)
             .zip((0..num_instructions).into_par_iter())
             .for_each(|(row, i)| {
-                set_fixed_val!(row, config.pc.0, E::BaseField::from(0 as u64));
-                set_fixed_val!(row, config.opcode.0, E::BaseField::from(0 as u64));
-                set_fixed_val!(row, config.rd.0, E::BaseField::from(0 as u64));
-                set_fixed_val!(row, config.funct3.0, E::BaseField::from(0 as u64));
-                set_fixed_val!(row, config.rs1.0, E::BaseField::from(0 as u64));
-                set_fixed_val!(row, config.rs2.0, E::BaseField::from(0 as u64));
-                set_fixed_val!(row, config.imm_or_funct7.0, E::BaseField::from(0 as u64));
+                let pc = pc_start + (i * WORD_SIZE) as u64;
+                set_fixed_val!(row, config.pc, E::BaseField::from(pc));
+                set_fixed_val!(row, config.opcode, E::BaseField::from(0 as u64));
+                set_fixed_val!(row, config.rd, E::BaseField::from(0 as u64));
+                set_fixed_val!(row, config.funct3, E::BaseField::from(0 as u64));
+                set_fixed_val!(row, config.rs1, E::BaseField::from(0 as u64));
+                set_fixed_val!(row, config.rs2, E::BaseField::from(0 as u64));
+                set_fixed_val!(row, config.imm_or_funct7, E::BaseField::from(0 as u64));
             });
 
         fixed
@@ -107,17 +110,21 @@ impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
         config: &Self::TableConfig,
         num_witin: usize,
         multiplicity: &[HashMap<u64, usize>],
+        num_instructions: &usize,
     ) -> Result<RowMajorMatrix<E::BaseField>, ZKVMError> {
-        // TODO: get instruction count.
-        tracing::debug!("num_witin: {}", num_witin);
-        let num_instructions = 1 << RANGE_CHIP_BIT_WIDTH;
+        tracing::debug!(
+            "num_instructions: {}. num_witin: {}",
+            *num_instructions,
+            num_witin,
+        );
 
         let multiplicity = &multiplicity[ROMType::Instruction as usize];
-        tracing::debug!("multiplicity: {:?}", multiplicity);
 
-        let mut prog_mlt = vec![0_usize; num_instructions];
-        for (limb, mlt) in multiplicity {
-            prog_mlt[*limb as usize] = *mlt;
+        let mut prog_mlt = vec![0_usize; *num_instructions];
+        for (pc, mlt) in multiplicity {
+            let i = (*pc as usize - CENO_PLATFORM.pc_start() as usize) / WORD_SIZE;
+            tracing::debug!("pc=0x{:x} index={} mlt={}", *pc, i, mlt);
+            prog_mlt[i] = *mlt;
         }
 
         let mut witness = RowMajorMatrix::<E::BaseField>::new(prog_mlt.len(), num_witin);
@@ -126,7 +133,7 @@ impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
             .with_min_len(MIN_PAR_SIZE)
             .zip(prog_mlt.into_par_iter())
             .for_each(|(row, mlt)| {
-                set_val!(row, config.mlt, E::BaseField::from(0 as u64)); // TODO
+                set_val!(row, config.mlt, E::BaseField::from(mlt as u64));
             });
 
         Ok(witness)
diff --git a/ceno_zkvm/src/tables/range.rs b/ceno_zkvm/src/tables/range.rs
index 0ae419aa4..2b195ea63 100644
--- a/ceno_zkvm/src/tables/range.rs
+++ b/ceno_zkvm/src/tables/range.rs
@@ -24,7 +24,8 @@ pub struct RangeTableCircuit<E>(PhantomData<E>);
 
 impl<E: ExtensionField> TableCircuit<E> for RangeTableCircuit<E> {
     type TableConfig = RangeTableConfig;
-    type Input = ();
+    type FixedInput = ();
+    type WitnessInput = ();
 
     fn name() -> String {
         "RANGE".into()
@@ -56,7 +57,7 @@ impl<E: ExtensionField> TableCircuit<E> for RangeTableCircuit<E> {
             .with_min_len(MIN_PAR_SIZE)
             .zip((0..num_u16s).into_par_iter())
             .for_each(|(row, i)| {
-                set_fixed_val!(row, config.u16_tbl.0, E::BaseField::from(i as u64));
+                set_fixed_val!(row, config.u16_tbl, E::BaseField::from(i as u64));
             });
 
         fixed
@@ -66,6 +67,7 @@ impl<E: ExtensionField> TableCircuit<E> for RangeTableCircuit<E> {
         config: &Self::TableConfig,
         num_witin: usize,
         multiplicity: &[HashMap<u64, usize>],
+        _input: &(),
     ) -> Result<RowMajorMatrix<E::BaseField>, ZKVMError> {
         let multiplicity = &multiplicity[ROMType::U16 as usize];
         let mut u16_mlt = vec![0; 1 << RANGE_CHIP_BIT_WIDTH];
diff --git a/ceno_zkvm/src/witness.rs b/ceno_zkvm/src/witness.rs
index a77d9c4c3..0f9e0d406 100644
--- a/ceno_zkvm/src/witness.rs
+++ b/ceno_zkvm/src/witness.rs
@@ -26,7 +26,7 @@ macro_rules! set_val {
 #[macro_export]
 macro_rules! set_fixed_val {
     ($ins:ident, $field:expr, $val:expr) => {
-        $ins[$field as usize] = MaybeUninit::new($val);
+        $ins[$field.0] = MaybeUninit::new($val);
     };
 }
 
@@ -142,6 +142,15 @@ impl LkMultiplicity {
             .or_default()) += 1;
     }
 
+    pub fn fetch(&mut self, pc: u32) {
+        let multiplicity = self
+            .multiplicity
+            .get_or(|| RefCell::new(array::from_fn(|_| HashMap::new())));
+        (*multiplicity.borrow_mut()[ROMType::Instruction as usize]
+            .entry(pc as u64)
+            .or_default()) += 1;
+    }
+
     /// merge result from multiple thread local to single result
     pub fn into_finalize_result(self) -> [HashMap<u64, usize>; mem::variant_count::<ROMType>()] {
         Arc::try_unwrap(self.multiplicity)

From 7e2919827230c62e25cafce1ff0bf0a7cdf7f3d7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Nicolas?= <info@nau.re>
Date: Wed, 11 Sep 2024 20:15:21 +0200
Subject: [PATCH 03/12] program-table: Example passing with PC and opcode

---
 ceno_zkvm/examples/riscv_add.rs            |  2 +-
 ceno_zkvm/src/chip_handler/general.rs      |  4 ++--
 ceno_zkvm/src/instructions/riscv/addsub.rs |  4 ++--
 ceno_zkvm/src/tables/program.rs            | 28 ++++++++++++++--------
 4 files changed, 23 insertions(+), 15 deletions(-)

diff --git a/ceno_zkvm/examples/riscv_add.rs b/ceno_zkvm/examples/riscv_add.rs
index f06c88b55..8dfa068ab 100644
--- a/ceno_zkvm/examples/riscv_add.rs
+++ b/ceno_zkvm/examples/riscv_add.rs
@@ -103,7 +103,7 @@ fn main() {
     let prover = ZKVMProver::new(pk);
     let verifier = ZKVMVerifier::new(vk);
 
-    // for instance_num_vars in 8..22 {
+    // for instance_num_vars in 8..22 { // TODO: restore.
     for instance_num_vars in 2..3 {
         let num_instances = 1 << instance_num_vars;
         let mut vm = VMState::new(CENO_PLATFORM);
diff --git a/ceno_zkvm/src/chip_handler/general.rs b/ceno_zkvm/src/chip_handler/general.rs
index 413ee417c..1b64eb0f0 100644
--- a/ceno_zkvm/src/chip_handler/general.rs
+++ b/ceno_zkvm/src/chip_handler/general.rs
@@ -56,12 +56,12 @@ impl<'a, E: ExtensionField> CircuitBuilder<'a, E> {
     }
 
     /// Fetch an instruction at a given PC from the Program table.
-    pub fn lk_fetch(&mut self, pc: Expression<E>) -> Result<(), ZKVMError> {
+    pub fn lk_fetch(&mut self, pc: Expression<E>, opcode: Expression<E>) -> Result<(), ZKVMError> {
         let insn_record = self.rlc_chip_record(vec![
             E::BaseField::from(ROMType::Instruction as u64).expr(),
             pc,
             // TODO: instruction fields.
-            0_usize.into(),
+            opcode,
             0_usize.into(),
             0_usize.into(),
             0_usize.into(),
diff --git a/ceno_zkvm/src/instructions/riscv/addsub.rs b/ceno_zkvm/src/instructions/riscv/addsub.rs
index bc92c510f..095b1f619 100644
--- a/ceno_zkvm/src/instructions/riscv/addsub.rs
+++ b/ceno_zkvm/src/instructions/riscv/addsub.rs
@@ -60,7 +60,7 @@ fn add_sub_gadget<E: ExtensionField, const IS_ADD: bool>(
     let next_pc = pc.expr() + PC_STEP_SIZE.into();
 
     // Instruction check.
-    circuit_builder.lk_fetch(pc.expr())?;
+    circuit_builder.lk_fetch(pc.expr(), 0x33.into())?;
 
     // Execution result = addend0 + addend1, with carry.
     let prev_rd_value = RegUInt::new_unchecked(|| "prev_rd_value", circuit_builder)?;
@@ -161,8 +161,8 @@ impl<E: ExtensionField> Instruction<E> for AddInstruction<E> {
         step: &StepRecord,
     ) -> Result<(), ZKVMError> {
         lk_multiplicity.fetch(step.pc().before.0);
+        set_val!(instance, config.pc, step.pc().before.0 as u64);
         // TODO use fields from step
-        set_val!(instance, config.pc, 1);
         set_val!(instance, config.ts, 2);
         let addend_0 = UIntValue::new_unchecked(step.rs1().unwrap().value);
         let addend_1 = UIntValue::new_unchecked(step.rs2().unwrap().value);
diff --git a/ceno_zkvm/src/tables/program.rs b/ceno_zkvm/src/tables/program.rs
index 2ce3ac468..eb4e536bf 100644
--- a/ceno_zkvm/src/tables/program.rs
+++ b/ceno_zkvm/src/tables/program.rs
@@ -10,7 +10,7 @@ use crate::{
     tables::TableCircuit,
     witness::RowMajorMatrix,
 };
-use ceno_emul::{CENO_PLATFORM, WORD_SIZE};
+use ceno_emul::{DecodedInstruction, Word, CENO_PLATFORM, WORD_SIZE};
 use ff_ext::ExtensionField;
 use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
 
@@ -80,10 +80,10 @@ impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
     fn generate_fixed_traces(
         config: &ProgramTableConfig,
         num_fixed: usize,
-        input: &[u32],
+        program: &[Word],
     ) -> RowMajorMatrix<E::BaseField> {
         // TODO: get bytecode of the program.
-        let num_instructions = input.len();
+        let num_instructions = program.len();
         let pc_start = CENO_PLATFORM.pc_start() as u64;
 
         let mut fixed = RowMajorMatrix::<E::BaseField>::new(num_instructions, num_fixed);
@@ -94,13 +94,21 @@ impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
             .zip((0..num_instructions).into_par_iter())
             .for_each(|(row, i)| {
                 let pc = pc_start + (i * WORD_SIZE) as u64;
-                set_fixed_val!(row, config.pc, E::BaseField::from(pc));
-                set_fixed_val!(row, config.opcode, E::BaseField::from(0 as u64));
-                set_fixed_val!(row, config.rd, E::BaseField::from(0 as u64));
-                set_fixed_val!(row, config.funct3, E::BaseField::from(0 as u64));
-                set_fixed_val!(row, config.rs1, E::BaseField::from(0 as u64));
-                set_fixed_val!(row, config.rs2, E::BaseField::from(0 as u64));
-                set_fixed_val!(row, config.imm_or_funct7, E::BaseField::from(0 as u64));
+                let insn = DecodedInstruction::new(program[i]);
+
+                tracing::debug!("pc=0x{:x} insn={:?}", pc, insn);
+
+                for (col, val) in [
+                    (&config.pc, pc as u32),
+                    (&config.opcode, insn.opcode()),
+                    (&config.rd, 0),
+                    (&config.funct3, 0),
+                    (&config.rs1, 0),
+                    (&config.rs2, 0),
+                    (&config.imm_or_funct7, 0),
+                ] {
+                    set_fixed_val!(row, *col, E::BaseField::from(val as u64));
+                }
             });
 
         fixed

From cab400f78a3c3e75959e41563f45065444f579cd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Nicolas?= <info@nau.re>
Date: Wed, 11 Sep 2024 21:25:17 +0200
Subject: [PATCH 04/12] program-table: Record type

---
 ceno_zkvm/src/chip_handler/general.rs      |  22 ++--
 ceno_zkvm/src/instructions/riscv/addsub.rs |  13 +-
 ceno_zkvm/src/tables/mod.rs                |   2 +-
 ceno_zkvm/src/tables/program.rs            | 142 +++++++++++++--------
 4 files changed, 107 insertions(+), 72 deletions(-)

diff --git a/ceno_zkvm/src/chip_handler/general.rs b/ceno_zkvm/src/chip_handler/general.rs
index 1b64eb0f0..96e8ce659 100644
--- a/ceno_zkvm/src/chip_handler/general.rs
+++ b/ceno_zkvm/src/chip_handler/general.rs
@@ -7,6 +7,7 @@ use crate::{
     error::ZKVMError,
     expression::{Expression, Fixed, ToExpr, WitIn},
     structs::ROMType,
+    tables::InsnRecord,
 };
 
 impl<'a, E: ExtensionField> CircuitBuilder<'a, E> {
@@ -56,19 +57,14 @@ impl<'a, E: ExtensionField> CircuitBuilder<'a, E> {
     }
 
     /// Fetch an instruction at a given PC from the Program table.
-    pub fn lk_fetch(&mut self, pc: Expression<E>, opcode: Expression<E>) -> Result<(), ZKVMError> {
-        let insn_record = self.rlc_chip_record(vec![
-            E::BaseField::from(ROMType::Instruction as u64).expr(),
-            pc,
-            // TODO: instruction fields.
-            opcode,
-            0_usize.into(),
-            0_usize.into(),
-            0_usize.into(),
-            0_usize.into(),
-            0_usize.into(),
-        ]);
-        self.cs.lk_record(|| "fetch", insn_record)
+    pub fn lk_fetch(&mut self, record: &InsnRecord<Expression<E>>) -> Result<(), ZKVMError> {
+        let rlc_record = {
+            let mut fields = vec![E::BaseField::from(ROMType::Instruction as u64).expr()];
+            fields.extend_from_slice(record.as_slice());
+            self.rlc_chip_record(fields)
+        };
+
+        self.cs.lk_record(|| "fetch", rlc_record)
     }
 
     pub fn read_record<NR, N>(
diff --git a/ceno_zkvm/src/instructions/riscv/addsub.rs b/ceno_zkvm/src/instructions/riscv/addsub.rs
index 095b1f619..cea1f0c8c 100644
--- a/ceno_zkvm/src/instructions/riscv/addsub.rs
+++ b/ceno_zkvm/src/instructions/riscv/addsub.rs
@@ -15,6 +15,7 @@ use crate::{
     expression::{ToExpr, WitIn},
     instructions::Instruction,
     set_val,
+    tables::InsnRecord,
     uint::UIntValue,
     witness::LkMultiplicity,
 };
@@ -59,8 +60,16 @@ fn add_sub_gadget<E: ExtensionField, const IS_ADD: bool>(
 
     let next_pc = pc.expr() + PC_STEP_SIZE.into();
 
-    // Instruction check.
-    circuit_builder.lk_fetch(pc.expr(), 0x33.into())?;
+    // Fetch the instruction.
+    circuit_builder.lk_fetch(&InsnRecord::new(
+        pc.expr(),
+        0x33.into(),
+        0x0.into(),
+        0x0.into(),
+        0x0.into(),
+        0x0.into(),
+        0x0.into(),
+    ))?;
 
     // Execution result = addend0 + addend1, with carry.
     let prev_rd_value = RegUInt::new_unchecked(|| "prev_rd_value", circuit_builder)?;
diff --git a/ceno_zkvm/src/tables/mod.rs b/ceno_zkvm/src/tables/mod.rs
index 89046b5f2..7c9a29ed4 100644
--- a/ceno_zkvm/src/tables/mod.rs
+++ b/ceno_zkvm/src/tables/mod.rs
@@ -6,7 +6,7 @@ mod range;
 pub use range::RangeTableCircuit;
 
 mod program;
-pub use program::ProgramTableCircuit;
+pub use program::{InsnRecord, ProgramTableCircuit};
 
 pub trait TableCircuit<E: ExtensionField> {
     type TableConfig: Send + Sync;
diff --git a/ceno_zkvm/src/tables/program.rs b/ceno_zkvm/src/tables/program.rs
index eb4e536bf..1dc1b60e5 100644
--- a/ceno_zkvm/src/tables/program.rs
+++ b/ceno_zkvm/src/tables/program.rs
@@ -12,20 +12,72 @@ use crate::{
 };
 use ceno_emul::{DecodedInstruction, Word, CENO_PLATFORM, WORD_SIZE};
 use ff_ext::ExtensionField;
+use itertools::Itertools;
 use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
 
 #[derive(Clone, Debug)]
-pub struct ProgramTableConfig {
-    pc: Fixed,
-    // Decoded instruction fields.
-    opcode: Fixed,
-    rd: Fixed,
-    funct3: Fixed,
-    rs1: Fixed,
-    rs2: Fixed,
+pub struct InsnRecord<T>([T; 7]);
+
+impl<T> InsnRecord<T> {
+    pub fn new(pc: T, opcode: T, rd: T, funct3: T, rs1: T, rs2: T, imm_or_funct7: T) -> Self {
+        InsnRecord([pc, opcode, rd, funct3, rs1, rs2, imm_or_funct7])
+    }
+
+    pub fn as_slice(&self) -> &[T] {
+        &self.0
+    }
+
+    pub fn pc(&self) -> &T {
+        &self.0[0]
+    }
+
+    pub fn opcode(&self) -> &T {
+        &self.0[1]
+    }
+
+    pub fn rd(&self) -> &T {
+        &self.0[2]
+    }
+
+    pub fn funct3(&self) -> &T {
+        &self.0[3]
+    }
+
+    pub fn rs1(&self) -> &T {
+        &self.0[4]
+    }
+
+    pub fn rs2(&self) -> &T {
+        &self.0[5]
+    }
+
     /// The complete immediate value, for instruction types I/S/B/U/J.
     /// Otherwise, the field funct7 of R-Type instructions.
-    imm_or_funct7: Fixed,
+    pub fn imm_or_funct7(&self) -> &T {
+        &self.0[6]
+    }
+}
+
+impl InsnRecord<u32> {
+    fn from_decoded(pc: u32, insn: &DecodedInstruction) -> Self {
+        InsnRecord::new(
+            pc,
+            insn.opcode(),
+            // TODO: implement and remove the 0s.
+            0 * insn.rd(),
+            0 * insn.func3(),
+            0 * insn.rs1(),
+            0 * insn.rs2(),
+            0 * insn.func7(), // TODO: get immediate for all types.
+        )
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct ProgramTableConfig {
+    /// The fixed table of instruction records.
+    record: InsnRecord<Fixed>,
+
     /// Multiplicity of the record - how many times an instruction is visited.
     mlt: WitIn,
 }
@@ -42,39 +94,32 @@ impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
     }
 
     fn construct_circuit(cb: &mut CircuitBuilder<E>) -> Result<ProgramTableConfig, ZKVMError> {
-        let pc = cb.create_fixed(|| "pc")?;
-        let opcode = cb.create_fixed(|| "opcode")?;
-        let rd = cb.create_fixed(|| "rd")?;
-        let funct3 = cb.create_fixed(|| "funct3")?;
-        let rs1 = cb.create_fixed(|| "rs1")?;
-        let rs2 = cb.create_fixed(|| "rs2")?;
-        let imm_or_funct7 = cb.create_fixed(|| "imm_or_funct7")?;
+        let record = InsnRecord([
+            cb.create_fixed(|| "pc")?,
+            cb.create_fixed(|| "opcode")?,
+            cb.create_fixed(|| "rd")?,
+            cb.create_fixed(|| "funct3")?,
+            cb.create_fixed(|| "rs1")?,
+            cb.create_fixed(|| "rs2")?,
+            cb.create_fixed(|| "imm_or_funct7")?,
+        ]);
 
         let mlt = cb.create_witin(|| "mlt")?;
 
-        let record_exprs = cb.rlc_chip_record(vec![
-            Expression::Constant(E::BaseField::from(ROMType::Instruction as u64)),
-            Expression::Fixed(pc.clone()),
-            Expression::Fixed(opcode.clone()),
-            Expression::Fixed(rd.clone()),
-            Expression::Fixed(funct3.clone()),
-            Expression::Fixed(rs1.clone()),
-            Expression::Fixed(rs2.clone()),
-            Expression::Fixed(imm_or_funct7.clone()),
-        ]);
+        let record_exprs = {
+            let mut fields = vec![E::BaseField::from(ROMType::Instruction as u64).expr()];
+            fields.extend(
+                record
+                    .as_slice()
+                    .iter()
+                    .map(|f| Expression::Fixed(f.clone())),
+            );
+            cb.rlc_chip_record(fields)
+        };
 
         cb.lk_table_record(|| "prog table", record_exprs, mlt.expr())?;
 
-        Ok(ProgramTableConfig {
-            pc,
-            opcode,
-            rd,
-            funct3,
-            rs1,
-            rs2,
-            imm_or_funct7,
-            mlt,
-        })
+        Ok(ProgramTableConfig { record, mlt })
     }
 
     fn generate_fixed_traces(
@@ -84,7 +129,7 @@ impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
     ) -> RowMajorMatrix<E::BaseField> {
         // TODO: get bytecode of the program.
         let num_instructions = program.len();
-        let pc_start = CENO_PLATFORM.pc_start() as u64;
+        let pc_start = CENO_PLATFORM.pc_start();
 
         let mut fixed = RowMajorMatrix::<E::BaseField>::new(num_instructions, num_fixed);
 
@@ -93,21 +138,12 @@ impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
             .with_min_len(MIN_PAR_SIZE)
             .zip((0..num_instructions).into_par_iter())
             .for_each(|(row, i)| {
-                let pc = pc_start + (i * WORD_SIZE) as u64;
+                let pc = pc_start + (i * WORD_SIZE) as u32;
                 let insn = DecodedInstruction::new(program[i]);
+                let values = InsnRecord::from_decoded(pc, &insn);
 
-                tracing::debug!("pc=0x{:x} insn={:?}", pc, insn);
-
-                for (col, val) in [
-                    (&config.pc, pc as u32),
-                    (&config.opcode, insn.opcode()),
-                    (&config.rd, 0),
-                    (&config.funct3, 0),
-                    (&config.rs1, 0),
-                    (&config.rs2, 0),
-                    (&config.imm_or_funct7, 0),
-                ] {
-                    set_fixed_val!(row, *col, E::BaseField::from(val as u64));
+                for (col, val) in config.record.as_slice().iter().zip_eq(values.as_slice()) {
+                    set_fixed_val!(row, *col, E::BaseField::from(*val as u64));
                 }
             });
 
@@ -120,12 +156,6 @@ impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
         multiplicity: &[HashMap<u64, usize>],
         num_instructions: &usize,
     ) -> Result<RowMajorMatrix<E::BaseField>, ZKVMError> {
-        tracing::debug!(
-            "num_instructions: {}. num_witin: {}",
-            *num_instructions,
-            num_witin,
-        );
-
         let multiplicity = &multiplicity[ROMType::Instruction as usize];
 
         let mut prog_mlt = vec![0_usize; *num_instructions];

From 3eeb0f776918776f916844deae4da0bbc9f9569d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Nicolas?= <info@nau.re>
Date: Wed, 11 Sep 2024 21:49:43 +0200
Subject: [PATCH 05/12] program-table: Support registers and funct

---
 ceno_emul/src/rv32im.rs                    |  4 +-
 ceno_zkvm/src/instructions/riscv/addsub.rs | 57 ++++++++++++++--------
 ceno_zkvm/src/tables/program.rs            | 11 ++---
 3 files changed, 44 insertions(+), 28 deletions(-)

diff --git a/ceno_emul/src/rv32im.rs b/ceno_emul/src/rv32im.rs
index 64e758f55..3e94ec751 100644
--- a/ceno_emul/src/rv32im.rs
+++ b/ceno_emul/src/rv32im.rs
@@ -215,7 +215,7 @@ impl DecodedInstruction {
         self.rd
     }
 
-    pub fn func3(&self) -> u32 {
+    pub fn funct3(&self) -> u32 {
         self.func3
     }
 
@@ -227,7 +227,7 @@ impl DecodedInstruction {
         self.rs2
     }
 
-    pub fn func7(&self) -> u32 {
+    pub fn funct7(&self) -> u32 {
         self.func7
     }
 
diff --git a/ceno_zkvm/src/instructions/riscv/addsub.rs b/ceno_zkvm/src/instructions/riscv/addsub.rs
index cea1f0c8c..75eb04ac1 100644
--- a/ceno_zkvm/src/instructions/riscv/addsub.rs
+++ b/ceno_zkvm/src/instructions/riscv/addsub.rs
@@ -60,17 +60,6 @@ fn add_sub_gadget<E: ExtensionField, const IS_ADD: bool>(
 
     let next_pc = pc.expr() + PC_STEP_SIZE.into();
 
-    // Fetch the instruction.
-    circuit_builder.lk_fetch(&InsnRecord::new(
-        pc.expr(),
-        0x33.into(),
-        0x0.into(),
-        0x0.into(),
-        0x0.into(),
-        0x0.into(),
-        0x0.into(),
-    ))?;
-
     // Execution result = addend0 + addend1, with carry.
     let prev_rd_value = RegUInt::new_unchecked(|| "prev_rd_value", circuit_builder)?;
 
@@ -101,6 +90,17 @@ fn add_sub_gadget<E: ExtensionField, const IS_ADD: bool>(
     let rs2_id = circuit_builder.create_witin(|| "rs2_id")?;
     let rd_id = circuit_builder.create_witin(|| "rd_id")?;
 
+    // Fetch the instruction.
+    circuit_builder.lk_fetch(&InsnRecord::new(
+        pc.expr(),
+        0x33.into(),
+        rd_id.expr(),
+        0x0.into(),
+        rs1_id.expr(),
+        rs2_id.expr(),
+        (if IS_ADD { 0 } else { 0x20 }).into(),
+    ))?;
+
     // TODO remove me, this is just for testing degree > 1 sumcheck in main constraints
     circuit_builder.require_zero(
         || "test_degree > 1",
@@ -171,6 +171,7 @@ impl<E: ExtensionField> Instruction<E> for AddInstruction<E> {
     ) -> Result<(), ZKVMError> {
         lk_multiplicity.fetch(step.pc().before.0);
         set_val!(instance, config.pc, step.pc().before.0 as u64);
+
         // TODO use fields from step
         set_val!(instance, config.ts, 2);
         let addend_0 = UIntValue::new_unchecked(step.rs1().unwrap().value);
@@ -193,10 +194,17 @@ impl<E: ExtensionField> Instruction<E> for AddInstruction<E> {
                 .map(|carry| E::BaseField::from(carry as u64))
                 .collect_vec(),
         );
+
+        // Register IDs.
+        for (reg_col, reg_id) in [
+            (config.rs1_id, step.rs1().unwrap().register_index()),
+            (config.rs2_id, step.rs2().unwrap().register_index()),
+            (config.rd_id, step.rd().unwrap().register_index()),
+        ] {
+            set_val!(instance, reg_col, reg_id as u64);
+        }
+
         // TODO #167
-        set_val!(instance, config.rs1_id, 2);
-        set_val!(instance, config.rs2_id, 2);
-        set_val!(instance, config.rd_id, 2);
         set_val!(instance, config.prev_rs1_ts, 2);
         set_val!(instance, config.prev_rs2_ts, 2);
         set_val!(instance, config.prev_rd_ts, 2);
@@ -220,11 +228,13 @@ impl<E: ExtensionField> Instruction<E> for SubInstruction<E> {
     fn assign_instance(
         config: &Self::InstructionConfig,
         instance: &mut [MaybeUninit<E::BaseField>],
-        _lk_multiplicity: &mut LkMultiplicity,
-        _step: &StepRecord,
+        lk_multiplicity: &mut LkMultiplicity,
+        step: &StepRecord,
     ) -> Result<(), ZKVMError> {
+        lk_multiplicity.fetch(step.pc().before.0);
+        set_val!(instance, config.pc, step.pc().before.0 as u64);
+
         // TODO use field from step
-        set_val!(instance, config.pc, _step.pc().before.0 as u64);
         set_val!(instance, config.ts, 2);
         config.prev_rd_value.wits_in().map(|prev_rd_value| {
             set_val!(instance, prev_rd_value[0], 4);
@@ -243,10 +253,17 @@ impl<E: ExtensionField> Instruction<E> for SubInstruction<E> {
             set_val!(instance, carry[0], 4);
             set_val!(instance, carry[1], 0);
         });
+
+        // Register IDs.
+        for (reg_col, reg_id) in [
+            (config.rs1_id, step.rs1().unwrap().register_index()),
+            (config.rs2_id, step.rs2().unwrap().register_index()),
+            (config.rd_id, step.rd().unwrap().register_index()),
+        ] {
+            set_val!(instance, reg_col, reg_id as u64);
+        }
+
         // TODO #167
-        set_val!(instance, config.rs1_id, 2);
-        set_val!(instance, config.rs2_id, 2);
-        set_val!(instance, config.rd_id, 2);
         set_val!(instance, config.prev_rs1_ts, 2);
         set_val!(instance, config.prev_rs2_ts, 2);
         set_val!(instance, config.prev_rd_ts, 2);
diff --git a/ceno_zkvm/src/tables/program.rs b/ceno_zkvm/src/tables/program.rs
index 1dc1b60e5..696c0e2c6 100644
--- a/ceno_zkvm/src/tables/program.rs
+++ b/ceno_zkvm/src/tables/program.rs
@@ -63,12 +63,11 @@ impl InsnRecord<u32> {
         InsnRecord::new(
             pc,
             insn.opcode(),
-            // TODO: implement and remove the 0s.
-            0 * insn.rd(),
-            0 * insn.func3(),
-            0 * insn.rs1(),
-            0 * insn.rs2(),
-            0 * insn.func7(), // TODO: get immediate for all types.
+            insn.rd(),
+            insn.funct3(),
+            insn.rs1(),
+            insn.rs2(),
+            insn.funct7(), // TODO: get immediate for all types.
         )
     }
 }

From 7e5c7180cffb0384f385a516e7cb15a33231671f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Nicolas?= <info@nau.re>
Date: Wed, 11 Sep 2024 22:06:39 +0200
Subject: [PATCH 06/12] program-table: put back example parameters

---
 ceno_zkvm/examples/riscv_add.rs | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/ceno_zkvm/examples/riscv_add.rs b/ceno_zkvm/examples/riscv_add.rs
index 8dfa068ab..14ef9bbcd 100644
--- a/ceno_zkvm/examples/riscv_add.rs
+++ b/ceno_zkvm/examples/riscv_add.rs
@@ -66,7 +66,6 @@ fn main() {
         .with(
             fmt::layer()
                 .compact()
-                .without_time()
                 .with_thread_ids(false)
                 .with_thread_names(false),
         )
@@ -103,8 +102,17 @@ fn main() {
     let prover = ZKVMProver::new(pk);
     let verifier = ZKVMVerifier::new(vk);
 
-    // for instance_num_vars in 8..22 { // TODO: restore.
-    for instance_num_vars in 2..3 {
+    let max_instance_num_vars = {
+        #[cfg(debug_assertions)]
+        {
+            tracing::info!("debug mode, testing only the smallest instance");
+            9
+        }
+        #[cfg(not(debug_assertions))]
+        22
+    };
+
+    for instance_num_vars in 8..max_instance_num_vars {
         let num_instances = 1 << instance_num_vars;
         let mut vm = VMState::new(CENO_PLATFORM);
         let pc_start = ByteAddr(CENO_PLATFORM.pc_start()).waddr();

From b7ec52645de4925fca12a4df9f53815dc9f597da Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Nicolas?= <info@nau.re>
Date: Wed, 11 Sep 2024 22:19:09 +0200
Subject: [PATCH 07/12] program-table: constants

---
 ceno_zkvm/src/instructions/riscv/addsub.rs | 11 ++++++++---
 ceno_zkvm/src/tables/program.rs            |  1 -
 2 files changed, 8 insertions(+), 4 deletions(-)

diff --git a/ceno_zkvm/src/instructions/riscv/addsub.rs b/ceno_zkvm/src/instructions/riscv/addsub.rs
index 75eb04ac1..d0e18cffa 100644
--- a/ceno_zkvm/src/instructions/riscv/addsub.rs
+++ b/ceno_zkvm/src/instructions/riscv/addsub.rs
@@ -21,6 +21,11 @@ use crate::{
 };
 use core::mem::MaybeUninit;
 
+const OPCODE_OP: usize = 0x33;
+const FUNCT3_ADD_SUB: usize = 0;
+const FUNCT7_ADD: usize = 0;
+const FUNCT7_SUB: usize = 0x20;
+
 pub struct AddInstruction<E>(PhantomData<E>);
 pub struct SubInstruction<E>(PhantomData<E>);
 
@@ -93,12 +98,12 @@ fn add_sub_gadget<E: ExtensionField, const IS_ADD: bool>(
     // Fetch the instruction.
     circuit_builder.lk_fetch(&InsnRecord::new(
         pc.expr(),
-        0x33.into(),
+        OPCODE_OP.into(),
         rd_id.expr(),
-        0x0.into(),
+        FUNCT3_ADD_SUB.into(),
         rs1_id.expr(),
         rs2_id.expr(),
-        (if IS_ADD { 0 } else { 0x20 }).into(),
+        (if IS_ADD { FUNCT7_ADD } else { FUNCT7_SUB }).into(),
     ))?;
 
     // TODO remove me, this is just for testing degree > 1 sumcheck in main constraints
diff --git a/ceno_zkvm/src/tables/program.rs b/ceno_zkvm/src/tables/program.rs
index 696c0e2c6..6acb8991c 100644
--- a/ceno_zkvm/src/tables/program.rs
+++ b/ceno_zkvm/src/tables/program.rs
@@ -160,7 +160,6 @@ impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
         let mut prog_mlt = vec![0_usize; *num_instructions];
         for (pc, mlt) in multiplicity {
             let i = (*pc as usize - CENO_PLATFORM.pc_start() as usize) / WORD_SIZE;
-            tracing::debug!("pc=0x{:x} index={} mlt={}", *pc, i, mlt);
             prog_mlt[i] = *mlt;
         }
 

From cffa410a8051f825be4bafa8e891439697933ed5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Nicolas?= <info@nau.re>
Date: Fri, 13 Sep 2024 13:52:50 +0200
Subject: [PATCH 08/12] program-table: move constants

---
 ceno_zkvm/examples/riscv_add.rs               |  1 +
 ceno_zkvm/src/instructions/riscv/addsub.rs    | 10 ++++------
 ceno_zkvm/src/instructions/riscv/constants.rs |  5 +++++
 ceno_zkvm/src/tables/program.rs               |  4 ++--
 4 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/ceno_zkvm/examples/riscv_add.rs b/ceno_zkvm/examples/riscv_add.rs
index 9c5766a21..e98429f1f 100644
--- a/ceno_zkvm/examples/riscv_add.rs
+++ b/ceno_zkvm/examples/riscv_add.rs
@@ -102,6 +102,7 @@ fn main() {
     let prover = ZKVMProver::new(pk);
     let verifier = ZKVMVerifier::new(vk);
 
+    // TODO: remove after #211.
     let max_instance_num_vars = {
         #[cfg(debug_assertions)]
         {
diff --git a/ceno_zkvm/src/instructions/riscv/addsub.rs b/ceno_zkvm/src/instructions/riscv/addsub.rs
index 708341b69..fa4d62368 100644
--- a/ceno_zkvm/src/instructions/riscv/addsub.rs
+++ b/ceno_zkvm/src/instructions/riscv/addsub.rs
@@ -6,7 +6,10 @@ use itertools::Itertools;
 
 use super::{
     config::ExprLtConfig,
-    constants::{OPType, OpcodeType, RegUInt, PC_STEP_SIZE},
+    constants::{
+        OPType, OpcodeType, RegUInt, FUNCT3_ADD_SUB, FUNCT7_ADD, FUNCT7_SUB, OPCODE_OP,
+        PC_STEP_SIZE,
+    },
     RIVInstruction,
 };
 use crate::{
@@ -22,11 +25,6 @@ use crate::{
 };
 use core::mem::MaybeUninit;
 
-const OPCODE_OP: usize = 0x33;
-const FUNCT3_ADD_SUB: usize = 0;
-const FUNCT7_ADD: usize = 0;
-const FUNCT7_SUB: usize = 0x20;
-
 pub struct AddInstruction<E>(PhantomData<E>);
 pub struct SubInstruction<E>(PhantomData<E>);
 
diff --git a/ceno_zkvm/src/instructions/riscv/constants.rs b/ceno_zkvm/src/instructions/riscv/constants.rs
index 37222be8b..813f2b340 100644
--- a/ceno_zkvm/src/instructions/riscv/constants.rs
+++ b/ceno_zkvm/src/instructions/riscv/constants.rs
@@ -4,6 +4,11 @@ use crate::uint::UInt;
 
 pub(crate) const PC_STEP_SIZE: usize = 4;
 
+pub const OPCODE_OP: usize = 0x33;
+pub const FUNCT3_ADD_SUB: usize = 0;
+pub const FUNCT7_ADD: usize = 0;
+pub const FUNCT7_SUB: usize = 0x20;
+
 #[allow(clippy::upper_case_acronyms)]
 #[derive(Debug, Clone, Copy)]
 pub enum OPType {
diff --git a/ceno_zkvm/src/tables/program.rs b/ceno_zkvm/src/tables/program.rs
index 6acb8991c..c8c66b777 100644
--- a/ceno_zkvm/src/tables/program.rs
+++ b/ceno_zkvm/src/tables/program.rs
@@ -135,8 +135,8 @@ impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
         fixed
             .par_iter_mut()
             .with_min_len(MIN_PAR_SIZE)
-            .zip((0..num_instructions).into_par_iter())
-            .for_each(|(row, i)| {
+            .enumerate()
+            .for_each(|(i, row)| {
                 let pc = pc_start + (i * WORD_SIZE) as u32;
                 let insn = DecodedInstruction::new(program[i]);
                 let values = InsnRecord::from_decoded(pc, &insn);

From b70aed797be519292405b8afaf64b6ec8019971f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Nicolas?= <info@nau.re>
Date: Fri, 13 Sep 2024 17:25:33 +0200
Subject: [PATCH 09/12] program-table: zip != enumerate

---
 ceno_zkvm/src/tables/program.rs | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/ceno_zkvm/src/tables/program.rs b/ceno_zkvm/src/tables/program.rs
index c8c66b777..6acb8991c 100644
--- a/ceno_zkvm/src/tables/program.rs
+++ b/ceno_zkvm/src/tables/program.rs
@@ -135,8 +135,8 @@ impl<E: ExtensionField> TableCircuit<E> for ProgramTableCircuit<E> {
         fixed
             .par_iter_mut()
             .with_min_len(MIN_PAR_SIZE)
-            .enumerate()
-            .for_each(|(i, row)| {
+            .zip((0..num_instructions).into_par_iter())
+            .for_each(|(row, i)| {
                 let pc = pc_start + (i * WORD_SIZE) as u32;
                 let insn = DecodedInstruction::new(program[i]);
                 let values = InsnRecord::from_decoded(pc, &insn);

From 5e99ca469a12c25c9a729b07177a3fb02faf3b40 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Nicolas?= <info@nau.re>
Date: Fri, 13 Sep 2024 19:29:46 +0200
Subject: [PATCH 10/12] program-table: Support for MockProver

---
 ceno_zkvm/src/instructions/riscv/addsub.rs | 20 +++++-----
 ceno_zkvm/src/scheme/mock_prover.rs        | 46 +++++++++++++++++++---
 ceno_zkvm/src/witness.rs                   |  6 ++-
 3 files changed, 57 insertions(+), 15 deletions(-)

diff --git a/ceno_zkvm/src/instructions/riscv/addsub.rs b/ceno_zkvm/src/instructions/riscv/addsub.rs
index fa4d62368..6e35d4895 100644
--- a/ceno_zkvm/src/instructions/riscv/addsub.rs
+++ b/ceno_zkvm/src/instructions/riscv/addsub.rs
@@ -300,15 +300,15 @@ impl<E: ExtensionField> Instruction<E> for SubInstruction<E> {
 
 #[cfg(test)]
 mod test {
-    use ceno_emul::{Change, ReadOp, StepRecord, WriteOp};
+    use ceno_emul::{Change, ReadOp, StepRecord, WriteOp, CENO_PLATFORM};
     use goldilocks::GoldilocksExt2;
     use itertools::Itertools;
     use multilinear_extensions::mle::IntoMLEs;
 
     use crate::{
         circuit_builder::{CircuitBuilder, ConstraintSystem},
-        instructions::Instruction,
-        scheme::mock_prover::MockProver,
+        instructions::{riscv::constants::PC_STEP_SIZE, Instruction},
+        scheme::mock_prover::{MockProver, MOCK_PC_ADD},
     };
 
     use super::AddInstruction;
@@ -333,18 +333,19 @@ mod test {
             &config,
             cb.cs.num_witin as usize,
             vec![StepRecord {
+                pc: Change::new(MOCK_PC_ADD, MOCK_PC_ADD + PC_STEP_SIZE),
                 rs1: Some(ReadOp {
-                    addr: 2.into(),
+                    addr: CENO_PLATFORM.register_vma(2).into(),
                     value: 11u32,
                     previous_cycle: 0,
                 }),
                 rs2: Some(ReadOp {
-                    addr: 3.into(),
+                    addr: CENO_PLATFORM.register_vma(3).into(),
                     value: 0xfffffffeu32,
                     previous_cycle: 0,
                 }),
                 rd: Some(WriteOp {
-                    addr: 4.into(),
+                    addr: CENO_PLATFORM.register_vma(4).into(),
                     value: Change {
                         before: 0u32,
                         after: 9u32,
@@ -388,18 +389,19 @@ mod test {
             &config,
             cb.cs.num_witin as usize,
             vec![StepRecord {
+                pc: Change::new(MOCK_PC_ADD, MOCK_PC_ADD + PC_STEP_SIZE),
                 rs1: Some(ReadOp {
-                    addr: 2.into(),
+                    addr: CENO_PLATFORM.register_vma(2).into(),
                     value: u32::MAX - 1,
                     previous_cycle: 0,
                 }),
                 rs2: Some(ReadOp {
-                    addr: 3.into(),
+                    addr: CENO_PLATFORM.register_vma(3).into(),
                     value: u32::MAX - 1,
                     previous_cycle: 0,
                 }),
                 rd: Some(WriteOp {
-                    addr: 4.into(),
+                    addr: CENO_PLATFORM.register_vma(4).into(),
                     value: Change {
                         before: 0u32,
                         after: u32::MAX - 2,
diff --git a/ceno_zkvm/src/scheme/mock_prover.rs b/ceno_zkvm/src/scheme/mock_prover.rs
index 3b9a8afcc..e8e403836 100644
--- a/ceno_zkvm/src/scheme/mock_prover.rs
+++ b/ceno_zkvm/src/scheme/mock_prover.rs
@@ -1,10 +1,13 @@
 use super::utils::{eval_by_expr, wit_infer_by_expr};
 use crate::{
-    circuit_builder::CircuitBuilder,
+    circuit_builder::{CircuitBuilder, ConstraintSystem},
     expression::Expression,
+    scheme::utils::eval_by_expr_with_fixed,
     structs::{ROMType, WitnessId},
+    tables::{ProgramTableCircuit, TableCircuit},
 };
 use ark_std::test_rng;
+use ceno_emul::{ByteAddr, CENO_PLATFORM};
 use ff_ext::ExtensionField;
 use generic_static::StaticTypeMap;
 use goldilocks::SmallField;
@@ -12,6 +15,15 @@ use itertools::Itertools;
 use multilinear_extensions::virtual_poly_v2::ArcMultilinearExtension;
 use std::{collections::HashSet, hash::Hash, marker::PhantomData, ops::Neg, sync::OnceLock};
 
+/// The program baked in the MockProver.
+/// TODO: Make this a parameter?
+const MOCK_PROGRAM: &[u32] = &[
+    // add x4, x2, x3
+    3 << 20 | 2 << 15 | 4 << 7 | 0x33,
+];
+// Addresses of particular instructions in the mock program.
+pub const MOCK_PC_ADD: ByteAddr = ByteAddr(CENO_PLATFORM.pc_start());
+
 #[allow(clippy::enum_variant_names)]
 #[derive(Debug, PartialEq, Clone)]
 pub(crate) enum MockProverError<E: ExtensionField> {
@@ -297,6 +309,29 @@ fn load_tables<E: ExtensionField>(cb: &CircuitBuilder<E>, challenge: [E; 2]) ->
         }
     }
 
+    fn load_program_table<E: ExtensionField>(
+        t_vec: &mut Vec<Vec<u8>>,
+        _cb: &CircuitBuilder<E>,
+        challenge: [E; 2],
+    ) {
+        let mut cs = ConstraintSystem::<E>::new(|| "mock_program");
+        let mut cb = CircuitBuilder::new(&mut cs);
+        let config = ProgramTableCircuit::construct_circuit(&mut cb).unwrap();
+        let fixed =
+            ProgramTableCircuit::<E>::generate_fixed_traces(&config, cs.num_fixed, MOCK_PROGRAM);
+        for table_expr in &cs.lk_table_expressions {
+            for row in fixed.iter_rows() {
+                // TODO: Find a better way to obtain the row content.
+                let row = row
+                    .iter()
+                    .map(|v| unsafe { v.clone().assume_init() }.into())
+                    .collect::<Vec<_>>();
+                let rlc_record = eval_by_expr_with_fixed(&row, &[], &challenge, &table_expr.values);
+                t_vec.push(rlc_record.to_repr().as_ref().to_vec());
+            }
+        }
+    }
+
     let mut table_vec = vec![];
     // TODO load more tables here
     load_u5_table(&mut table_vec, cb, challenge);
@@ -304,6 +339,7 @@ fn load_tables<E: ExtensionField>(cb: &CircuitBuilder<E>, challenge: [E; 2]) ->
     load_lt_table(&mut table_vec, cb, challenge);
     load_and_table(&mut table_vec, cb, challenge);
     load_ltu_table(&mut table_vec, cb, challenge);
+    load_program_table(&mut table_vec, cb, challenge);
     HashSet::from_iter(table_vec)
 }
 
@@ -336,7 +372,7 @@ fn load_once_tables<E: ExtensionField + 'static + Sync + Send>(
 
 impl<'a, E: ExtensionField + Hash> MockProver<E> {
     pub fn run_with_challenge(
-        cb: &mut CircuitBuilder<E>,
+        cb: &CircuitBuilder<E>,
         wits_in: &[ArcMultilinearExtension<'a, E>],
         challenge: [E; 2],
     ) -> Result<(), Vec<MockProverError<E>>> {
@@ -344,14 +380,14 @@ impl<'a, E: ExtensionField + Hash> MockProver<E> {
     }
 
     pub fn run(
-        cb: &mut CircuitBuilder<E>,
+        cb: &CircuitBuilder<E>,
         wits_in: &[ArcMultilinearExtension<'a, E>],
     ) -> Result<(), Vec<MockProverError<E>>> {
         Self::run_maybe_challenge(cb, wits_in, None)
     }
 
     fn run_maybe_challenge(
-        cb: &mut CircuitBuilder<E>,
+        cb: &CircuitBuilder<E>,
         wits_in: &[ArcMultilinearExtension<'a, E>],
         challenge: Option<[E; 2]>,
     ) -> Result<(), Vec<MockProverError<E>>> {
@@ -451,7 +487,7 @@ impl<'a, E: ExtensionField + Hash> MockProver<E> {
     }
 
     pub fn assert_satisfied(
-        cb: &mut CircuitBuilder<E>,
+        cb: &CircuitBuilder<E>,
         wits_in: &[ArcMultilinearExtension<'a, E>],
         challenge: Option<[E; 2]>,
     ) {
diff --git a/ceno_zkvm/src/witness.rs b/ceno_zkvm/src/witness.rs
index 0f9e0d406..263b6645c 100644
--- a/ceno_zkvm/src/witness.rs
+++ b/ceno_zkvm/src/witness.rs
@@ -3,7 +3,7 @@ use std::{
     cell::RefCell,
     collections::HashMap,
     mem::{self, MaybeUninit},
-    slice::ChunksMut,
+    slice::{Chunks, ChunksMut},
     sync::Arc,
 };
 
@@ -52,6 +52,10 @@ impl<T: Sized + Sync + Clone + Send> RowMajorMatrix<T> {
         self.values.len() / self.num_col - self.num_padding_rows
     }
 
+    pub fn iter_rows(&self) -> Chunks<MaybeUninit<T>> {
+        self.values.chunks(self.num_col)
+    }
+
     pub fn iter_mut(&mut self) -> ChunksMut<MaybeUninit<T>> {
         self.values.chunks_mut(self.num_col)
     }

From 4c19e864a5c686508c4de5ea991ad18477b3dc87 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Nicolas?= <info@nau.re>
Date: Sat, 14 Sep 2024 12:51:33 +0200
Subject: [PATCH 11/12] program-table: fix after merge

---
 ceno_zkvm/src/instructions/riscv/addsub.rs | 39 ++--------------------
 1 file changed, 2 insertions(+), 37 deletions(-)

diff --git a/ceno_zkvm/src/instructions/riscv/addsub.rs b/ceno_zkvm/src/instructions/riscv/addsub.rs
index 99d91f808..44b672f2e 100644
--- a/ceno_zkvm/src/instructions/riscv/addsub.rs
+++ b/ceno_zkvm/src/instructions/riscv/addsub.rs
@@ -160,6 +160,7 @@ fn add_sub_assignment<E: ExtensionField, const IS_ADD: bool>(
     lk_multiplicity: &mut LkMultiplicity,
     step: &StepRecord,
 ) -> Result<(), ZKVMError> {
+    lk_multiplicity.fetch(step.pc().before.0);
     set_val!(instance, config.pc, step.pc().before.0 as u64);
     set_val!(instance, config.ts, step.cycle());
     let addend_1 = UIntValue::new_unchecked(step.rs2().unwrap().value);
@@ -277,43 +278,7 @@ impl<E: ExtensionField> Instruction<E> for SubInstruction<E> {
         lk_multiplicity: &mut LkMultiplicity,
         step: &StepRecord,
     ) -> Result<(), ZKVMError> {
-        lk_multiplicity.fetch(step.pc().before.0);
-        set_val!(instance, config.pc, step.pc().before.0 as u64);
-
-        // TODO use field from step
-        set_val!(instance, config.ts, 2);
-        config.prev_rd_value.wits_in().map(|prev_rd_value| {
-            set_val!(instance, prev_rd_value[0], 4);
-            set_val!(instance, prev_rd_value[1], 4);
-        });
-        config.addend_0.wits_in().map(|addend_0| {
-            set_val!(instance, addend_0[0], 4);
-            set_val!(instance, addend_0[1], 4);
-        });
-        config.addend_1.wits_in().map(|addend_1| {
-            set_val!(instance, addend_1[0], 4);
-            set_val!(instance, addend_1[1], 4);
-        });
-        // TODO #174
-        config.outcome.carries.as_ref().map(|carry| {
-            set_val!(instance, carry[0], 4);
-            set_val!(instance, carry[1], 0);
-        });
-
-        // Register IDs.
-        for (reg_col, reg_id) in [
-            (config.rs1_id, step.rs1().unwrap().register_index()),
-            (config.rs2_id, step.rs2().unwrap().register_index()),
-            (config.rd_id, step.rd().unwrap().register_index()),
-        ] {
-            set_val!(instance, reg_col, reg_id as u64);
-        }
-
-        // TODO #167
-        set_val!(instance, config.prev_rs1_ts, 2);
-        set_val!(instance, config.prev_rs2_ts, 2);
-        set_val!(instance, config.prev_rd_ts, 2);
-        Ok(())
+        add_sub_assignment::<_, false>(config, instance, lk_multiplicity, step)
     }
 }
 

From a5dde7c435adf9119121c18e3ec6b8a2f1c51ddf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Aur=C3=A9lien=20Nicolas?= <info@nau.re>
Date: Sat, 14 Sep 2024 13:04:58 +0200
Subject: [PATCH 12/12] program-table: support SUB in MockProver

---
 ceno_zkvm/src/instructions/riscv/addsub.rs | 10 +++++++---
 ceno_zkvm/src/scheme/mock_prover.rs        |  7 +++++--
 2 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/ceno_zkvm/src/instructions/riscv/addsub.rs b/ceno_zkvm/src/instructions/riscv/addsub.rs
index 44b672f2e..993cc5b6a 100644
--- a/ceno_zkvm/src/instructions/riscv/addsub.rs
+++ b/ceno_zkvm/src/instructions/riscv/addsub.rs
@@ -292,7 +292,7 @@ mod test {
     use crate::{
         circuit_builder::{CircuitBuilder, ConstraintSystem},
         instructions::{riscv::constants::PC_STEP_SIZE, Instruction},
-        scheme::mock_prover::{MockProver, MOCK_PC_ADD},
+        scheme::mock_prover::{MockProver, MOCK_PC_ADD, MOCK_PC_SUB, MOCK_PROGRAM},
     };
 
     use super::{AddInstruction, SubInstruction};
@@ -319,6 +319,7 @@ mod test {
             vec![StepRecord {
                 cycle: 3,
                 pc: Change::new(MOCK_PC_ADD, MOCK_PC_ADD + PC_STEP_SIZE),
+                insn_code: MOCK_PROGRAM[0],
                 rs1: Some(ReadOp {
                     addr: CENO_PLATFORM.register_vma(2).into(),
                     value: 11u32,
@@ -376,6 +377,7 @@ mod test {
             vec![StepRecord {
                 cycle: 3,
                 pc: Change::new(MOCK_PC_ADD, MOCK_PC_ADD + PC_STEP_SIZE),
+                insn_code: MOCK_PROGRAM[0],
                 rs1: Some(ReadOp {
                     addr: CENO_PLATFORM.register_vma(2).into(),
                     value: u32::MAX - 1,
@@ -432,7 +434,8 @@ mod test {
             cb.cs.num_witin as usize,
             vec![StepRecord {
                 cycle: 3,
-                pc: Change::new(MOCK_PC_ADD, MOCK_PC_ADD + PC_STEP_SIZE), // TODO
+                pc: Change::new(MOCK_PC_SUB, MOCK_PC_SUB + PC_STEP_SIZE),
+                insn_code: MOCK_PROGRAM[1],
                 rs1: Some(ReadOp {
                     addr: CENO_PLATFORM.register_vma(2).into(),
                     value: 11u32,
@@ -489,7 +492,8 @@ mod test {
             cb.cs.num_witin as usize,
             vec![StepRecord {
                 cycle: 3,
-                pc: Change::new(MOCK_PC_ADD, MOCK_PC_ADD + PC_STEP_SIZE), // TODO
+                pc: Change::new(MOCK_PC_SUB, MOCK_PC_SUB + PC_STEP_SIZE),
+                insn_code: MOCK_PROGRAM[1],
                 rs1: Some(ReadOp {
                     addr: CENO_PLATFORM.register_vma(2).into(),
                     value: 3u32,
diff --git a/ceno_zkvm/src/scheme/mock_prover.rs b/ceno_zkvm/src/scheme/mock_prover.rs
index 8047b155d..8dea82b70 100644
--- a/ceno_zkvm/src/scheme/mock_prover.rs
+++ b/ceno_zkvm/src/scheme/mock_prover.rs
@@ -17,12 +17,15 @@ use std::{collections::HashSet, hash::Hash, marker::PhantomData, ops::Neg, sync:
 
 /// The program baked in the MockProver.
 /// TODO: Make this a parameter?
-const MOCK_PROGRAM: &[u32] = &[
+pub const MOCK_PROGRAM: &[u32] = &[
     // add x4, x2, x3
-    3 << 20 | 2 << 15 | 4 << 7 | 0x33,
+    0x00 << 25 | 3 << 20 | 2 << 15 | 4 << 7 | 0x33,
+    // sub  x4, x2, x3
+    0x20 << 25 | 3 << 20 | 2 << 15 | 4 << 7 | 0x33,
 ];
 // Addresses of particular instructions in the mock program.
 pub const MOCK_PC_ADD: ByteAddr = ByteAddr(CENO_PLATFORM.pc_start());
+pub const MOCK_PC_SUB: ByteAddr = ByteAddr(CENO_PLATFORM.pc_start() + 4);
 
 #[allow(clippy::enum_variant_names)]
 #[derive(Debug, PartialEq, Clone)]