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.

Author: tekknolagi

zjit/src/hir.rs

@@ -7282,6 +7282,18 @@ 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 });
+                        local_inval = false;
                     }
                 }
                 YARVINSN_sendforward => {
@@ -7304,6 +7316,15 @@ 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 });
+                        local_inval = false;
                     }
                 }
                 YARVINSN_invokesuper => {
@@ -7325,6 +7346,9 @@ 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 });
+                        local_inval = false;
                     }
                 }
                 YARVINSN_invokesuperforward => {
@@ -7346,6 +7370,9 @@ 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 });
+                        local_inval = false;
                     }
                 }
                 YARVINSN_invokeblock => {

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
         ");
     }
 

zjit/src/hir/tests.rs

@@ -1668,6 +1668,7 @@ pub mod hir_build_tests {
           Jump bb2(v5, v6)
         bb2(v8:BasicObject, v9:BasicObject):
           v14:BasicObject = Send v9, 0x1000, :each # SendFallbackReason: Uncategorized(send)
+          PatchPoint NoEPEscape(test)
           CheckInterrupts
           Return v14
         ");
@@ -1926,6 +1927,7 @@ pub mod hir_build_tests {
           Jump bb2(v5, v6)
         bb2(v8:BasicObject, v9:BasicObject):
           v15:BasicObject = InvokeSuperForward v8, 0x1000, v9 # SendFallbackReason: Uncategorized(invokesuperforward)
+          PatchPoint NoEPEscape(test)
           CheckInterrupts
           Return v15
         ");
@@ -4366,18 +4368,19 @@ pub mod hir_build_tests {
         bb2(v8:BasicObject, v9:BasicObject):
           v14:BasicObject = Send v8, 0x1000, :tap # SendFallbackReason: Uncategorized(send)
           v15:BasicObject = GetLocal :block, l0, EP@3
-          v19:CBool = IsBlockParamModified l0
-          IfTrue v19, bb3(v8, v15)
+          PatchPoint NoEPEscape(test)
+          v21:CBool = IsBlockParamModified l0
+          IfTrue v21, bb3(v8, v15)
           Jump bb4(v8, v15)
-        bb3(v20:BasicObject, v21:BasicObject):
-          v28:BasicObject = GetLocal :block, l0, EP@3
-          Jump bb5(v20, v28, v28)
-        bb4(v23:BasicObject, v24:BasicObject):
-          v30:BasicObject = GetBlockParam :block, l0, EP@3
-          Jump bb5(v23, v30, v30)
-        bb5(v32:BasicObject, v33:BasicObject, v34:BasicObject):
-          CheckInterrupts
-          Return v34
+        bb3(v22:BasicObject, v23:BasicObject):
+          v30:BasicObject = GetLocal :block, l0, EP@3
+          Jump bb5(v22, v30, v30)
+        bb4(v25:BasicObject, v26:BasicObject):
+          v32:BasicObject = GetBlockParam :block, l0, EP@3
+          Jump bb5(v25, v32, v32)
+        bb5(v34:BasicObject, v35:BasicObject, v36:BasicObject):
+          CheckInterrupts
+          Return v36
         ");
     }
  }