ZJIT: Use rb_vm_env_write() for hir::Insn::SetLocal

XrXr · XrXr · commit ad84032800ba · 2025-07-22T17:03:15.000-04:00
We weren't firing write barriers before when writing to imemo/env
objects. Wbcheck caught this with test/ruby/test_refinement.rb:

    ruby -v: ruby 3.5.0dev (2025-07-22T17:05:58Z wbcheck 2569a80954) +ZJIT dev +PRISM +GC[wbcheck] [x86_64-linux]
    WBCHECK ERROR: Missed write barrier detected!
      Parent object: 0x558de9f4e6e0 (wb_protected: true)
        rb_obj_info_dump: 0x0000558de9f4e6e0 T_IMEMO/&lt;env&gt;
      Reference counts - snapshot: 3, writebarrier: 0, current: 4, missed: 1
      Missing reference to: 0x558decf37c30
        rb_obj_info_dump: 0x0000558decf37c30 method/UnboundMethod method

    WBCHECK SUMMARY: Found 1 objects with missed write barriers (1 total violations)
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
@@ -1,5 +1,6 @@
 use std::cell::Cell;
 use std::rc::Rc;
+use std::ffi::{c_int};
 
 use crate::asm::Label;
 use crate::backend::current::{Reg, ALLOC_REGS};
@@ -446,8 +447,18 @@ fn gen_getlocal_with_ep(asm: &mut Assembler, local_ep_offset: u32, level: u32) -
 /// can't optimize the level=0 case using the SP register.
 fn gen_setlocal_with_ep(asm: &mut Assembler, val: Opnd, local_ep_offset: u32, level: u32) -> Option<()> {
     let ep = gen_get_ep(asm, level);
-    let offset = -(SIZEOF_VALUE_I32 * i32::try_from(local_ep_offset).ok()?);
-    asm.mov(Opnd::mem(64, ep, offset), val);
+    let local_index = c_int::try_from(local_ep_offset).ok().and_then(|idx| idx.checked_mul(-1))?;
+    match val {
+        // If we're writing a constant, non-heap VALUE, do a raw memory write without
+        // running write barrier.
+        lir::Opnd::Value(const_val) if const_val.special_const_p() => {
+            let offset = -(SIZEOF_VALUE_I32 * i32::try_from(local_ep_offset).ok()?);
+            asm.mov(Opnd::mem(64, ep, offset), val);
+        }
+        // We're potentially writing a reference to an IMEMO/env object,
+        // so take care of the write barrier with a function.
+        _ => { asm_ccall!(asm, rb_vm_env_write, ep, local_index.into(), val); }
+    }
     Some(())
 }
 
