YJIT: Add Opnd#with_num_bits to use only 8 bits (#6359)

* YJIT: Add Opnd#sub_opnd to use only 8 bits

* Add with_num_bits and let arm64_split use it

* Add another assertion to with_num_bits

* Use only with_num_bits

7 files changed, 61 insertions, 16 deletions

diff --git a/yjit/src/asm/arm64/inst/load_store.rs b/yjit/src/asm/arm64/inst/load_store.rs
index ea42f2d17f..e877c6de77 100644
--- a/yjit/src/asm/arm64/inst/load_store.rs
+++ b/yjit/src/asm/arm64/inst/load_store.rs
@@ -2,6 +2,7 @@ use super::super::arg::truncate_imm;
 
 /// The size of the operands being operated on.
 enum Size {
+    Size8 = 0b00,
     Size32 = 0b10,
     Size64 = 0b11,
 }
@@ -81,6 +82,12 @@ impl LoadStore {
         Self { rt, rn, idx: Index::None, imm9, opc: Opc::LDR, size: num_bits.into() }
     }
 
+    /// LDURB (load register, byte, unscaled)
+    /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/LDURB--Load-Register-Byte--unscaled--?lang=en
+    pub fn ldurb(rt: u8, rn: u8, imm9: i16) -> Self {
+        Self { rt, rn, idx: Index::None, imm9, opc: Opc::LDR, size: Size::Size8 }
+    }
+
     /// LDURSW (load register, unscaled, signed)
     /// https://developer.arm.com/documentation/ddi0596/2021-12/Base-Instructions/LDURSW--Load-Register-Signed-Word--unscaled--?lang=en
     pub fn ldursw(rt: u8, rn: u8, imm9: i16) -> Self {
@@ -158,6 +165,13 @@ mod tests {
     }
 
     #[test]
+    fn test_ldurb() {
+        let inst = LoadStore::ldurb(0, 1, 0);
+        let result: u32 = inst.into();
+        assert_eq!(0x38400020, result);
+    }
+
+    #[test]
     fn test_ldur_with_imm() {
         let inst = LoadStore::ldur(0, 1, 123, 64);
         let result: u32 = inst.into();
diff --git a/yjit/src/asm/arm64/mod.rs b/yjit/src/asm/arm64/mod.rs
index 420151c6d1..d97452a045 100644
--- a/yjit/src/asm/arm64/mod.rs
+++ b/yjit/src/asm/arm64/mod.rs
@@ -507,6 +507,22 @@ pub fn ldur(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
     cb.write_bytes(&bytes);
 }
 
+/// LDURB - load a byte from memory, zero-extend it, and write it to a register
+pub fn ldurb(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
+    let bytes: [u8; 4] = match (rt, rn) {
+        (A64Opnd::Reg(rt), A64Opnd::Mem(rn)) => {
+            assert!(rt.num_bits == rn.num_bits, "Expected registers to be the same size");
+            assert!(rt.num_bits == 8, "Expected registers to have size 8");
+            assert!(mem_disp_fits_bits(rn.disp), "Expected displacement to be 9 bits or less");
+
+            LoadStore::ldurb(rt.reg_no, rn.base_reg_no, rn.disp as i16).into()
+        },
+        _ => panic!("Invalid operands for LDURB")
+    };
+
+    cb.write_bytes(&bytes);
+}
+
 /// LDURSW - load a 32-bit memory address into a register and sign-extend it
 pub fn ldursw(cb: &mut CodeBlock, rt: A64Opnd, rn: A64Opnd) {
     let bytes: [u8; 4] = match (rt, rn) {
diff --git a/yjit/src/asm/arm64/opnd.rs b/yjit/src/asm/arm64/opnd.rs
index c89481fb03..52b2a84637 100644
--- a/yjit/src/asm/arm64/opnd.rs
+++ b/yjit/src/asm/arm64/opnd.rs
@@ -12,10 +12,8 @@ pub struct A64Reg
 }
 
 impl A64Reg {
-    pub fn sub_reg(&self, num_bits: u8) -> Self {
-        assert!(num_bits == 32 || num_bits == 64);
-        assert!(num_bits <= self.num_bits);
-
+    pub fn with_num_bits(&self, num_bits: u8) -> Self {
+        assert!(num_bits == 8 || num_bits == 16 || num_bits == 32 || num_bits == 64);
         Self { num_bits, reg_no: self.reg_no }
     }
 }
diff --git a/yjit/src/asm/x86_64/mod.rs b/yjit/src/asm/x86_64/mod.rs
index 3f865b82a5..d310e3bf12 100644
--- a/yjit/src/asm/x86_64/mod.rs
+++ b/yjit/src/asm/x86_64/mod.rs
@@ -89,16 +89,13 @@ pub enum X86Opnd
 }
 
 impl X86Reg {
-    pub fn sub_reg(&self, num_bits: u8) -> Self {
+    pub fn with_num_bits(&self, num_bits: u8) -> Self {
         assert!(
             num_bits == 8 ||
             num_bits == 16 ||
             num_bits == 32 ||
             num_bits == 64
         );
-
-        assert!(num_bits <= self.num_bits);
-
         Self {
             num_bits,
             reg_type: self.reg_type,
diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index 446332788a..9d19c65114 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -140,7 +140,14 @@ impl Assembler
                 Opnd::Reg(_) | Opnd::InsnOut { .. } => opnd,
                 Opnd::Mem(_) => {
                     let split_opnd = split_memory_address(asm, opnd);
-                    asm.load(split_opnd)
+                    let out_opnd = asm.load(split_opnd);
+                    // Many Arm insns support only 32-bit or 64-bit operands. asm.load with fewer
+                    // bits zero-extends the value, so it's safe to recognize it as a 32-bit value.
+                    if out_opnd.rm_num_bits() < 32 {
+                        out_opnd.with_num_bits(32).unwrap()
+                    } else {
+                        out_opnd
+                    }
                 },
                 _ => asm.load(opnd)
             }
@@ -747,7 +754,11 @@ impl Assembler
                             emit_load_value(cb, out.into(), imm as u64);
                         },
                         Opnd::Mem(_) => {
-                            ldur(cb, out.into(), opnd.into());
+                            match opnd.rm_num_bits() {
+                                64 | 32 => ldur(cb, out.into(), opnd.into()),
+                                8 => ldurb(cb, out.into(), opnd.into()),
+                                num_bits => panic!("unexpected num_bits: {}", num_bits)
+                            };
                         },
                         Opnd::Value(value) => {
                             // We dont need to check if it's a special const
diff --git a/yjit/src/backend/ir.rs b/yjit/src/backend/ir.rs
index ee6499ff64..609ca8eaf4 100644
--- a/yjit/src/backend/ir.rs
+++ b/yjit/src/backend/ir.rs
@@ -151,6 +151,16 @@ impl Opnd
         }
     }
 
+    pub fn with_num_bits(&self, num_bits: u8) -> Option<Opnd> {
+        assert!(num_bits == 8 || num_bits == 16 || num_bits == 32 || num_bits == 64);
+        match *self {
+            Opnd::Reg(reg) => Some(Opnd::Reg(reg.with_num_bits(num_bits))),
+            Opnd::Mem(Mem { base, disp, .. }) => Some(Opnd::Mem(Mem { base, disp, num_bits })),
+            Opnd::InsnOut { idx, .. } => Some(Opnd::InsnOut { idx, num_bits }),
+            _ => None,
+        }
+    }
+
     /// Get the size in bits for register/memory operands.
     pub fn rm_num_bits(&self) -> u8 {
         self.num_bits().unwrap()
@@ -1052,21 +1062,21 @@ impl Assembler
                 // output operand on this instruction because the live range
                 // extends beyond the index of the instruction.
                 let out = insn.out_opnd_mut().unwrap();
-                *out = Opnd::Reg(out_reg.unwrap().sub_reg(out_num_bits));
+                *out = Opnd::Reg(out_reg.unwrap().with_num_bits(out_num_bits));
             }
 
             // Replace InsnOut operands by their corresponding register
             let mut opnd_iter = insn.opnd_iter_mut();
             while let Some(opnd) = opnd_iter.next() {
                 match *opnd {
-                    Opnd::InsnOut { idx, .. } => {
-                        *opnd = *asm.insns[idx].out_opnd().unwrap();
+                    Opnd::InsnOut { idx, num_bits } => {
+                        *opnd = (*asm.insns[idx].out_opnd().unwrap()).with_num_bits(num_bits).unwrap();
                     },
                     Opnd::Mem(Mem { base: MemBase::InsnOut(idx), disp, num_bits }) => {
                         let base = MemBase::Reg(asm.insns[idx].out_opnd().unwrap().unwrap_reg().reg_no);
                         *opnd = Opnd::Mem(Mem { base, disp, num_bits });
                     }
-                     _ => {},
+                    _ => {},
                 }
             }
 
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 1c48a1b040..ca9ec655d1 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -3477,8 +3477,7 @@ fn jit_guard_known_klass(
 
             asm.comment("guard object is static symbol");
             assert!(RUBY_SPECIAL_SHIFT == 8);
-            let flag_bits = asm.and(obj_opnd, Opnd::UImm(0xf));
-            asm.cmp(flag_bits, Opnd::UImm(RUBY_SYMBOL_FLAG as u64));
+            asm.cmp(obj_opnd.with_num_bits(8).unwrap(), Opnd::UImm(RUBY_SYMBOL_FLAG as u64));
             jit_chain_guard(JCC_JNE, jit, ctx, asm, ocb, max_chain_depth, side_exit);
             ctx.upgrade_opnd_type(insn_opnd, Type::ImmSymbol);
         }