tests, but at what cost

src/arch/x86_64/encode.rs

@@ -15,7 +15,8 @@ pub struct Code {
 #[derive(Clone, Copy)]
 pub struct Mem {
     pub reg: Reg,
-    pub disp: u32,
+    pub disp: i32,
+    pub width: Width,
 }
 
 #[derive(Clone, Copy)]
@@ -31,8 +32,8 @@ pub enum RegMem {
     Mem(Mem),
 }
 
-pub fn mem(reg: Reg, disp: u32) -> Mem {
-    Mem { reg, disp }
+pub fn mem(reg: Reg, disp: i32, width: Width) -> Mem {
+    Mem { reg, disp, width }
 }
 
 impl Code {
@@ -40,7 +41,7 @@ impl Code {
         let dst = dst.into();
         let src = src.into();
         match dst {
-            RegMem::Reg(dst) => match src {
+            RegMem::Reg(mut dst) => match src {
                 RegImmMem::Reg(src) => {
                     if dst.width() != src.width() {
                         return Err("src and dst are not same width".into());
@@ -57,13 +58,17 @@ impl Code {
                     self.bytes.push(modrm_regs(src, dst));
                 }
                 RegImmMem::Imm(src) => {
+                    let src_width = Width::fit(src);
+                    if src_width > dst.width() {
+                        return Err("immediate cannot fit in register".into());
+                    }
                     self.prefix16(dst);
+                    if src_width <= Width::B32 {
+                        dst.lower64();
+                    }
                     if dst.requires_rex() {
                         self.bytes.push(rex(dst.width(), 0, 0, dst));
                     }
-                    if src > dst.width().max() {
-                        return Err("immediate cannot fit in register".into());
-                    }
                     let opcode = 0xb0 | ((dst.width().gt8() as u8) << 3);
                     self.bytes.push(opcode | dst.base());
                     self.bytes.extend(&src.to_le_bytes()[..dst.width().bytes()]);
@@ -73,14 +78,22 @@ impl Code {
             RegMem::Mem(dst) => match src {
                 RegImmMem::Reg(src) => todo!(),
                 RegImmMem::Imm(src) => {
-                    if src > u32::MAX as u64 {
+                    let src_width = Width::fit(src);
+                    if src_width == Width::B64 {
                         return Err("cannot move 64 bit immediate into memory".into());
                     }
-                    let src = src as u32;
-
-                    self.bytes.extend([rex(1, dst.reg, 0, 0), 0xc7]);
+                    match dst.reg.width() {
+                        Width::B8 | Width::B16 => return Err("invalid register width".into()),
+                        Width::B32 => self.bytes.push(0x67),
+                        Width::B64 => (),
+                    }
+                    self.prefix16(src_width);
+                    if dst.reg.requires_mem_rex() {
+                        self.bytes.push(rex(src_width, 0, 0, dst.reg));
+                    }
+                    self.bytes.push(0xc6 | (src_width != Width::B8) as u8);
                     self.modrm_regdisp(dst.reg, dst.disp);
-                    self.bytes.extend(src.to_le_bytes());
+                    self.bytes.extend(&src.to_le_bytes()[..src_width.bytes()]);
                 }
                 RegImmMem::Mem(_) => return Err("cannot move memory to memory".into()),
             },
@@ -100,18 +113,16 @@ impl Code {
                 Width::B16 => {}
                 _ => return Err("register must be 64 or 16 bit".into()),
             },
-            RegImmMem::Imm(imm) => match imm.try_into() {
-                Ok(imm) => {
-                    const U8: u32 = 2 << 8;
-                    if let 0..U8 = imm {
-                        self.bytes.push(0x6a);
-                        self.bytes.push(imm as u8);
-                    } else {
-                        self.bytes.push(0x68);
-                        self.bytes.extend(imm.to_le_bytes());
-                    }
+            RegImmMem::Imm(imm) => match Width::fit(imm) {
+                Width::B8 => {
+                    self.bytes.push(0x6a);
+                    self.bytes.push(imm as u8);
                 }
-                Err(_) => return Err("immediate must be 32 bit".into()),
+                Width::B16 | Width::B32 => {
+                    self.bytes.push(0x68);
+                    self.bytes.extend((imm as u32).to_le_bytes());
+                }
+                Width::B64 => return Err("immediate must be 32 bit or less".into()),
             },
             RegImmMem::Mem(mem) => todo!(),
         }
@@ -170,18 +181,24 @@ impl Code {
         }
     }
 
-    fn modrm_regdisp(&mut self, reg: Reg, disp: u32) {
-        let disp8 = disp < u8::MAX as u32;
-        let mod_ = if disp8 { 0b01 } else { 0b10 };
+    fn modrm_regdisp(&mut self, reg: Reg, disp: i32) {
+        const I8_MIN: i32 = i8::MIN as i32;
+        const I8_MAX: i32 = i8::MAX as i32;
+        let mod_ = match disp {
+            0 => 0b00,
+            I8_MIN..=I8_MAX => 0b01,
+            _ => 0b10,
+        };
         self.bytes.push(modrm(mod_, 0, reg.base()));
         if reg.val() == rsp.val() {
             // SIB
             self.bytes.push(0x24);
         }
-        if disp8 {
-            self.bytes.push(disp as u8);
-        } else {
-            self.bytes.extend(disp.to_le_bytes());
+        match mod_ {
+            0b00 => (),
+            0b01 => self.bytes.push(disp as u8),
+            0b10 => self.bytes.extend(disp.to_le_bytes()),
+            _ => unreachable!(),
         }
     }
 
@@ -245,6 +262,6 @@ impl From<i64> for RegImmMem {
 
 impl From<i32> for RegImmMem {
     fn from(value: i32) -> Self {
-        Self::Imm(value as u32 as u64)
+        Self::Imm(value as u64)
     }
 }