ZJIT: Add NoEPEscape guard after send-with-block

tekknolagi · tekknolagi · commit 1c418035ce15 · 2026-02-11T21:01:49.000-05:00
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.
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
@@ -7282,6 +7282,17 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
 
                     if !blockiseq.is_null() {
                         reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq);
+                        // Guard against EP escape (e.g. eval) during the block. This must
+                        // be placed before any subsequent gen_spill_locals overwrites the
+                        // frame with stale local values. The snapshot uses the post-send
+                        // state without locals, pointing to the next instruction, so a
+                        // side-exit continues from the right place with locals from the frame.
+                        let len = insn_len(opcode as usize) as usize;
+                        let mut pp_state = state.clone();
+                        pp_state.pc = unsafe { pc.add(len) };
+                        pp_state.insn_idx = insn_idx as usize + len;
+                        let pp_exit_id = fun.push_insn(block, Insn::Snapshot { state: pp_state.without_locals() });
+                        fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: pp_exit_id });
                     }
                 }
                 YARVINSN_sendforward => {
@@ -7304,6 +7315,14 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
 
                     if !blockiseq.is_null() {
                         reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq);
+                        let next_insn_idx = insn_idx + insn_len(opcode as usize) as u32;
+                        let next_pc = unsafe { rb_iseq_pc_at_idx(iseq, next_insn_idx) };
+                        let mut pp_state = state.clone();
+                        pp_state.pc = next_pc;
+                        pp_state.insn_idx = next_insn_idx as usize;
+                        let pp_exit = pp_state.without_locals();
+                        let pp_exit_id = fun.push_insn(block, Insn::Snapshot { state: pp_exit });
+                        fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: pp_exit_id });
                     }
                 }
                 YARVINSN_invokesuper => {
@@ -7325,6 +7344,8 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
 
                     if !blockiseq.is_null() {
                         reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq);
+                        let pp_exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state.without_locals() });
+                        fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: pp_exit_id });
                     }
                 }
                 YARVINSN_invokesuperforward => {
@@ -7346,6 +7367,8 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
 
                     if !blockiseq.is_null() {
                         reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq);
+                        let pp_exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state.without_locals() });
+                        fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: pp_exit_id });
                     }
                 }
                 YARVINSN_invokeblock => {
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
@@ -823,10 +823,11 @@ mod hir_opt_tests {
         bb2(v8:BasicObject, v9:BasicObject):
           PatchPoint NoSingletonClass(C@0x1000)
           PatchPoint MethodRedefined(C@0x1000, fun_new_map@0x1008, cme:0x1010)
-          v21:ArraySubclass[class_exact:C] = GuardType v9, ArraySubclass[class_exact:C]
-          v22:BasicObject = SendDirect v21, 0x1038, :fun_new_map (0x1048)
+          v23:ArraySubclass[class_exact:C] = GuardType v9, ArraySubclass[class_exact:C]
+          v24:BasicObject = SendDirect v23, 0x1038, :fun_new_map (0x1048)
+          PatchPoint NoEPEscape(test)
           CheckInterrupts
-          Return v22
+          Return v24
         ");
     }
 
@@ -854,10 +855,11 @@ mod hir_opt_tests {
         bb2(v8:BasicObject, v9:BasicObject):
           PatchPoint NoSingletonClass(C@0x1000)
           PatchPoint MethodRedefined(C@0x1000, bar@0x1008, cme:0x1010)
-          v22:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C]
-          v23:BasicObject = CCallWithFrame v22, :Enumerable#bar@0x1038, block=0x1040
+          v24:HeapObject[class_exact:C] = GuardType v9, HeapObject[class_exact:C]
+          v25:BasicObject = CCallWithFrame v24, :Enumerable#bar@0x1038, block=0x1040
+          PatchPoint NoEPEscape(test)
           CheckInterrupts
-          Return v23
+          Return v25
         ");
     }
 
@@ -2919,10 +2921,11 @@ mod hir_opt_tests {
           v13:Fixnum[2] = Const Value(2)
           PatchPoint NoSingletonClass(Object@0x1000)
           PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
-          v22:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
-          v23:BasicObject = SendDirect v22, 0x1038, :foo (0x1048), v11, v13
+          v24:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
+          v25:BasicObject = SendDirect v24, 0x1038, :foo (0x1048), v11, v13
+          PatchPoint NoEPEscape(test)
           CheckInterrupts
-          Return v23
+          Return v25
         ");
     }
 
@@ -2955,6 +2958,7 @@ mod hir_opt_tests {
           PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
           v30:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v8, HeapObject[class_exact*:Object@VALUE(0x1000)]
           IncrCounter inline_iseq_optimized_send_count
+          v33:Fixnum[1] = Const Value(1)
           PatchPoint NoEPEscape(test)
           CheckInterrupts
           Return v13
@@ -3038,21 +3042,21 @@ mod hir_opt_tests {
           SetLocal :a, l0, EP@3, v13
           PatchPoint SingleRactorMode
           PatchPoint StableConstantNames(0x1000, EvalLocal)
-          v59:Class[EvalLocal@0x1008] = Const Value(VALUE(0x1008))
+          v60:Class[EvalLocal@0x1008] = Const Value(VALUE(0x1008))
           v21:NilClass = Const Value(nil)
           PatchPoint MethodRedefined(EvalLocal@0x1008, new@0x1009, cme:0x1010)
-          v62:HeapObject[class_exact:EvalLocal] = ObjectAllocClass EvalLocal:VALUE(0x1008)
+          v63:HeapObject[class_exact:EvalLocal] = ObjectAllocClass EvalLocal:VALUE(0x1008)
           PatchPoint NoSingletonClass(EvalLocal@0x1008)
           PatchPoint MethodRedefined(EvalLocal@0x1008, initialize@0x1038, cme:0x1040)
-          v70:NilClass = Const Value(nil)
+          v71:NilClass = Const Value(nil)
           IncrCounter inline_cfunc_optimized_send_count
           CheckInterrupts
           PatchPoint NoSingletonClass(EvalLocal@0x1008)
           PatchPoint MethodRedefined(EvalLocal@0x1008, f@0x1068, cme:0x1070)
-          v66:BasicObject = SendDirect v62, 0x1098, :f (0x10a8)
-          v52:BasicObject = GetLocal :a, l0, EP@3
+          v67:BasicObject = SendDirect v63, 0x1098, :f (0x10a8)
+          PatchPoint NoEPEscape(test)
           CheckInterrupts
-          Return v52
+          Return v13
         ");
     }
 
@@ -5631,6 +5635,7 @@ mod hir_opt_tests {
           Jump bb2(v4)
         bb2(v6:BasicObject):
           v11:BasicObject = Send v6, 0x1000, :bmethod # SendFallbackReason: Send: unsupported method type Bmethod
+          PatchPoint NoEPEscape(test)
           CheckInterrupts
           Return v11
         ");
@@ -6605,9 +6610,10 @@ mod hir_opt_tests {
           v11:ArrayExact = ArrayDup v10
           PatchPoint NoSingletonClass(Array@0x1008)
           PatchPoint MethodRedefined(Array@0x1008, map@0x1010, cme:0x1018)
-          v20:BasicObject = SendDirect v11, 0x1040, :map (0x1050)
+          v22:BasicObject = SendDirect v11, 0x1040, :map (0x1050)
+          PatchPoint NoEPEscape(test)
           CheckInterrupts
-          Return v20
+          Return v22
         ");
     }
 
@@ -6772,10 +6778,11 @@ mod hir_opt_tests {
         bb2(v6:BasicObject):
           PatchPoint NoSingletonClass(Object@0x1000)
           PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
-          v18:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
-          v19:BasicObject = SendDirect v18, 0x1038, :foo (0x1048)
+          v20:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
+          v21:BasicObject = SendDirect v20, 0x1038, :foo (0x1048)
+          PatchPoint NoEPEscape(test)
           CheckInterrupts
-          Return v19
+          Return v21
         ");
     }
 
@@ -9688,11 +9695,12 @@ mod hir_opt_tests {
         bb2(v6:BasicObject):
           PatchPoint NoSingletonClass(Object@0x1000)
           PatchPoint MethodRedefined(Object@0x1000, callee@0x1008, cme:0x1010)
-          v18:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
+          v20:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
           IncrCounter inline_iseq_optimized_send_count
-          v21:Fixnum[123] = Const Value(123)
+          v23:Fixnum[123] = Const Value(123)
+          PatchPoint NoEPEscape(test)
           CheckInterrupts
-          Return v21
+          Return v23
         ");
     }
 
@@ -9718,11 +9726,12 @@ mod hir_opt_tests {
         bb2(v6:BasicObject):
           PatchPoint NoSingletonClass(Object@0x1000)
           PatchPoint MethodRedefined(Object@0x1000, callee@0x1008, cme:0x1010)
-          v18:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
+          v20:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
           IncrCounter inline_iseq_optimized_send_count
-          v21:Fixnum[123] = Const Value(123)
+          v23:Fixnum[123] = Const Value(123)
+          PatchPoint NoEPEscape(test)
           CheckInterrupts
-          Return v21
+          Return v23
         ");
     }
 
@@ -9748,11 +9757,12 @@ mod hir_opt_tests {
         bb2(v6:BasicObject):
           PatchPoint NoSingletonClass(Object@0x1000)
           PatchPoint MethodRedefined(Object@0x1000, callee@0x1008, cme:0x1010)
-          v18:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
+          v20:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
           IncrCounter inline_iseq_optimized_send_count
-          v21:Fixnum[123] = Const Value(123)
+          v23:Fixnum[123] = Const Value(123)
+          PatchPoint NoEPEscape(test)
           CheckInterrupts
-          Return v21
+          Return v23
         ");
     }
 
@@ -10674,19 +10684,20 @@ mod hir_opt_tests {
          CheckInterrupts
          SetLocal :formatted, l0, EP@3, v15
          PatchPoint SingleRactorMode
-         v53:HeapBasicObject = GuardType v14, HeapBasicObject
-         v54:CShape = LoadField v53, :_shape_id@0x1000
-         v55:CShape[0x1001] = GuardBitEquals v54, CShape(0x1001)
-         StoreField v53, :@formatted@0x1002, v15
-         WriteBarrier v53, v15
-         v58:CShape[0x1003] = Const CShape(0x1003)
-         StoreField v53, :_shape_id@0x1000, v58
+         v55:HeapBasicObject = GuardType v14, HeapBasicObject
+         v56:CShape = LoadField v55, :_shape_id@0x1000
+         v57:CShape[0x1001] = GuardBitEquals v56, CShape(0x1001)
+         StoreField v55, :@formatted@0x1002, v15
+         WriteBarrier v55, v15
+         v60:CShape[0x1003] = Const CShape(0x1003)
+         StoreField v55, :_shape_id@0x1000, v60
          v46:Class[VMFrozenCore] = Const Value(VALUE(0x1008))
          PatchPoint NoSingletonClass(Class@0x1010)
          PatchPoint MethodRedefined(Class@0x1010, lambda@0x1018, cme:0x1020)
-         v63:BasicObject = CCallWithFrame v46, :RubyVM::FrozenCore.lambda@0x1048, block=0x1050
+         v65:BasicObject = CCallWithFrame v46, :RubyVM::FrozenCore.lambda@0x1048, block=0x1050
+         PatchPoint NoEPEscape(read_nil_local)
          CheckInterrupts
-         Return v63
+         Return v65
        ");
     }
 
@@ -11519,6 +11530,7 @@ mod hir_opt_tests {
           Jump bb2(v4)
         bb2(v6:BasicObject):
           v11:BasicObject = InvokeSuper v6, 0x1000 # SendFallbackReason: super: call made with a block
+          PatchPoint NoEPEscape(foo)
           CheckInterrupts
           Return v11
         ");
@@ -11685,13 +11697,13 @@ mod hir_opt_tests {
         bb2(v10:BasicObject, v11:BasicObject, v12:NilClass):
           PatchPoint NoSingletonClass(B@0x1000)
           PatchPoint MethodRedefined(B@0x1000, proc@0x1008, cme:0x1010)
-          v33:HeapObject[class_exact:B] = GuardType v10, HeapObject[class_exact:B]
-          v34:BasicObject = CCallWithFrame v33, :Kernel#proc@0x1038, block=0x1040
-          SetLocal :other_block, l0, EP@3, v34
-          v23:BasicObject = GetLocal :other_block, l0, EP@3
-          v25:BasicObject = InvokeSuper v10, 0x1048, v23 # SendFallbackReason: super: complex argument passing to `super` call
+          v34:HeapObject[class_exact:B] = GuardType v10, HeapObject[class_exact:B]
+          v35:BasicObject = CCallWithFrame v34, :Kernel#proc@0x1038, block=0x1040
+          PatchPoint NoEPEscape(foo)
+          SetLocal :other_block, l0, EP@3, v35
+          v26:BasicObject = InvokeSuper v10, 0x1048, v35 # SendFallbackReason: super: complex argument passing to `super` call
           CheckInterrupts
-          Return v25
+          Return v26
         ");
     }
 
@@ -12065,20 +12077,22 @@ mod hir_opt_tests {
           v70:Falsy = RefineType v59, Falsy
           PatchPoint NoSingletonClass(Object@0x1018)
           PatchPoint MethodRedefined(Object@0x1018, lambda@0x1020, cme:0x1028)
-          v106:HeapObject[class_exact*:Object@VALUE(0x1018)] = GuardType v57, HeapObject[class_exact*:Object@VALUE(0x1018)]
-          v107:BasicObject = CCallWithFrame v106, :Kernel#lambda@0x1050, block=0x1058
-          SetLocal :sep, l0, EP@5, v107
-          Jump bb7(v57, v58, v107, v60, v61)
-        bb7(v77:BasicObject, v78:BasicObject, v79:BasicObject, v80:BasicObject, v81:NilClass):
+          v110:HeapObject[class_exact*:Object@VALUE(0x1018)] = GuardType v57, HeapObject[class_exact*:Object@VALUE(0x1018)]
+          v111:BasicObject = CCallWithFrame v110, :Kernel#lambda@0x1050, block=0x1058
+          PatchPoint NoEPEscape(test_ep_escape)
+          SetLocal :sep, l0, EP@5, v111
+          Jump bb7(v57, v58, v111, v60, v61)
+        bb7(v79:BasicObject, v80:BasicObject, v81:BasicObject, v82:BasicObject, v83:NilClass):
           PatchPoint SingleRactorMode
           PatchPoint StableConstantNames(0x1060, CONST)
-          v102:HashExact[VALUE(0x1068)] = Const Value(VALUE(0x1068))
-          SetLocal :kwsplat, l0, EP@3, v102
-          v91:BasicObject = GetLocal :list, l0, EP@6
-          v93:BasicObject = GetLocal :iter_method, l0, EP@4
-          v95:BasicObject = Send v91, 0x1070, :__send__, v93 # SendFallbackReason: Send: unsupported method type Optimized
-          CheckInterrupts
-          Return v95
+          v106:HashExact[VALUE(0x1068)] = Const Value(VALUE(0x1068))
+          SetLocal :kwsplat, l0, EP@3, v106
+          v93:BasicObject = GetLocal :list, l0, EP@6
+          v95:BasicObject = GetLocal :iter_method, l0, EP@4
+          v97:BasicObject = Send v93, 0x1070, :__send__, v95 # SendFallbackReason: Uncategorized(send)
+          PatchPoint NoEPEscape(test_ep_escape)
+          CheckInterrupts
+          Return v97
         ");
     }
 
diff --git a/zjit/src/hir/tests.rs b/zjit/src/hir/tests.rs

Original file line number	Diff line number	Diff line change
`@@ -7282,6 +7282,17 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {`
`7282`	`7282`
`7283`	`7283`	`if !blockiseq.is_null() {`
`7284`	`7284`	`reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq);`
	`7285`	`+ // Guard against EP escape (e.g. eval) during the block. This must`
	`7286`	`+ // be placed before any subsequent gen_spill_locals overwrites the`
	`7287`	`+ // frame with stale local values. The snapshot uses the post-send`
	`7288`	`+ // state without locals, pointing to the next instruction, so a`
	`7289`	`+ // side-exit continues from the right place with locals from the frame.`
	`7290`	`+ let len = insn_len(opcode as usize) as usize;`
	`7291`	`+ let mut pp_state = state.clone();`
	`7292`	`+ pp_state.pc = unsafe { pc.add(len) };`
	`7293`	`+ pp_state.insn_idx = insn_idx as usize + len;`
	`7294`	`+ let pp_exit_id = fun.push_insn(block, Insn::Snapshot { state: pp_state.without_locals() });`
	`7295`	`+ fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: pp_exit_id });`
`7285`	`7296`	`}`
`7286`	`7297`	`}`
`7287`	`7298`	`YARVINSN_sendforward => {`
`@@ -7304,6 +7315,14 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {`
`7304`	`7315`
`7305`	`7316`	`if !blockiseq.is_null() {`
`7306`	`7317`	`reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq);`
	`7318`	`+ let next_insn_idx = insn_idx + insn_len(opcode as usize) as u32;`
	`7319`	`+ let next_pc = unsafe { rb_iseq_pc_at_idx(iseq, next_insn_idx) };`
	`7320`	`+ let mut pp_state = state.clone();`
	`7321`	`+ pp_state.pc = next_pc;`
	`7322`	`+ pp_state.insn_idx = next_insn_idx as usize;`
	`7323`	`+ let pp_exit = pp_state.without_locals();`
	`7324`	`+ let pp_exit_id = fun.push_insn(block, Insn::Snapshot { state: pp_exit });`
	`7325`	`+ fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: pp_exit_id });`
`7307`	`7326`	`}`
`7308`	`7327`	`}`
`7309`	`7328`	`YARVINSN_invokesuper => {`
`@@ -7325,6 +7344,8 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {`
`7325`	`7344`
`7326`	`7345`	`if !blockiseq.is_null() {`
`7327`	`7346`	`reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq);`
	`7347`	`+ let pp_exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state.without_locals() });`
	`7348`	`+ fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: pp_exit_id });`
`7328`	`7349`	`}`
`7329`	`7350`	`}`
`7330`	`7351`	`YARVINSN_invokesuperforward => {`
`@@ -7346,6 +7367,8 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {`
`7346`	`7367`
`7347`	`7368`	`if !blockiseq.is_null() {`
`7348`	`7369`	`reload_modified_locals(&mut fun, block, &mut state, iseq, blockiseq);`
	`7370`	`+ let pp_exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state.without_locals() });`
	`7371`	`+ fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: pp_exit_id });`
`7349`	`7372`	`}`
`7350`	`7373`	`}`
`7351`	`7374`	`YARVINSN_invokeblock => {`