ZJIT: Precise GC writebarriers

jhawthorn · tekknolagi · st0012 · jhawthorn · commit cb33f22f5b64 · 2025-07-17T11:50:13.000-07:00
This issues writebarriers for objects added via gc_offsets or by
profiling. This may be slower than writebarrier_remember, but we would
like it to be more debuggable.

Co-authored-by: Max Bernstein &lt;ruby@bernsteinbear.com&gt;
Co-authored-by: Stan Lo &lt;stan001212@gmail.com&gt;
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
@@ -4,7 +4,7 @@ use std::rc::Rc;
 use crate::asm::Label;
 use crate::backend::current::{Reg, ALLOC_REGS};
 use crate::invariants::track_bop_assumption;
-use crate::gc::get_or_create_iseq_payload;
+use crate::gc::{get_or_create_iseq_payload, append_gc_offsets};
 use crate::state::ZJITState;
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::backend::lir::{self, asm_comment, asm_ccall, Assembler, Opnd, SideExitContext, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, NATIVE_STACK_PTR, SP};
@@ -124,7 +124,7 @@ fn gen_iseq_entry_point_body(cb: &mut CodeBlock, iseq: IseqPtr) -> *const u8 {
             // Remember the block address to reuse it later
             let payload = get_or_create_iseq_payload(iseq);
             payload.start_ptr = Some(start_ptr);
-            payload.gc_offsets.extend(gc_offsets);
+            append_gc_offsets(iseq, &gc_offsets);
 
             // Compile an entry point to the JIT code
             (gen_entry(cb, iseq, &function, start_ptr, jit.c_stack_bytes), jit.branch_iseqs)
@@ -193,7 +193,7 @@ fn gen_iseq(cb: &mut CodeBlock, iseq: IseqPtr) -> Option<(CodePtr, Vec<(Rc<Branc
     let result = gen_function(cb, iseq, &function);
     if let Some((start_ptr, gc_offsets, jit)) = result {
         payload.start_ptr = Some(start_ptr);
-        payload.gc_offsets.extend(gc_offsets);
+        append_gc_offsets(iseq, &gc_offsets);
         Some((start_ptr, jit.branch_iseqs))
     } else {
         None
@@ -999,7 +999,7 @@ fn gen_guard_type(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, guard
 
 /// Compile an identity check with a side exit
 fn gen_guard_bit_equals(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, expected: VALUE, state: &FrameState) -> Option<lir::Opnd> {
-    asm.cmp(val, Opnd::UImm(expected.into()));
+    asm.cmp(val, Opnd::Value(expected));
     asm.jnz(side_exit(jit, state, GuardBitEquals(expected))?);
     Some(val)
 }
diff --git a/zjit/src/gc.rs b/zjit/src/gc.rs
@@ -90,6 +90,23 @@ pub extern "C" fn rb_zjit_iseq_mark(payload: *mut c_void) {
     }
 }
 
+/// Append a set of gc_offsets to the iseq's payload
+pub fn append_gc_offsets(iseq: IseqPtr, offsets: &Vec<CodePtr>) {
+    let payload = get_or_create_iseq_payload(iseq);
+    payload.gc_offsets.extend(offsets);
+
+    // Call writebarrier on each newly added value
+    let cb = ZJITState::get_code_block();
+    for &offset in offsets.iter() {
+        let value_ptr: *const u8 = offset.raw_ptr(cb);
+        let value_ptr = value_ptr as *const VALUE;
+        unsafe {
+            let object = value_ptr.read_unaligned();
+            rb_gc_writebarrier(iseq.into(), object);
+        }
+    }
+}
+
 /// GC callback for updating GC objects in the per-iseq payload.
 /// This is a mirror of [rb_zjit_iseq_mark].
 #[unsafe(no_mangle)]
diff --git a/zjit/src/profile.rs b/zjit/src/profile.rs
@@ -88,11 +88,10 @@ fn profile_operands(profiler: &mut Profiler, profile: &mut IseqProfile, n: usize
     for i in 0..n {
         let opnd_type = Type::from_value(profiler.peek_at_stack((n - i - 1) as isize));
         types[i] = types[i].union(opnd_type);
+        if let Some(object) = types[i].gc_object() {
+            unsafe { rb_gc_writebarrier(profiler.iseq.into(), object) };
+        }
     }
-    // In the loop above, we probably added a new reference to the profile through
-    // the VALUE in Type. It's messy and relatively slow to conditionally run a
-    // write barrier for each Type, so just remember to re-mark the iseq.
-    unsafe { rb_gc_writebarrier_remember(profiler.iseq.into()) };
 }
 
 #[derive(Debug)]

Original file line number	Diff line number	Diff line change
`@@ -88,11 +88,10 @@ fn profile_operands(profiler: &mut Profiler, profile: &mut IseqProfile, n: usize`
`88`	`88`	`for i in 0..n {`
`89`	`89`	`let opnd_type = Type::from_value(profiler.peek_at_stack((n - i - 1) as isize));`
`90`	`90`	`types[i] = types[i].union(opnd_type);`
	`91`	`+ if let Some(object) = types[i].gc_object() {`
	`92`	`+ unsafe { rb_gc_writebarrier(profiler.iseq.into(), object) };`
	`93`	`+ }`
`91`	`94`	`}`
`92`		`- // In the loop above, we probably added a new reference to the profile through`
`93`		`- // the VALUE in Type. It's messy and relatively slow to conditionally run a`
`94`		`- // write barrier for each Type, so just remember to re-mark the iseq.`
`95`		`- unsafe { rb_gc_writebarrier_remember(profiler.iseq.into()) };`
`96`	`95`	`}`
`97`	`96`
`98`	`97`	`#[derive(Debug)]`