ZJIT: For JIT-to-JIT send, avoid loading uninitialized local through EP

JIT-to-JIT sends don't blit locals to nil in the callee's
EP memory region because HIR is aware of this initial state and
memory ops are only done when necessary. Previously, we
read from this initialized memory by emitting `GetLocal` in e.g. BBs
that are immediate successor to an entrypoint.

The entry points sets up the frame state properly and we also reload
locals if necessary after an operation that potentially makes the
environment escape. So, listen to the frame state when it's supposed to
be up-to-date (`!local_inval`).

zjit/src/hir.rs

@@ -5230,19 +5230,24 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                 }
                 YARVINSN_getlocal_WC_0 => {
                     let ep_offset = get_arg(pc, 0).as_u32();
-                    if ep_escaped || has_blockiseq { // TODO: figure out how to drop has_blockiseq here
+                    if !local_inval {
+                        // The FrameState is the source of truth for locals until invalidated.
+                        // In case of JIT-to-JIT send locals might never end up in EP memory.
+                        let val = state.getlocal(ep_offset);
+                        state.stack_push(val);
+                    } else if ep_escaped || has_blockiseq { // TODO: figure out how to drop has_blockiseq here
                         // Read the local using EP
                         let val = fun.push_insn(block, Insn::GetLocal { ep_offset, level: 0, use_sp: false, rest_param: false });
                         state.setlocal(ep_offset, val); // remember the result to spill on side-exits
                         state.stack_push(val);
                     } 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.
-                            let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state.without_locals() }); // skip spilling locals
-                            fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: exit_id });
-                            local_inval = false;
-                        }
+                        assert!(local_inval); // if check above
+                        // There has been some non-leaf call since JIT entry or the last patch point,
+                        // so add a patch point to make sure locals have not been escaped.
+                        let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state.without_locals() }); // skip spilling locals
+                        fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoEPEscape(iseq), state: exit_id });
+                        local_inval = false;
+
                         // Read the local from FrameState
                         let val = state.getlocal(ep_offset);
                         state.stack_push(val);

zjit/src/hir/opt_tests.rs

@@ -780,14 +780,13 @@ mod hir_opt_tests {
           EntryPoint JIT(0)
           Jump bb2(v5, v6)
         bb2(v8:BasicObject, v9:BasicObject):
-          v13:BasicObject = GetLocal l0, EP@3
           PatchPoint MethodRedefined(C@0x1000, fun_new_map@0x1008, cme:0x1010)
           PatchPoint NoSingletonClass(C@0x1000)
-          v24:ArraySubclass[class_exact:C] = GuardType v13, ArraySubclass[class_exact:C]
-          v25:BasicObject = CCallWithFrame C#fun_new_map@0x1038, v24, block=0x1040
-          v16:BasicObject = GetLocal l0, EP@3
+          v23:ArraySubclass[class_exact:C] = GuardType v9, ArraySubclass[class_exact:C]
+          v24:BasicObject = CCallWithFrame C#fun_new_map@0x1038, v23, block=0x1040
+          v15:BasicObject = GetLocal l0, EP@3
           CheckInterrupts
-          Return v25
+          Return v24
         ");
     }
 
@@ -8425,4 +8424,57 @@ mod hir_opt_tests {
           Return v32
         ");
     }
+
+    #[test]
+    fn no_load_from_ep_right_after_entrypoint() {
+      let formatted = eval("
+          def read_nil_local(a, _b, _c)
+            formatted ||= a
+            @formatted = formatted
+            -> { formatted } # the environment escapes
+          end
+
+          def call
+            puts [], [], [], []     # fill VM stack with junk
+            read_nil_local(true, 1, 1) # expected direct send
+          end
+
+          call # profile
+          call # compile
+          @formatted
+       ");
+       assert_eq!(Qtrue, formatted, "{}", formatted.obj_info());
+       assert_snapshot!(hir_string("read_nil_local"), @r"
+       fn read_nil_local@<compiled>:3:
+       bb0():
+         EntryPoint interpreter
+         v1:BasicObject = LoadSelf
+         v2:BasicObject = GetLocal l0, SP@7
+         v3:BasicObject = GetLocal l0, SP@6
+         v4:BasicObject = GetLocal l0, SP@5
+         v5:NilClass = Const Value(nil)
+         Jump bb2(v1, v2, v3, v4, v5)
+       bb1(v8:BasicObject, v9:BasicObject, v10:BasicObject, v11:BasicObject):
+         EntryPoint JIT(0)
+         v12:NilClass = Const Value(nil)
+         Jump bb2(v8, v9, v10, v11, v12)
+       bb2(v14:BasicObject, v15:BasicObject, v16:BasicObject, v17:BasicObject, v18:NilClass):
+         CheckInterrupts
+         v27:BasicObject = GetLocal l0, EP@6
+         SetLocal l0, EP@3, v27
+         v39:BasicObject = GetLocal l0, EP@3
+         PatchPoint SingleRactorMode
+         SetIvar v14, :@formatted, v39
+         v45:Class[VMFrozenCore] = Const Value(VALUE(0x1000))
+         PatchPoint MethodRedefined(Class@0x1008, lambda@0x1010, cme:0x1018)
+         PatchPoint NoSingletonClass(Class@0x1008)
+         v59:BasicObject = CCallWithFrame RubyVM::FrozenCore.lambda@0x1040, v45, block=0x1048
+         v48:BasicObject = GetLocal l0, EP@6
+         v49:BasicObject = GetLocal l0, EP@5
+         v50:BasicObject = GetLocal l0, EP@4
+         v51:BasicObject = GetLocal l0, EP@3
+         CheckInterrupts
+         Return v59
+       ");
+    }
 }

zjit/src/hir/tests.rs

@@ -1500,11 +1500,10 @@ pub mod hir_build_tests {
           EntryPoint JIT(0)
           Jump bb2(v5, v6)
         bb2(v8:BasicObject, v9:BasicObject):
-          v13:BasicObject = GetLocal l0, EP@3
-          v15:BasicObject = Send v13, 0x1000, :each
-          v16:BasicObject = GetLocal l0, EP@3
+          v14:BasicObject = Send v9, 0x1000, :each
+          v15:BasicObject = GetLocal l0, EP@3
           CheckInterrupts
-          Return v15
+          Return v14
         ");
     }