Add guard to YARVINSN_setlocal even though it breaks a test

Author: jacob-shops

zjit/src/codegen.rs

@@ -561,7 +561,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::GetLocal { ep_offset, level, use_sp, .. } => gen_getlocal(asm, ep_offset, level, use_sp),
         &Insn::IsBlockParamModified { level } => gen_is_block_param_modified(asm, level),
         &Insn::GetBlockParam { ep_offset, level, state } => gen_getblockparam(jit, asm, ep_offset, level, &function.frame_state(state)),
-        &Insn::SetLocal { val, ep_offset, level, state } => no_output!(gen_setlocal(asm, opnd!(val), function.type_of(val), ep_offset, level, &function.frame_state(state))),
+        &Insn::SetLocal { val, ep_offset, level } => no_output!(gen_setlocal(jit, asm, opnd!(val), function.type_of(val), ep_offset, level )),
         Insn::GetConstantPath { ic, state } => gen_get_constant_path(jit, asm, *ic, &function.frame_state(*state)),
         Insn::GetClassVar { id, ic, state } => gen_getclassvar(jit, asm, *id, *ic, &function.frame_state(*state)),
         Insn::SetClassVar { id, val, ic, state } => no_output!(gen_setclassvar(jit, asm, *id, opnd!(val), *ic, &function.frame_state(*state))),
@@ -735,32 +735,15 @@ fn gen_getlocal(asm: &mut Assembler, local_ep_offset: u32, level: u32, use_sp: b
 /// Set a local variable from a higher scope or the heap. `local_ep_offset` is in number of VALUEs.
 /// We generate this instruction with level=0 only when the local variable is on the heap, so we
 /// can't optimize the level=0 case using the SP register.
-fn gen_setlocal(asm: &mut Assembler, val: Opnd, val_type: Type, local_ep_offset: u32, level: u32, state: &FrameState) {
+// TODO(Jacob): remove superfluous arguments from function
+fn gen_setlocal(jit: &mut JITState, asm: &mut Assembler, val: Opnd, val_type: Type, local_ep_offset: u32, level: u32) {
     let local_ep_offset = c_int::try_from(local_ep_offset).unwrap_or_else(|_| panic!("Could not convert local_ep_offset {local_ep_offset} to i32"));
     if level > 0 {
         gen_incr_counter(asm, Counter::vm_write_to_parent_iseq_local_count);
     }
     let ep = gen_get_ep(asm, level);
-
-    // When we've proved that we're writing an immediate,
-    // we can skip the write barrier.
-    if val_type.is_immediate() {
-        let offset = -(SIZEOF_VALUE_I32 * local_ep_offset);
-        asm.mov(Opnd::mem(64, ep, offset), val);
-    } else {
-        // Side exit if the write barrier is required.
-        // TODO(Jacob): Convert to guard, and maybe fix up other getblock case that uses this
-        // TODO(Jacob): Figure out what modified `VM_ENV_FLAG_WB_REQUIRED`
-        let ep = gen_get_ep(asm, level);
-        let flags = Opnd::mem(VALUE_BITS, ep, SIZEOF_VALUE_I32 * (VM_ENV_DATA_INDEX_FLAGS as i32));
-        // asm.test(flags, VM_ENV_FLAG_WB_REQUIRED.into());
-        // asm.jnz(side_exit(jit, state, SideExitReason::WriteBarrierRequired));
-        // TODO(Jacob): Remove the write barrier check that is somehow buried in these next few lines
-        // We're potentially writing a reference to an IMEMO/env object,
-        // so take care of the write barrier with a function.
-        let local_index = -local_ep_offset;
-        asm_ccall!(asm, rb_vm_env_write, ep, local_index.into(), val);
-    }
+    let offset = -(SIZEOF_VALUE_I32 * local_ep_offset);
+    asm.mov(Opnd::mem(64, ep, offset), val);
 }
 
 /// Returns 1 (as CBool) when VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM is set; returns 0 otherwise.

zjit/src/hir.rs

@@ -850,7 +850,7 @@ pub enum Insn {
     /// Get the block parameter as a Proc.
     GetBlockParam { level: u32, ep_offset: u32, state: InsnId },
     /// Set a local variable in a higher scope or the heap
-    SetLocal { level: u32, ep_offset: u32, val: InsnId, state: InsnId },
+    SetLocal { level: u32, ep_offset: u32, val: InsnId },
     GetSpecialSymbol { symbol_type: SpecialBackrefSymbol, state: InsnId },
     GetSpecialNumber { nth: u64, state: InsnId },
 
@@ -2404,7 +2404,7 @@ impl Function {
             &SetIvar { self_val, id, ic, val, state } => SetIvar { self_val: find!(self_val), id, ic, val: find!(val), state },
             &GetClassVar { id, ic, state } => GetClassVar { id, ic, state },
             &SetClassVar { id, val, ic, state } => SetClassVar { id, val: find!(val), ic, state },
-            &SetLocal { val, ep_offset, level, state } => SetLocal { val: find!(val), ep_offset, level, state },
+            &SetLocal { val, ep_offset, level } => SetLocal { val: find!(val), ep_offset, level },
             &GetSpecialSymbol { symbol_type, state } => GetSpecialSymbol { symbol_type, state },
             &GetSpecialNumber { nth, state } => GetSpecialNumber { nth, state },
             &ToArray { val, state } => ToArray { val: find!(val), state },
@@ -6723,11 +6723,12 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     }
                 }
                 YARVINSN_setlocal_WC_0 => {
+                    // TODO(Jacob): Modify this to use guards
                     let ep_offset = get_arg(pc, 0).as_u32();
                     let val = state.stack_pop()?;
                     if ep_escaped || has_blockiseq { // TODO: figure out how to drop has_blockiseq here
                         // Write the local using EP
-                        fun.push_insn(block, Insn::SetLocal { val, ep_offset, level: 0, state: exit_id });
+                        fun.push_insn(block, Insn::SetLocal { val, ep_offset, level: 0 });
                     } else if local_inval {
                         // If there has been any non-leaf call since JIT entry or the last patch point,
                         // add a patch point to make sure locals have not been escaped.
@@ -6743,20 +6744,24 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     state.stack_push(fun.push_insn(block, Insn::GetLocal { ep_offset, level: 1, use_sp: false, rest_param: false }));
                 }
                 YARVINSN_setlocal_WC_1 => {
+                    // TODO(Jacob): Modify this to use guards
                     let ep_offset = get_arg(pc, 0).as_u32();
-                    fun.push_insn(block, Insn::SetLocal { val: state.stack_pop()?, ep_offset, level: 1, state: exit_id });
+                    fun.push_insn(block, Insn::SetLocal { val: state.stack_pop()?, ep_offset, level: 1 });
                 }
                 YARVINSN_getlocal => {
                     let ep_offset = get_arg(pc, 0).as_u32();
                     let level = get_arg(pc, 1).as_u32();
                     state.stack_push(fun.push_insn(block, Insn::GetLocal { ep_offset, level, use_sp: false, rest_param: false }));
                 }
                 YARVINSN_setlocal => {
+                    // TODO(Jacob): Make sure this works
+                    // TODO(Jacob): Clean up the offset, levels, ep chasing, etc
                     let ep_offset = get_arg(pc, 0).as_u32();
                     let level = get_arg(pc, 1).as_u32();
-                    // TODO(Jacob): Add guard here
-                    // TODO(Jacob): Figure out if we should change ALL of the setlocals and what the differences are
-                    fun.push_insn(block, Insn::SetLocal { val: state.stack_pop()?, ep_offset, level, state: exit_id });
+                    let ep = fun.push_insn(block, Insn::GetEP { level });
+                    let flags = fun.push_insn(block, Insn::LoadField { recv: ep, id: ID!(_env_data_index_flags), offset: SIZEOF_VALUE_I32 * (VM_ENV_DATA_INDEX_FLAGS as i32), return_type: types::CInt64 });
+                    fun.push_insn(block, Insn::GuardNoBitsSet { val: flags, mask: Const::CUInt64(VM_ENV_FLAG_WB_REQUIRED.into()), reason: SideExitReason::WriteBarrierRequired, state: exit_id });
+                    fun.push_insn(block, Insn::SetLocal { val: state.stack_pop()?, ep_offset, level });
                 }
                 YARVINSN_getblockparamproxy => {
                     let level = get_arg(pc, 1).as_u32();