ZJIT: Handle ISEQ moves (#14250)

* ZJIT: Handle ISEQ moves in IseqCall

* ZJIT: Handle ISEQ moves in Invariants

* Let gen_iseq_call take a reference

* Avoid unneeded iter()

3 files changed, 89 insertions, 32 deletions

diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 467a45846e..d2810cddb7 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -1,4 +1,4 @@
-use std::cell::Cell;
+use std::cell::{Cell, RefCell};
 use std::rc::Rc;
 use std::ffi::{c_int, c_long, c_void};
 
@@ -27,7 +27,7 @@ struct JITState {
     labels: Vec<Option<Target>>,
 
     /// ISEQ calls that need to be compiled later
-    iseq_calls: Vec<Rc<IseqCall>>,
+    iseq_calls: Vec<Rc<RefCell<IseqCall>>>,
 
     /// The number of bytes allocated for basic block arguments spilled onto the C stack
     c_stack_slots: usize,
@@ -126,13 +126,14 @@ fn gen_iseq_entry_point_body(cb: &mut CodeBlock, iseq: IseqPtr) -> Option<CodePt
     };
 
     // Stub callee ISEQs for JIT-to-JIT calls
-    for iseq_call in jit.iseq_calls.into_iter() {
+    for iseq_call in jit.iseq_calls.iter() {
         gen_iseq_call(cb, iseq, iseq_call)?;
     }
 
     // Remember the block address to reuse it later
     let payload = get_or_create_iseq_payload(iseq);
     payload.status = IseqStatus::Compiled(start_ptr);
+    payload.iseq_calls.extend(jit.iseq_calls);
     append_gc_offsets(iseq, &gc_offsets);
 
     // Return a JIT code address
@@ -140,19 +141,19 @@ fn gen_iseq_entry_point_body(cb: &mut CodeBlock, iseq: IseqPtr) -> Option<CodePt
 }
 
 /// Stub a branch for a JIT-to-JIT call
-fn gen_iseq_call(cb: &mut CodeBlock, caller_iseq: IseqPtr, iseq_call: Rc<IseqCall>) -> Option<()> {
+fn gen_iseq_call(cb: &mut CodeBlock, caller_iseq: IseqPtr, iseq_call: &Rc<RefCell<IseqCall>>) -> Option<()> {
     // Compile a function stub
     let Some(stub_ptr) = gen_function_stub(cb, iseq_call.clone()) else {
         // Failed to compile the stub. Bail out of compiling the caller ISEQ.
         debug!("Failed to compile iseq: could not compile stub: {} -> {}",
-               iseq_get_location(caller_iseq, 0), iseq_get_location(iseq_call.iseq, 0));
+               iseq_get_location(caller_iseq, 0), iseq_get_location(iseq_call.borrow().iseq, 0));
         return None;
     };
 
     // Update the JIT-to-JIT call to call the stub
     let stub_addr = stub_ptr.raw_ptr(cb);
-    iseq_call.regenerate(cb, |asm| {
-        asm_comment!(asm, "call function stub: {}", iseq_get_location(iseq_call.iseq, 0));
+    iseq_call.borrow_mut().regenerate(cb, |asm| {
+        asm_comment!(asm, "call function stub: {}", iseq_get_location(iseq_call.borrow().iseq, 0));
         asm.ccall(stub_addr, vec![]);
     });
     Some(())
@@ -205,7 +206,7 @@ fn gen_entry(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function, function_pt
 }
 
 /// Compile an ISEQ into machine code
-fn gen_iseq(cb: &mut CodeBlock, iseq: IseqPtr) -> Option<(CodePtr, Vec<Rc<IseqCall>>)> {
+fn gen_iseq(cb: &mut CodeBlock, iseq: IseqPtr) -> Option<(CodePtr, Vec<Rc<RefCell<IseqCall>>>)> {
     // Return an existing pointer if it's already compiled
     let payload = get_or_create_iseq_payload(iseq);
     match payload.status {
@@ -227,6 +228,7 @@ fn gen_iseq(cb: &mut CodeBlock, iseq: IseqPtr) -> Option<(CodePtr, Vec<Rc<IseqCa
     let result = gen_function(cb, iseq, &function);
     if let Some((start_ptr, gc_offsets, jit)) = result {
         payload.status = IseqStatus::Compiled(start_ptr);
+        payload.iseq_calls.extend(jit.iseq_calls.clone());
         append_gc_offsets(iseq, &gc_offsets);
         Some((start_ptr, jit.iseq_calls))
     } else {
@@ -1470,9 +1472,9 @@ c_callable! {
         with_vm_lock(src_loc!(), || {
             // gen_push_frame() doesn't set PC and SP, so we need to set them before exit.
             // function_stub_hit_body() may allocate and call gc_validate_pc(), so we always set PC.
-            let iseq_call = unsafe { Rc::from_raw(iseq_call_ptr as *const IseqCall) };
+            let iseq_call = unsafe { Rc::from_raw(iseq_call_ptr as *const RefCell<IseqCall>) };
             let cfp = unsafe { get_ec_cfp(ec) };
-            let pc = unsafe { rb_iseq_pc_at_idx(iseq_call.iseq, 0) }; // TODO: handle opt_pc once supported
+            let pc = unsafe { rb_iseq_pc_at_idx(iseq_call.borrow().iseq, 0) }; // TODO: handle opt_pc once supported
             unsafe { rb_set_cfp_pc(cfp, pc) };
             unsafe { rb_set_cfp_sp(cfp, sp) };
 
@@ -1480,10 +1482,10 @@ c_callable! {
             // TODO: Alan thinks the payload status part of this check can happen without the VM lock, since the whole
             // code path can be made read-only. But you still need the check as is while holding the VM lock in any case.
             let cb = ZJITState::get_code_block();
-            let payload = get_or_create_iseq_payload(iseq_call.iseq);
+            let payload = get_or_create_iseq_payload(iseq_call.borrow().iseq);
             if cb.has_dropped_bytes() || payload.status == IseqStatus::CantCompile {
                 // We'll use this Rc again, so increment the ref count decremented by from_raw.
-                unsafe { Rc::increment_strong_count(iseq_call_ptr as *const IseqCall); }
+                unsafe { Rc::increment_strong_count(iseq_call_ptr as *const RefCell<IseqCall>); }
 
                 // Exit to the interpreter
                 return ZJITState::get_exit_trampoline().raw_ptr(cb);
@@ -1504,22 +1506,22 @@ c_callable! {
 }
 
 /// Compile an ISEQ for a function stub
-fn function_stub_hit_body(cb: &mut CodeBlock, iseq_call: &Rc<IseqCall>) -> Option<CodePtr> {
+fn function_stub_hit_body(cb: &mut CodeBlock, iseq_call: &Rc<RefCell<IseqCall>>) -> Option<CodePtr> {
     // Compile the stubbed ISEQ
-    let Some((code_ptr, iseq_calls)) = gen_iseq(cb, iseq_call.iseq) else {
-        debug!("Failed to compile iseq: gen_iseq failed: {}", iseq_get_location(iseq_call.iseq, 0));
+    let Some((code_ptr, iseq_calls)) = gen_iseq(cb, iseq_call.borrow().iseq) else {
+        debug!("Failed to compile iseq: gen_iseq failed: {}", iseq_get_location(iseq_call.borrow().iseq, 0));
         return None;
     };
 
     // Stub callee ISEQs for JIT-to-JIT calls
-    for callee_iseq_call in iseq_calls.into_iter() {
-        gen_iseq_call(cb, iseq_call.iseq, callee_iseq_call)?;
+    for callee_iseq_call in iseq_calls.iter() {
+        gen_iseq_call(cb, iseq_call.borrow().iseq, callee_iseq_call)?;
     }
 
     // Update the stub to call the code pointer
     let code_addr = code_ptr.raw_ptr(cb);
-    iseq_call.regenerate(cb, |asm| {
-        asm_comment!(asm, "call compiled function: {}", iseq_get_location(iseq_call.iseq, 0));
+    iseq_call.borrow_mut().regenerate(cb, |asm| {
+        asm_comment!(asm, "call compiled function: {}", iseq_get_location(iseq_call.borrow().iseq, 0));
         asm.ccall(code_addr, vec![]);
     });
 
@@ -1527,9 +1529,9 @@ fn function_stub_hit_body(cb: &mut CodeBlock, iseq_call: &Rc<IseqCall>) -> Optio
 }
 
 /// Compile a stub for an ISEQ called by SendWithoutBlockDirect
-fn gen_function_stub(cb: &mut CodeBlock, iseq_call: Rc<IseqCall>) -> Option<CodePtr> {
+fn gen_function_stub(cb: &mut CodeBlock, iseq_call: Rc<RefCell<IseqCall>>) -> Option<CodePtr> {
     let mut asm = Assembler::new();
-    asm_comment!(asm, "Stub: {}", iseq_get_location(iseq_call.iseq, 0));
+    asm_comment!(asm, "Stub: {}", iseq_get_location(iseq_call.borrow().iseq, 0));
 
     // Call function_stub_hit using the shared trampoline. See `gen_function_stub_hit_trampoline`.
     // Use load_into instead of mov, which is split on arm64, to avoid clobbering ALLOC_REGS.
@@ -1636,7 +1638,7 @@ fn aligned_stack_bytes(num_slots: usize) -> usize {
 
 impl Assembler {
     /// Make a C call while marking the start and end positions for IseqCall
-    fn ccall_with_iseq_call(&mut self, fptr: *const u8, opnds: Vec<Opnd>, iseq_call: &Rc<IseqCall>) -> Opnd {
+    fn ccall_with_iseq_call(&mut self, fptr: *const u8, opnds: Vec<Opnd>, iseq_call: &Rc<RefCell<IseqCall>>) -> Opnd {
         // We need to create our own branch rc objects so that we can move the closure below
         let start_iseq_call = iseq_call.clone();
         let end_iseq_call = iseq_call.clone();
@@ -1645,10 +1647,10 @@ impl Assembler {
             fptr,
             opnds,
             move |code_ptr, _| {
-                start_iseq_call.start_addr.set(Some(code_ptr));
+                start_iseq_call.borrow_mut().start_addr.set(Some(code_ptr));
             },
             move |code_ptr, _| {
-                end_iseq_call.end_addr.set(Some(code_ptr));
+                end_iseq_call.borrow_mut().end_addr.set(Some(code_ptr));
             },
         )
     }
@@ -1656,9 +1658,9 @@ impl Assembler {
 
 /// Store info about a JIT-to-JIT call
 #[derive(Debug)]
-struct IseqCall {
+pub struct IseqCall {
     /// Callee ISEQ that start_addr jumps to
-    iseq: IseqPtr,
+    pub iseq: IseqPtr,
 
     /// Position where the call instruction starts
     start_addr: Cell<Option<CodePtr>>,
@@ -1669,12 +1671,13 @@ struct IseqCall {
 
 impl IseqCall {
     /// Allocate a new IseqCall
-    fn new(iseq: IseqPtr) -> Rc<Self> {
-        Rc::new(IseqCall {
+    fn new(iseq: IseqPtr) -> Rc<RefCell<Self>> {
+        let iseq_call = IseqCall {
             iseq,
             start_addr: Cell::new(None),
             end_addr: Cell::new(None),
-        })
+        };
+        Rc::new(RefCell::new(iseq_call))
     }
 
     /// Regenerate a IseqCall with a given callback
diff --git a/zjit/src/gc.rs b/zjit/src/gc.rs
index 52a036d49e..3462b80232 100644
--- a/zjit/src/gc.rs
+++ b/zjit/src/gc.rs
@@ -1,6 +1,9 @@
 // This module is responsible for marking/moving objects on GC.
 
+use std::cell::RefCell;
+use std::rc::Rc;
 use std::{ffi::c_void, ops::Range};
+use crate::codegen::IseqCall;
 use crate::{cruby::*, profile::IseqProfile, state::ZJITState, stats::with_time_stat, virtualmem::CodePtr};
 use crate::stats::Counter::gc_time_ns;
 
@@ -15,6 +18,9 @@ pub struct IseqPayload {
 
     /// GC offsets of the JIT code. These are the addresses of objects that need to be marked.
     pub gc_offsets: Vec<CodePtr>,
+
+    /// JIT-to-JIT calls in the ISEQ. The IseqPayload's ISEQ is the caller of it.
+    pub iseq_calls: Vec<Rc<RefCell<IseqCall>>>,
 }
 
 impl IseqPayload {
@@ -23,6 +29,7 @@ impl IseqPayload {
             status: IseqStatus::NotCompiled,
             profile: IseqProfile::new(iseq_size),
             gc_offsets: vec![],
+            iseq_calls: vec![],
         }
     }
 }
@@ -112,6 +119,16 @@ pub extern "C" fn rb_zjit_iseq_update_references(payload: *mut c_void) {
     with_time_stat(gc_time_ns, || iseq_update_references(payload));
 }
 
+/// GC callback for updating object references after all object moves
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_root_update_references() {
+    if !ZJITState::has_instance() {
+        return;
+    }
+    let invariants = ZJITState::get_invariants();
+    invariants.update_references();
+}
+
 fn iseq_mark(payload: &IseqPayload) {
     // Mark objects retained by profiling instructions
     payload.profile.each_object(|object| {
@@ -135,10 +152,22 @@ fn iseq_mark(payload: &IseqPayload) {
 /// This is a mirror of [iseq_mark].
 fn iseq_update_references(payload: &mut IseqPayload) {
     // Move objects retained by profiling instructions
-    payload.profile.each_object_mut(|object| {
-        *object = unsafe { rb_gc_location(*object) };
+    payload.profile.each_object_mut(|old_object| {
+        let new_object = unsafe { rb_gc_location(*old_object) };
+        if *old_object != new_object {
+            *old_object = new_object;
+        }
     });
 
+    // Move ISEQ references in IseqCall
+    for iseq_call in payload.iseq_calls.iter_mut() {
+        let old_iseq = iseq_call.borrow().iseq;
+        let new_iseq = unsafe { rb_gc_location(VALUE(old_iseq as usize)) }.0 as IseqPtr;
+        if old_iseq != new_iseq {
+            iseq_call.borrow_mut().iseq = new_iseq;
+        }
+    }
+
     // Move objects baked in JIT code
     let cb = ZJITState::get_code_block();
     for &offset in payload.gc_offsets.iter() {
diff --git a/zjit/src/invariants.rs b/zjit/src/invariants.rs
index 3f291415be..14fea76d1b 100644
--- a/zjit/src/invariants.rs
+++ b/zjit/src/invariants.rs
@@ -1,6 +1,6 @@
 use std::{collections::{HashMap, HashSet}, mem};
 
-use crate::{backend::lir::{asm_comment, Assembler}, cruby::{rb_callable_method_entry_t, ruby_basic_operators, src_loc, with_vm_lock, IseqPtr, RedefinitionFlag, ID}, gc::IseqPayload, hir::Invariant, options::debug, state::{zjit_enabled_p, ZJITState}, virtualmem::CodePtr};
+use crate::{backend::lir::{asm_comment, Assembler}, cruby::{rb_callable_method_entry_t, rb_gc_location, ruby_basic_operators, src_loc, with_vm_lock, IseqPtr, RedefinitionFlag, ID, VALUE}, gc::IseqPayload, hir::Invariant, options::debug, state::{zjit_enabled_p, ZJITState}, virtualmem::CodePtr};
 use crate::stats::with_time_stat;
 use crate::stats::Counter::invalidation_time_ns;
 use crate::gc::remove_gc_offsets;
@@ -56,6 +56,31 @@ pub struct Invariants {
     single_ractor_patch_points: HashSet<PatchPoint>,
 }
 
+impl Invariants {
+    /// Update object references in Invariants
+    pub fn update_references(&mut self) {
+        Self::update_iseq_references(&mut self.ep_escape_iseqs);
+        Self::update_iseq_references(&mut self.no_ep_escape_iseqs);
+    }
+
+    /// Update ISEQ references in a given HashSet<IseqPtr>
+    fn update_iseq_references(iseqs: &mut HashSet<IseqPtr>) {
+        let mut moved: Vec<IseqPtr> = Vec::with_capacity(iseqs.len());
+
+        iseqs.retain(|&old_iseq| {
+            let new_iseq = unsafe { rb_gc_location(VALUE(old_iseq as usize)) }.0 as IseqPtr;
+            if old_iseq != new_iseq {
+                moved.push(new_iseq);
+            }
+            old_iseq == new_iseq
+        });
+
+        for new_iseq in moved {
+            iseqs.insert(new_iseq);
+        }
+    }
+}
+
 /// Called when a basic operator is redefined. Note that all the blocks assuming
 /// the stability of different operators are invalidated together and we don't
 /// do fine-grained tracking.