From ffaece6ca8e00774a3879848869fc09057f83a8f Mon Sep 17 00:00:00 2001
From: KimiWu <leonhartwu@gmail.com>
Date: Thu, 17 Oct 2024 15:25:19 +0800
Subject: [PATCH] apply to J instruction

---
 ceno_emul/src/rv32im.rs                       |  3 +-
 ceno_emul/src/tracer.rs                       |  4 +--
 ceno_zkvm/src/instructions/riscv/jump/test.rs | 30 +++++++++++++------
 3 files changed, 24 insertions(+), 13 deletions(-)

diff --git a/ceno_emul/src/rv32im.rs b/ceno_emul/src/rv32im.rs
index a673aa454..9a95f9cfe 100644
--- a/ceno_emul/src/rv32im.rs
+++ b/ceno_emul/src/rv32im.rs
@@ -256,7 +256,6 @@ impl EncodedInstruction {
         let opcode = kind.codes().opcode;
         let imm_1_4 = (imm >> 1) & 0xf; // skip imm[0]
         let imm_5_10 = (imm >> 5) & 0x3f;
-        println!("{:#b}: {:#b}, {:#b}", imm, imm_1_4, imm_5_10);
         ((imm >> 12) & 1) << 31
             | imm_5_10 << 25
             | rs2 << 20
@@ -269,7 +268,7 @@ impl EncodedInstruction {
     fn encode_j(kind: InsnKind, rd: u32, imm: u32) -> u32 {
         let rd = rd & 0x1f;
         let opcode = kind.codes().opcode;
-        let imm_1_10 = (imm >> 1) & 0x03ff; // 10-bits mask
+        let imm_1_10 = (imm >> 1) & 0x03ff; // 10-bits mask, skip imm[0]
         let imm_12_19 = (imm >> 12) & 0xff;
         ((imm >> 20) & 1) << 31
             | imm_1_10 << 21
diff --git a/ceno_emul/src/tracer.rs b/ceno_emul/src/tracer.rs
index 6f9d623f7..862c15782 100644
--- a/ceno_emul/src/tracer.rs
+++ b/ceno_emul/src/tracer.rs
@@ -127,11 +127,11 @@ impl StepRecord {
     pub fn new_j_instruction(
         cycle: Cycle,
         pc: Change<ByteAddr>,
-        insn_code: Word,
+        insn_code: u32,
         rd: Change<Word>,
         prev_cycle: Cycle,
     ) -> StepRecord {
-        StepRecord::new_insn(cycle, pc, insn_code, None, None, Some(rd), prev_cycle)
+        StepRecord::new_insn2(cycle, pc, insn_code, None, None, Some(rd), prev_cycle)
     }
 
     fn new_insn(
diff --git a/ceno_zkvm/src/instructions/riscv/jump/test.rs b/ceno_zkvm/src/instructions/riscv/jump/test.rs
index 8db48d7b4..85c94f72b 100644
--- a/ceno_zkvm/src/instructions/riscv/jump/test.rs
+++ b/ceno_zkvm/src/instructions/riscv/jump/test.rs
@@ -11,6 +11,16 @@ use crate::{
 
 use super::{AuipcInstruction, JalInstruction, LuiInstruction};
 
+fn imm_j(imm: i32) -> u32 {
+    // imm is 21 bits in J-type
+    const IMM_MAX: i32 = 2i32.pow(21);
+    let imm = if imm.is_negative() {
+        (IMM_MAX + imm) as u32
+    } else {
+        imm as u32
+    };
+    imm
+}
 #[test]
 fn test_opcode_jal() {
     let mut cs = ConstraintSystem::<GoldilocksExt2>::new(|| "riscv");
@@ -26,22 +36,23 @@ fn test_opcode_jal() {
         .unwrap()
         .unwrap();
 
-    let pc_offset: i32 = -4i32;
-    let new_pc: ByteAddr = ByteAddr(MOCK_PC_JAL.0.wrapping_add_signed(pc_offset));
+    let pc_offset: i32 = -8i32;
+    let new_pc: ByteAddr = ByteAddr(MOCK_PC_START.0.wrapping_add_signed(pc_offset));
+    let insn_code = EncodedInstruction::encode(InsnKind::JAL, 0, 0, 4, imm_j(pc_offset));
     let (raw_witin, lkm) = JalInstruction::<GoldilocksExt2>::assign_instances(
         &config,
         cb.cs.num_witin as usize,
         vec![StepRecord::new_j_instruction(
             4,
-            Change::new(MOCK_PC_JAL, new_pc),
-            MOCK_PROGRAM[21],
-            Change::new(0, (MOCK_PC_JAL + PC_STEP_SIZE).into()),
+            Change::new(MOCK_PC_START, new_pc),
+            insn_code,
+            Change::new(0, (MOCK_PC_START + PC_STEP_SIZE).into()),
             0,
         )],
     )
     .unwrap();
 
-    MockProver::assert_satisfied(
+    MockProver::assert_satisfied_with_program(
         &cb,
         &raw_witin
             .de_interleaving()
@@ -49,12 +60,13 @@ fn test_opcode_jal() {
             .into_iter()
             .map(|v| v.into())
             .collect_vec(),
+        &[insn_code],
         None,
         Some(lkm),
     );
 }
 
-fn imm(imm: u32) -> u32 {
+fn imm_u(imm: u32) -> u32 {
     // valid imm is imm[12:31] in U-type
     imm << 12
 }
@@ -73,7 +85,7 @@ fn test_opcode_lui() {
         .unwrap()
         .unwrap();
 
-    let imm_value = imm(0x90005);
+    let imm_value = imm_u(0x90005);
     let insn_code = EncodedInstruction::encode(InsnKind::LUI, 0, 0, 4, imm_value);
     let (raw_witin, lkm) = LuiInstruction::<GoldilocksExt2>::assign_instances(
         &config,
@@ -117,7 +129,7 @@ fn test_opcode_auipc() {
         .unwrap()
         .unwrap();
 
-    let imm_value = imm(0x90005);
+    let imm_value = imm_u(0x90005);
     let insn_code = EncodedInstruction::encode(InsnKind::AUIPC, 0, 0, 4, imm_value);
     let (raw_witin, lkm) = AuipcInstruction::<GoldilocksExt2>::assign_instances(
         &config,