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

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/<env> 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)
ruby · XrXr · Jul 22, 2025 · Jul 22, 2025 · Jul 22, 2025 · e2ab7e10a4f61cbd58a0b1cf7f7141b9eda3e8b0
commit e2ab7e10a4f61cbd58a0b1cf7f7141b9eda3e8b0
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,20 @@ 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);
+    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.
+        _ => {
+            let local_index = c_int::try_from(local_ep_offset).ok().and_then(|idx| idx.checked_mul(-1))?;
+            asm_ccall!(asm, rb_vm_env_write, ep, local_index.into(), val);
+        }
+    }
     Some(())
 }