From fa2a6db2e24b09f5b23d692ae39f12fb59047ab3 Mon Sep 17 00:00:00 2001
From: Shadow Cat <shadowcatzero@protonmail.com>
Date: Fri, 12 Jun 2026 22:39:58 -0400
Subject: [PATCH] edit rex fn

---
 src/arch/x86_64/encode.rs | 10 +++----
 src/arch/x86_64/util.rs   | 56 +++++++++++++++++++++++++++------------
 2 files changed, 44 insertions(+), 22 deletions(-)

diff --git a/src/arch/x86_64/encode.rs b/src/arch/x86_64/encode.rs
index b877ee0..96a57e8 100644
--- a/src/arch/x86_64/encode.rs
+++ b/src/arch/x86_64/encode.rs
@@ -40,14 +40,14 @@ impl Code {
                     if dst.width() == Width::B64 && src_width <= Width::B32 && src.0 < 0 {
                         // use different op that sign extends for less bytes
                         self.bytes
-                            .extend([rex(dst.width(), 0, 0, dst), 0xc7, 0xc0 | dst.base()]);
+                            .extend([rex(dst, 0, 0, dst), 0xc7, 0xc0 | dst.base()]);
                         self.imm(src, Width::B32);
                     } else {
                         if src_width <= Width::B32 {
                             dst = dst.lower64();
                         }
                         if dst.requires_rex() {
-                            self.bytes.push(rex(dst.width(), 0, 0, dst));
+                            self.bytes.push(rex(dst, 0, 0, dst));
                         }
                         self.bytes.push(0xb0 | (dst.not8() << 3) | dst.base());
                         self.imm(src, dst.width());
@@ -63,7 +63,7 @@ impl Code {
                     self.prefix32(&src)?;
                     self.prefix16(dst);
                     if dst.requires_rex() | src.reg.gt8() {
-                        self.bytes.push(rex(dst.width(), dst, 0, src.reg));
+                        self.bytes.push(rex(dst, dst, 0, src));
                     }
                     self.bytes.push(0x8a | dst.not8());
                     self.modrm_regdisp(dst, src);
@@ -80,7 +80,7 @@ impl Code {
                     self.prefix32(&dst)?;
                     self.prefix16(src);
                     if src.requires_rex() | dst.reg.gt8() {
-                        self.bytes.push(rex(src.width(), src, 0, dst.reg));
+                        self.bytes.push(rex(src, src, 0, dst));
                     }
                     self.bytes.push(0x88 | src.not8());
                     self.modrm_regdisp(src, dst);
@@ -101,7 +101,7 @@ impl Code {
                     self.prefix32(&dst)?;
                     self.prefix16(encode_width);
                     if dst.reg.requires_mem_rex() || dst.width == Width::B64 {
-                        self.bytes.push(rex(dst.width, 0, 0, dst.reg));
+                        self.bytes.push(rex(dst, 0, 0, dst));
                     }
                     self.bytes.push(0xc6 | encode_width.not8());
                     self.modrm_regdisp(None, dst);
diff --git a/src/arch/x86_64/util.rs b/src/arch/x86_64/util.rs
index e0194a3..5796def 100644
--- a/src/arch/x86_64/util.rs
+++ b/src/arch/x86_64/util.rs
@@ -16,20 +16,13 @@ pub fn modrm(mod_: u8, reg: u8, rm: u8) -> u8 {
 }
 
 #[inline(always)]
-pub fn rex(w: impl RexBit, r: impl RexBit, x: impl RexBit, b: impl RexBit) -> u8 {
-    0b0100_0000 | bit(w, 3) | bit(r, 2) | bit(x, 1) | bit(b, 0)
+pub fn rex(w: impl RexW, r: impl RexBit, x: u8, b: impl RexBit) -> u8 {
+    0b0100_0000 | bit(w.rexw(), 3) | bit(r.rex(), 2) | bit(x.rex(), 1) | bit(b.rex(), 0)
 }
 
 #[inline(always)]
-pub fn bit(val: impl RexBit, pos: u8) -> u8 {
-    (val.rex() as u8) << pos
-}
-
-/// assumes the next instruction is directly after
-pub fn addr_offset(pos: usize, addr: u64) -> [u8; 4] {
-    let pos = (pos + 4) as i32;
-    let offset = addr as i32 - pos;
-    offset.to_le_bytes()
+fn bit(val: bool, pos: u8) -> u8 {
+    (val as u8) << pos
 }
 
 pub trait RexBit {
@@ -42,24 +35,53 @@ impl RexBit for u8 {
     }
 }
 
-impl RexBit for bool {
+impl RexBit for Reg {
     fn rex(self) -> bool {
-        self
+        self.gt8()
     }
 }
 
-impl<R: Into<Reg>> RexBit for R {
+impl RexBit for Mem {
     fn rex(self) -> bool {
-        self.into().gt8()
+        self.reg.rex()
     }
 }
 
-impl RexBit for Width {
-    fn rex(self) -> bool {
+pub trait RexW {
+    fn rexw(self) -> bool;
+}
+
+impl RexW for Width {
+    fn rexw(self) -> bool {
         self == Width::B64
     }
 }
 
+impl RexW for Reg {
+    fn rexw(self) -> bool {
+        self.width().rexw()
+    }
+}
+
+impl RexW for u8 {
+    fn rexw(self) -> bool {
+        self == 1
+    }
+}
+
+impl RexW for Mem {
+    fn rexw(self) -> bool {
+        self.width.rexw()
+    }
+}
+
+/// assumes the next instruction is directly after
+pub fn addr_offset(pos: usize, addr: u64) -> [u8; 4] {
+    let pos = (pos + 4) as i32;
+    let offset = addr as i32 - pos;
+    offset.to_le_bytes()
+}
+
 pub struct SignedHex {
     pub val: i128,
     pub op: bool,