support specific i type imm (#267)

Co-authored-by: Aurélien Nicolas <[email protected]>

ceno_emul/src/rv32im.rs

@@ -305,15 +305,19 @@ impl DecodedInstruction {
         self.func7
     }
 
-    /// Get the decoded immediate, or the funct7 field depending on the instruction format.
+    /// Get the decoded immediate, or 2^shift, or the funct7 field, depending on the instruction format.
     pub fn imm_or_funct7(&self) -> u32 {
-        match self.codes().format {
-            R => self.func7,
-            I => self.imm_i(),
-            S => self.imm_s(),
-            B => self.imm_b(),
-            U => self.imm_u(),
-            J => self.imm_j(),
+        match self.codes() {
+            InsnCodes { format: R, .. } => self.func7,
+            InsnCodes {
+                kind: SLLI | SRLI | SRAI,
+                ..
+            } => 1 << self.rs2(), // decode the shift as a multiplication by 2.pow(rs2)
+            InsnCodes { format: I, .. } => self.imm_i(),
+            InsnCodes { format: S, .. } => self.imm_s(),
+            InsnCodes { format: B, .. } => self.imm_b(),
+            InsnCodes { format: U, .. } => self.imm_u(),
+            InsnCodes { format: J, .. } => self.imm_j(),
         }
     }
 
@@ -359,6 +363,32 @@ impl DecodedInstruction {
     }
 }
 
+#[cfg(test)]
+#[test]
+fn test_decode_imm() {
+    for (i, expected) in [
+        // Example of I-type: ADDI.
+        // imm    | rs1     | funct3      | rd     | opcode
+        (89 << 20 | 1 << 15 | 0b000 << 12 | 1 << 7 | 0x13, 89),
+        // Shifts get a precomputed power of 2: SLLI, SRLI, SRAI.
+        (31 << 20 | 1 << 15 | 0b001 << 12 | 1 << 7 | 0x13, 1 << 31),
+        (31 << 20 | 1 << 15 | 0b101 << 12 | 1 << 7 | 0x13, 1 << 31),
+        (
+            1 << 30 | 31 << 20 | 1 << 15 | 0b101 << 12 | 1 << 7 | 0x13,
+            1 << 31,
+        ),
+        // Example of R-type with funct7: SUB.
+        // funct7     | rs2    | rs1     | funct3      | rd     | opcode
+        (
+            0x20 << 25 | 1 << 20 | 1 << 15 | 0 << 12 | 1 << 7 | 0x33,
+            0x20,
+        ),
+    ] {
+        let imm = DecodedInstruction::new(i).imm_or_funct7();
+        assert_eq!(imm, expected);
+    }
+}
+
 const fn insn(
     format: InsnFormat,
     kind: InsnKind,