YJIT: Always encode Opnd::Value in 64 bits on x86_64 for GC offsets (#6733)

* YJIT: Always encode Opnd::Value in 64 bits on x86_64

for GC offsets

Co-authored-by: Alan Wu <alansi.xingwu@shopify.com>

* Introduce heap_object_p

* Leave original mov intact

* Remove unneeded branches

* Add a test for movabs

Co-authored-by: Alan Wu <alansi.xingwu@shopify.com>

yjit/src/asm/x86_64/mod.rs

@@ -1034,6 +1034,20 @@ pub fn mov(cb: &mut CodeBlock, dst: X86Opnd, src: X86Opnd) {
     };
 }
 
+/// A variant of mov used for always writing the value in 64 bits for GC offsets.
+pub fn movabs(cb: &mut CodeBlock, dst: X86Opnd, value: u64) {
+    match dst {
+        X86Opnd::Reg(reg) => {
+            assert_eq!(reg.num_bits, 64);
+            write_rex(cb, true, 0, 0, reg.reg_no);
+
+            write_opcode(cb, 0xb8, reg);
+            cb.write_int(value, 64);
+        },
+        _ => unreachable!()
+    }
+}
+
 /// movsx - Move with sign extension (signed integers)
 pub fn movsx(cb: &mut CodeBlock, dst: X86Opnd, src: X86Opnd) {
     if let X86Opnd::Reg(_dst_reg) = dst {

yjit/src/asm/x86_64/tests.rs

@@ -188,6 +188,12 @@ fn test_mov() {
     check_bytes("48c742f8f4ffffff", |cb| mov(cb, mem_opnd(64, RDX, -8), imm_opnd(-12)));
 }
 
+#[test]
+fn test_movabs() {
+    check_bytes("49b83400000000000000", |cb| movabs(cb, R8, 0x34));
+    check_bytes("49b80000008000000000", |cb| movabs(cb, R8, 0x80000000));
+}
+
 #[test]
 fn test_mov_unsigned() {
     // MOV AL, imm8

yjit/src/backend/x86_64/mod.rs

@@ -465,15 +465,15 @@ impl Assembler
                 // This assumes only load instructions can contain references to GC'd Value operands
                 Insn::Load { opnd, out } |
                 Insn::LoadInto { dest: out, opnd } => {
-                    mov(cb, out.into(), opnd.into());
-
-                    // If the value being loaded is a heap object
-                    if let Opnd::Value(val) = opnd {
-                        if !val.special_const_p() {
+                    match opnd {
+                        Opnd::Value(val) if val.heap_object_p() => {
+                            // Using movabs because mov might write value in 32 bits
+                            movabs(cb, out.into(), val.0 as _);
                             // The pointer immediate is encoded as the last part of the mov written out
                             let ptr_offset: u32 = (cb.get_write_pos() as u32) - (SIZEOF_VALUE as u32);
                             insn_gc_offsets.push(ptr_offset);
                         }
+                        _ => mov(cb, out.into(), opnd.into())
                     }
                 },
 

yjit/src/cruby.rs

@@ -333,6 +333,11 @@ impl VALUE {
         self.immediate_p() || !self.test()
     }
 
+    /// Return true if the value is a heap object
+    pub fn heap_object_p(self) -> bool {
+        !self.special_const_p()
+    }
+
     /// Return true if the value is a Ruby Fixnum (immediate-size integer)
     pub fn fixnum_p(self) -> bool {
         let VALUE(cval) = self;