1 files changed, 11 insertions, 15 deletions

diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index bbc5e800c0..6d6877f273 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -1,6 +1,7 @@
 use crate::asm::x86_64::*;
 use crate::asm::*;
 use crate::codegen::*;
+use crate::virtualmem::CodePtr;
 use crate::cruby::*;
 use crate::options::*;
 use crate::stats::*;
@@ -9,7 +10,6 @@ use core::ffi::c_void;
 use std::cell::*;
 use std::hash::{Hash, Hasher};
 use std::mem;
-use std::mem::size_of;
 use std::rc::{Rc};
 use InsnOpnd::*;
 use TempMapping::*;
@@ -573,24 +573,22 @@ pub extern "C" fn rb_yjit_iseq_update_references(payload: *mut c_void) {
             // Walk over references to objects in generated code.
             for offset in &block.gc_object_offsets {
                 let offset_to_value = offset.as_usize();
-                let value_address: *const u8 = cb.get_ptr(offset_to_value).raw_ptr();
+                let value_code_ptr = cb.get_ptr(offset_to_value);
+                let value_ptr: *const u8 = value_code_ptr.raw_ptr();
                 // Creating an unaligned pointer is well defined unlike in C.
-                let value_address = value_address as *mut VALUE;
+                let value_ptr = value_ptr as *mut VALUE;
 
                 // SAFETY: these point to YJIT's code buffer
-                let object = unsafe { value_address.read_unaligned() };
+                let object = unsafe { value_ptr.read_unaligned() };
                 let new_addr = unsafe { rb_gc_location(object) };
 
-                // Only write when the VALUE moves, to be CoW friendly.
+                // Only write when the VALUE moves, to be copy-on-write friendly.
                 if new_addr != object {
-                    // Possibly unlock the page we need to update
-                    cb.mark_position_writable(offset_to_value);
-
-                    // Object could cross a page boundary, so unlock there as well
-                    cb.mark_position_writable(offset_to_value + size_of::<VALUE>() - 1);
-
-                    // SAFETY: we just made this address writable
-                    unsafe { value_address.write_unaligned(new_addr) };
+                    for (byte_idx, &byte) in new_addr.as_u64().to_le_bytes().iter().enumerate() {
+                        let byte_code_ptr = value_code_ptr.add_bytes(byte_idx);
+                        cb.get_mem().write_byte(byte_code_ptr, byte)
+                            .expect("patching existing code should be within bounds");
+                    }
                 }
             }
         }
@@ -599,8 +597,6 @@ pub extern "C" fn rb_yjit_iseq_update_references(payload: *mut c_void) {
     // Note that we would have returned already if YJIT is off.
     cb.mark_all_executable();
 
-    // I guess we need to make the outlined block executable as well because
-    // we don't split the two at exact page boundaries.
     CodegenGlobals::get_outlined_cb()
         .unwrap()
         .mark_all_executable();