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>

4 files changed, 30 insertions, 5 deletions

diff --git a/yjit/src/asm/x86_64/mod.rs b/yjit/src/asm/x86_64/mod.rs
index 616a54ba7b..30fe4072b4 100644
--- a/yjit/src/asm/x86_64/mod.rs
+++ b/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 {
diff --git a/yjit/src/asm/x86_64/tests.rs b/yjit/src/asm/x86_64/tests.rs
index 57cc080710..c6d49e084d 100644
--- a/yjit/src/asm/x86_64/tests.rs
+++ b/yjit/src/asm/x86_64/tests.rs
@@ -189,6 +189,12 @@ fn test_mov() {
 }
 
 #[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
     check_bytes("b001", |cb| mov(cb, AL, uimm_opnd(1)));
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index dc5f21221d..68cd35574f 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/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())
                     }
                 },
 
diff --git a/yjit/src/cruby.rs b/yjit/src/cruby.rs
index 168443e6f0..d37af38320 100644
--- a/yjit/src/cruby.rs
+++ b/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;