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

1 files changed, 14 insertions, 4 deletions

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 });
                     }
-                     _ => {},
+                    _ => {},
                 }
             }