ZJIT: Have JITFrame mark and update its own GC pointers

Add mark() and update_references() methods to JITFrame so GC logic
is encapsulated. Store *mut JITFrame in ZJITState to avoid unsafe
const-to-mut casts in gc.rs.

3 files changed, 25 insertions, 16 deletions

diff --git a/zjit/src/gc.rs b/zjit/src/gc.rs
index 1a0990bc6d..239b71d5f4 100644
--- a/zjit/src/gc.rs
+++ b/zjit/src/gc.rs
@@ -3,7 +3,7 @@
 use std::ptr::null;
 use std::{ffi::c_void, ops::Range};
 use crate::{cruby::*, state::ZJITState, stats::with_time_stat, virtualmem::CodePtr};
-use crate::payload::{IseqPayload, IseqVersionRef, JITFrame, get_or_create_iseq_payload};
+use crate::payload::{IseqPayload, IseqVersionRef, get_or_create_iseq_payload};
 use crate::stats::Counter::gc_time_ns;
 
 /// GC callback for marking GC objects in the per-ISEQ payload.
@@ -95,13 +95,7 @@ pub extern "C" fn rb_zjit_root_update_references() {
     // but JITFrames not on the stack also need their iseq pointers updated
     // because the JIT code will reuse them on the next call.
     for &jit_frame in ZJITState::get_jit_frames().iter() {
-        let old_iseq = unsafe { (*jit_frame).iseq };
-        if !old_iseq.is_null() {
-            let new_iseq = unsafe { rb_gc_location(VALUE::from(old_iseq)) }.as_iseq();
-            if old_iseq != new_iseq {
-                unsafe { (*(jit_frame as *mut JITFrame)).iseq = new_iseq; }
-            }
-        }
+        unsafe { &mut *jit_frame }.update_references();
     }
 }
 
@@ -227,8 +221,6 @@ pub extern "C" fn rb_zjit_root_mark() {
         return;
     }
     for &jit_frame in ZJITState::get_jit_frames().iter() {
-        if !unsafe { (*jit_frame).iseq }.is_null() {
-            unsafe { rb_gc_mark_movable(VALUE::from((*jit_frame).iseq)); }
-        }
+        unsafe { &*jit_frame }.mark();
     }
 }
diff --git a/zjit/src/jit_frame.rs b/zjit/src/jit_frame.rs
index 7626100137..1ab56d8360 100644
--- a/zjit/src/jit_frame.rs
+++ b/zjit/src/jit_frame.rs
@@ -1,4 +1,4 @@
-use crate::cruby::{IseqPtr, VALUE};
+use crate::cruby::{IseqPtr, VALUE, rb_gc_mark_movable, rb_gc_location};
 use crate::state::ZJITState;
 
 #[derive(Debug)]
@@ -21,9 +21,9 @@ impl JITFrame {
     /// Allocate a JITFrame on the heap, register it with ZJITState, and return
     /// a raw pointer that remains valid for the lifetime of the process.
     fn alloc(jit_frame: JITFrame) -> *const Self {
-        let raw_ptr = Box::into_raw(Box::new(jit_frame)) as *const _;
+        let raw_ptr = Box::into_raw(Box::new(jit_frame));
         ZJITState::get_jit_frames().push(raw_ptr);
-        raw_ptr
+        raw_ptr as *const _
     }
 
     /// Create a JITFrame for an ISEQ frame.
@@ -35,6 +35,23 @@ impl JITFrame {
     pub fn new_cfunc() -> *const Self {
         Self::alloc(JITFrame { pc: std::ptr::null(), iseq: std::ptr::null(), materialize_block_code: false })
     }
+
+    /// Mark the iseq pointer for GC. Called from rb_zjit_root_mark.
+    pub fn mark(&self) {
+        if !self.iseq.is_null() {
+            unsafe { rb_gc_mark_movable(VALUE::from(self.iseq)); }
+        }
+    }
+
+    /// Update the iseq pointer after GC compaction.
+    pub fn update_references(&mut self) {
+        if !self.iseq.is_null() {
+            let new_iseq = unsafe { rb_gc_location(VALUE::from(self.iseq)) }.as_iseq();
+            if self.iseq != new_iseq {
+                self.iseq = new_iseq;
+            }
+        }
+    }
 }
 
 #[cfg(test)]
diff --git a/zjit/src/state.rs b/zjit/src/state.rs
index 1f5c568835..813e23dd41 100644
--- a/zjit/src/state.rs
+++ b/zjit/src/state.rs
@@ -73,7 +73,7 @@ pub struct ZJITState {
     perfetto_tracer: Option<PerfettoTracer>,
 
     /// Frame metadata for ISEQ and C calls that are known at compile time
-    jit_frames: Vec<*const JITFrame>,
+    jit_frames: Vec<*mut JITFrame>,
 }
 
 /// Tracks the initialization progress
@@ -191,7 +191,7 @@ impl ZJITState {
         &mut ZJITState::get_instance().invariants
     }
 
-    pub fn get_jit_frames() -> &'static mut Vec<*const JITFrame> {
+    pub fn get_jit_frames() -> &'static mut Vec<*mut JITFrame> {
         &mut ZJITState::get_instance().jit_frames
     }