From f6e655e41c44cac887dba7e15f96efca65874aab Mon Sep 17 00:00:00 2001
From: Benjamin Kyd <benjamin.kyd@intel.com>
Date: Thu, 13 Jul 2023 01:04:10 +0100
Subject: [PATCH] traits are so ergonomic and an awesome way to implement imm
 encoding

---
 src/ext/decode.rs   |  1 +
 src/ext/encoding.rs | 71 +++++++++++++++++++++++++++++++++++++++++++--
 src/ext/i/mod.rs    |  7 +++--
 3 files changed, 73 insertions(+), 6 deletions(-)

diff --git a/src/ext/decode.rs b/src/ext/decode.rs
index 5588b6e..e27d45b 100644
--- a/src/ext/decode.rs
+++ b/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,
diff --git a/src/ext/encoding.rs b/src/ext/encoding.rs
index 31bfc9f..33e3d0e 100644
--- a/src/ext/encoding.rs
+++ b/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>,
 }
diff --git a/src/ext/i/mod.rs b/src/ext/i/mod.rs
index de5c039..8a6bd29 100644
--- a/src/ext/i/mod.rs
+++ b/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),
 }
-