traits are so ergonomic and an awesome way to implement imm encoding

src/ext/decode.rs

@@ -15,6 +15,7 @@ impl DecodeCycle {
         DecodeCycle { extensions: ext }
     }
 
+    // TODO: speed this up by matching based on the opcode field and then fn3
     pub fn decode_exec_inst(
         &self,
         inst: rv32::Word,

src/ext/encoding.rs

@@ -3,6 +3,7 @@ use enum_dispatch::*;
 
 use crate::system::rv32;
 use crate::cpu;
+use crate::helpers;
 
 #[enum_dispatch]
 pub trait Instruction {
@@ -11,6 +12,10 @@ pub trait Instruction {
     fn step(&self, inst: GenInstruction, state: &mut cpu::CPUState);
 }
 
+pub trait ImmediateMode {
+    fn sext_imm(&self) -> rv32::XLen;
+    fn full_imm(&self) -> rv32::XLen;
+}
 
 // Null undecided type
 #[bitfield]
@@ -43,15 +48,75 @@ pub struct IType {
     pub imm: B12,
 }
 
-enum EncodingType {
-    R(RType),
-    I(IType),
+impl ImmediateMode for IType {
+    fn sext_imm(&self) -> rv32::XLen {
+        helpers::sext(self.full_imm(), 12)
+    }
+
+    fn full_imm(&self) -> rv32::XLen {
+        self.imm() as rv32::XLen
+    }
 }
 
+#[bitfield]
+pub struct SType {
+    opcode: B7,
+    imm_4_0: B5,
+    funct3: B3,
+    rs1: B5,
+    rs2: B5,
+    imm_11_5: B7,
+}
+
+// imm[11:5] = inst[31:25], imm[4:0] = inst[11:7]
+impl ImmediateMode for SType {
+    fn sext_imm(&self) -> rv32::XLen {
+        helpers::sext(self.full_imm(), 12)
+    }
+
+    fn full_imm(&self) -> rv32::XLen {
+        ((self.imm_11_5() as rv32::XLen) << 5) | self.imm_4_0() as rv32::XLen
+    }
+}
+
+#[bitfield]
+pub struct BType {
+    opcode: B7,
+    imm_11: B1,
+    imm_4_1: B4,
+    funct3: B3,
+    rs1: B5,
+    rs2: B5,
+    imm_10_5: B6,
+    imm_12: B1,
+}
+
+#[bitfield]
+pub struct UType {
+    opcode: B7,
+    rd: B5,
+    imm: B20,
+}
+
+#[bitfield]
+pub struct JType {
+    opcode: B7,
+    rd: B5,
+    imm_19_12: B8,
+    imm_11: B1,
+    imm_10_1: B10,
+    imm_20: B1,
+}
+
+
 #[repr(align(8))]
 pub union GenInstruction {
     pub inst: rv32::Word,
     pub null: std::mem::ManuallyDrop<NullType>,
     pub R: std::mem::ManuallyDrop<RType>,
     pub I: std::mem::ManuallyDrop<IType>,
+    pub S: std::mem::ManuallyDrop<SType>,
+    pub B: std::mem::ManuallyDrop<BType>,
+    pub U: std::mem::ManuallyDrop<UType>,
+    pub J: std::mem::ManuallyDrop<JType>,
 }

src/ext/i/mod.rs

@@ -3,8 +3,9 @@ use enum_dispatch::*;
 use strum::EnumIter;
 
 use super::encoding::{GenInstruction, Instruction};
-use crate::helpers::sext;
 use crate::cpu;
+use crate::ext::encoding::ImmediateMode;
+use crate::helpers::sext;
 use crate::system::rv32;
 
 #[derive(Default, Copy, Clone)]
@@ -25,7 +26,8 @@ impl Instruction for ADDI {
     fn step(&self, inst: GenInstruction, state: &mut cpu::CPUState) {
         println!("VM > Executing ADDI");
         let inst = unsafe { inst.I };
-        state.x[inst.rd() as usize] = state.x[inst.rs1() as usize].wrapping_add(sext(inst.imm() as u32, 32));
+        state.x[inst.rd() as usize] =
+            state.x[inst.rs1() as usize].wrapping_add(inst.sext_imm())
     }
 }
 
@@ -58,4 +60,3 @@ pub enum ExtensionI {
     ADDI(ADDI),
     ADD(ADD),
 }
-