Skip to content

Commit 5b42da5

Browse files
tekknolagitekknolagi
committed
ZJIT: Add NoEPEscape guard after send-with-block
After reload_modified_locals, locals not written by the block retain stale FrameState values. If eval ran inside the block, these stale values would be spilled to the frame by gen_spill_locals on the next non-leaf call, overwriting the correct values set by eval. Add a NoEPEscape PatchPoint immediately after reload_modified_locals. If the EP escaped during the block (e.g. via eval), this side-exits before any stale locals can be spilled. The snapshot uses the post-send state without locals so the interpreter reads them from the frame, which has the correct values at that point.
1 parent f5ec572 commit 5b42da5

File tree

3 files changed

+54
-28
lines changed

3 files changed

+54
-28
lines changed

‎zjit/src/hir.rs‎

Lines changed: 27 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6307,20 +6307,41 @@ fn locals_written_in_block(parent_iseq: IseqPtr, blockiseq: IseqPtr, target_dept
63076307
}
63086308
}
63096309

6310-
/// Reload locals that may have been modified by the blockiseq.
6311-
fn reload_modified_locals(fun: &mut Function, block: BlockId, state: &mut FrameState, iseq: IseqPtr, blockiseq: IseqPtr) {
6310+
/// Reload locals that may have been modified by the blockiseq, and if any were
6311+
/// reloaded, guard against EP escape (e.g. eval writing to locals we can't
6312+
/// detect by scanning bytecode).
6313+
///
6314+
/// `pc` must point to the send instruction and `opcode` must be its decoded opcode.
6315+
/// We build a side-exit snapshot pointing to the *next* instruction so the interpreter
6316+
/// doesn't re-execute the send. We compute next_pc with pointer arithmetic
6317+
/// (`pc + insn_len`) rather than `rb_iseq_pc_at_idx` to avoid a bounds-check assertion.
6318+
fn reload_modified_locals(fun: &mut Function, block: BlockId, state: &mut FrameState, iseq: IseqPtr, blockiseq: IseqPtr, pc: *const VALUE, opcode: u32, insn_idx: u32) {
63126319
// TODO: Avoid reloading locals that are not referenced by the blockiseq
63136320
// or not used after this. Max thinks we could eventually DCE them.
63146321
let mut locals = BitSet::with_capacity(state.locals.len());
63156322
locals_written_in_block(iseq, blockiseq, 1, &mut locals);
6323+
let mut reloaded = false;
63166324
for local_idx in 0..state.locals.len() {
63176325
if locals.get(local_idx) {
63186326
let ep_offset = local_idx_to_ep_offset(iseq, local_idx) as u32;
63196327
// TODO: We could use `use_sp: true` with PatchPoint
63206328
let val = fun.push_insn(block, Insn::GetLocal { ep_offset, level: 0, use_sp: false, rest_param: false });
63216329
state.setlocal(ep_offset, val);
6330+
reloaded = true;
63226331
}
63236332
}
6333+
6334+
// Only add a NoEPEscape guard when we actually reloaded locals. If we didn't
6335+
// reload any, the subsequent getlocal PatchPoints (from local_inval) handle it.
6336+
// The snapshot omits locals so the interpreter reads them from the frame.
6337+
if reloaded {
6338+
let len = insn_len(opcode as usize) as usize;
6339+
let mut pp_state = state.clone();
6340+
pp_state.pc = unsafe { pc.add(len) };
6341+
pp_state.insn_idx = insn_idx as usize + len;
6342+
let pp_exit_id = fun.push_insn(block, Insn::Snapshot { state: pp_state.without_locals() });
6343+
fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: pp_exit_id });
6344+
}
63246345
}
63256346

63266347
/// Compile ISEQ into High-level IR
@@ -7281,7 +7302,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
72817302
state.stack_push(send);
72827303

72837304
if !blockiseq.is_null() {
7284-
reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq);
7305+
reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq, pc, opcode, insn_idx);
72857306
}
72867307
}
72877308
YARVINSN_sendforward => {
@@ -7303,7 +7324,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
73037324
state.stack_push(send_forward);
73047325

73057326
if !blockiseq.is_null() {
7306-
reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq);
7327+
reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq, pc, opcode, insn_idx);
73077328
}
73087329
}
73097330
YARVINSN_invokesuper => {
@@ -7324,7 +7345,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
73247345
state.stack_push(result);
73257346

73267347
if !blockiseq.is_null() {
7327-
reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq);
7348+
reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq, pc, opcode, insn_idx);
73287349
}
73297350
}
73307351
YARVINSN_invokesuperforward => {
@@ -7345,7 +7366,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
73457366
state.stack_push(result);
73467367

73477368
if !blockiseq.is_null() {
7348-
reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq);
7369+
reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq, pc, opcode, insn_idx);
73497370
}
73507371
}
73517372
YARVINSN_invokeblock => {

‎zjit/src/hir/opt_tests.rs‎

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2997,16 +2997,17 @@ mod hir_opt_tests {
29972997
v20:Fixnum[2] = Const Value(2)
29982998
PatchPoint NoSingletonClass(Object@0x1000)
29992999
PatchPoint MethodRedefined(Object@0x1000, tap@0x1008, cme:0x1010)
3000-
v42:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v10, HeapObject[class_exact*:Object@VALUE(0x1000)]
3001-
v43:BasicObject = SendDirect v42, 0x1038, :tap (0x1048)
3000+
v44:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v10, HeapObject[class_exact*:Object@VALUE(0x1000)]
3001+
v45:BasicObject = SendDirect v44, 0x1038, :tap (0x1048)
30023002
v26:BasicObject = GetLocal :b, l0, EP@3
30033003
PatchPoint NoEPEscape(test)
3004+
PatchPoint NoEPEscape(test)
30043005
PatchPoint MethodRedefined(Integer@0x1050, +@0x1058, cme:0x1060)
3005-
v46:Fixnum = GuardType v26, Fixnum
3006-
v47:Fixnum = FixnumAdd v16, v46
3006+
v48:Fixnum = GuardType v26, Fixnum
3007+
v49:Fixnum = FixnumAdd v16, v48
30073008
IncrCounter inline_cfunc_optimized_send_count
30083009
CheckInterrupts
3009-
Return v47
3010+
Return v49
30103011
");
30113012
}
30123013

‎zjit/src/hir/tests.rs‎

Lines changed: 21 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -4216,9 +4216,10 @@ pub mod hir_build_tests {
42164216
v25:BasicObject = Send v10, 0x1000, :tap # SendFallbackReason: Uncategorized(send)
42174217
v26:BasicObject = GetLocal :b, l0, EP@3
42184218
PatchPoint NoEPEscape(test)
4219-
v35:BasicObject = SendWithoutBlock v16, :+, v26 # SendFallbackReason: Uncategorized(opt_plus)
4219+
PatchPoint NoEPEscape(test)
4220+
v37:BasicObject = SendWithoutBlock v16, :+, v26 # SendFallbackReason: Uncategorized(opt_plus)
42204221
CheckInterrupts
4221-
Return v35
4222+
Return v37
42224223
");
42234224
}
42244225

@@ -4252,9 +4253,10 @@ pub mod hir_build_tests {
42524253
v25:BasicObject = Send v10, 0x1000, :tap # SendFallbackReason: Uncategorized(send)
42534254
v26:BasicObject = GetLocal :b, l0, EP@3
42544255
PatchPoint NoEPEscape(test)
4255-
v35:BasicObject = SendWithoutBlock v16, :+, v26 # SendFallbackReason: Uncategorized(opt_plus)
4256+
PatchPoint NoEPEscape(test)
4257+
v37:BasicObject = SendWithoutBlock v16, :+, v26 # SendFallbackReason: Uncategorized(opt_plus)
42564258
CheckInterrupts
4257-
Return v35
4259+
Return v37
42584260
");
42594261
}
42604262

@@ -4294,9 +4296,10 @@ pub mod hir_build_tests {
42944296
v25:BasicObject = Send v10, 0x1000, :tap # SendFallbackReason: Uncategorized(send)
42954297
v26:BasicObject = GetLocal :b, l0, EP@3
42964298
PatchPoint NoEPEscape(test)
4297-
v35:BasicObject = SendWithoutBlock v16, :+, v26 # SendFallbackReason: Uncategorized(opt_plus)
4299+
PatchPoint NoEPEscape(test)
4300+
v37:BasicObject = SendWithoutBlock v16, :+, v26 # SendFallbackReason: Uncategorized(opt_plus)
42984301
CheckInterrupts
4299-
Return v35
4302+
Return v37
43004303
");
43014304
}
43024305

@@ -4366,18 +4369,19 @@ pub mod hir_build_tests {
43664369
bb2(v8:BasicObject, v9:BasicObject):
43674370
v14:BasicObject = Send v8, 0x1000, :tap # SendFallbackReason: Uncategorized(send)
43684371
v15:BasicObject = GetLocal :block, l0, EP@3
4369-
v19:CBool = IsBlockParamModified l0
4370-
IfTrue v19, bb3(v8, v15)
4372+
PatchPoint NoEPEscape(test)
4373+
v21:CBool = IsBlockParamModified l0
4374+
IfTrue v21, bb3(v8, v15)
43714375
Jump bb4(v8, v15)
4372-
bb3(v20:BasicObject, v21:BasicObject):
4373-
v28:BasicObject = GetLocal :block, l0, EP@3
4374-
Jump bb5(v20, v28, v28)
4375-
bb4(v23:BasicObject, v24:BasicObject):
4376-
v30:BasicObject = GetBlockParam :block, l0, EP@3
4377-
Jump bb5(v23, v30, v30)
4378-
bb5(v32:BasicObject, v33:BasicObject, v34:BasicObject):
4379-
CheckInterrupts
4380-
Return v34
4376+
bb3(v22:BasicObject, v23:BasicObject):
4377+
v30:BasicObject = GetLocal :block, l0, EP@3
4378+
Jump bb5(v22, v30, v30)
4379+
bb4(v25:BasicObject, v26:BasicObject):
4380+
v32:BasicObject = GetBlockParam :block, l0, EP@3
4381+
Jump bb5(v25, v32, v32)
4382+
bb5(v34:BasicObject, v35:BasicObject, v36:BasicObject):
4383+
CheckInterrupts
4384+
Return v36
43814385
");
43824386
}
43834387
}

0 commit comments

Comments
 (0)