From 0d462d5b3d57052d47bf2443eb74011f9cbf33a7 Mon Sep 17 00:00:00 2001
From: "sm.wu" <hero78119@gmail.com>
Date: Wed, 30 Oct 2024 01:02:53 +0800
Subject: [PATCH] fix sign extension check on imm and function renaming

---
 ceno_emul/src/rv32im.rs         | 21 ++++++++++++---------
 ceno_zkvm/src/tables/program.rs |  2 +-
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/ceno_emul/src/rv32im.rs b/ceno_emul/src/rv32im.rs
index 096ede73f..cbb2afa88 100644
--- a/ceno_emul/src/rv32im.rs
+++ b/ceno_emul/src/rv32im.rs
@@ -299,7 +299,7 @@ impl DecodedInstruction {
     }
 
     /// Indicate whether the immediate is interpreted as a signed integer, and it is negative.
-    pub fn imm_is_negative(&self) -> bool {
+    pub fn imm_field_is_negative(&self) -> bool {
         match self.codes() {
             InsnCodes { format: R | U, .. } => false,
             InsnCodes {
@@ -323,27 +323,30 @@ impl DecodedInstruction {
         (i.category, i.kind)
     }
 
-    pub fn imm12_sign_ext(&self) -> u32 {
-        self.imm_is_negative() as u32 * 0xfffff000
-    }
-
     pub fn imm_b(&self) -> u32 {
-        self.imm12_sign_ext()
+        (self.top_bit * 0xfffff000)
             | ((self.rd & 1) << 11)
             | ((self.func7 & 0x3f) << 5)
             | (self.rd & 0x1e)
     }
 
+    /// checks if the imm requires sign extension
+    pub fn requires_sign_ext(&self) -> bool {
+        !matches!(self.codes(), InsnCodes { kind: SLTIU, .. })
+    }
+
     pub fn imm_i(&self) -> u32 {
-        self.imm12_sign_ext() | (self.func7 << 5) | self.rs2
+        ((self.requires_sign_ext() as u32) * self.top_bit * 0xfffff000)
+            | (self.func7 << 5)
+            | self.rs2
     }
 
     pub fn imm_s(&self) -> u32 {
-        self.imm12_sign_ext() | (self.func7 << 5) | self.rd
+        (self.top_bit * 0xfffff000) | (self.func7 << 5) | self.rd
     }
 
     pub fn imm_j(&self) -> u32 {
-        self.imm12_sign_ext()
+        (self.top_bit * 0xfff00000)
             | (self.rs1 << 15)
             | (self.func3 << 12)
             | ((self.rs2 & 1) << 11)
diff --git a/ceno_zkvm/src/tables/program.rs b/ceno_zkvm/src/tables/program.rs
index d5b629957..13291aa88 100644
--- a/ceno_zkvm/src/tables/program.rs
+++ b/ceno_zkvm/src/tables/program.rs
@@ -94,7 +94,7 @@ impl InsnRecord<u32> {
     /// Interpret the immediate or funct7 as unsigned or signed depending on the instruction.
     /// Convert negative values from two's complement to field.
     pub fn imm_or_funct7_field<F: SmallField>(insn: &DecodedInstruction) -> F {
-        if insn.imm_is_negative() {
+        if insn.imm_field_is_negative() {
             -F::from(-(insn.imm_or_funct7() as i32) as u64)
         } else {
             F::from(insn.imm_or_funct7() as u64)