ZJIT: Locals in stackmap

Author: XrXr

vm.c

@@ -1128,6 +1128,8 @@ vm_make_env_each(const rb_execution_context_t * const ec, rb_control_frame_t *co
     if (VM_FRAME_RUBYFRAME_P(cfp)) {
         rb_yjit_invalidate_ep_is_bp(iseq);
         rb_zjit_invalidate_no_ep_escape(iseq);
+        // Also use the stack map to put values into the format expected below
+        rb_zjit_spill_frame(cfp);
     }
 
     /*
@@ -2883,6 +2885,35 @@ vm_exec_loop(rb_execution_context_t *ec, enum ruby_tag_type state,
 }
 
 #if USE_ZJIT
+// Write a ZJIT frame's off-stack Ruby slots (operands and locals) into VM-stack
+// memory using its stack map, decoding downward from cfp->sp. Does not touch
+// any frame metadata (pc/_iseq/block_code/jit_return).
+static void
+zjit_write_frame_stack(rb_control_frame_t *cfp, const zjit_jit_frame_t *jit_frame)
+{
+    int32_t stack_size = (int32_t)jit_frame->stack_size;
+    if (stack_size > 0) {
+        VALUE *stack = cfp->sp;
+        for (int32_t i = 0; i < stack_size; i++) {
+            VALUE entry = jit_frame->stack[i];
+            if (ZJIT_STACK_MAP_VREG_P(entry)) {
+                // Decode a native stack slot offset generated by ZJIT's backend.
+                // It's an offset from NATIVE_BASE_PTR, which is copied into
+                // cfp->jit_return, to the encoded stack slot.
+                stack--;
+                *stack = ((VALUE *)cfp->jit_return)[-(ssize_t)ZJIT_STACK_MAP_VREG_INDEX(entry)];
+            }
+            else if (ZJIT_STACK_MAP_SKIP_P(entry)) {
+                stack -= ZJIT_STACK_MAP_SKIP_SIZE(entry);
+            }
+            else {
+                stack--;
+                *stack = entry;
+            }
+        }
+    }
+}
+
 // Materialize JITFrame-enabled CFP into interpreter-compatible CFP
 static void
 zjit_materialize_frames(const rb_execution_context_t *ec, rb_control_frame_t *cfp, bool materialize_target)
@@ -2916,27 +2947,7 @@ zjit_materialize_frames(const rb_execution_context_t *ec, rb_control_frame_t *cf
             // the exiting frame is already written by compile_exit_save_state()
             // and skipped here after materialize_exit_trampoline clears its
             // jit_return, so this restores older ZJIT frames from stack maps.
-            int32_t stack_size = (int32_t)jit_frame->stack_size;
-            if (stack_size > 0) {
-                VALUE *stack = cfp->sp;
-                for (int32_t i = 0; i < stack_size; i++) {
-                    VALUE entry = jit_frame->stack[i];
-                    if (ZJIT_STACK_MAP_VREG_P(entry)) {
-                        // Decode a native stack slot offset generated by ZJIT's backend.
-                        // It's an offset from NATIVE_BASE_PTR, which is copied into
-                        // cfp->jit_return, to the encoded stack slot.
-                        stack--;
-                        *stack = ((VALUE *)cfp->jit_return)[-(ssize_t)ZJIT_STACK_MAP_VREG_INDEX(entry)];
-                    }
-                    else if (ZJIT_STACK_MAP_SKIP_P(entry)) {
-                        stack -= ZJIT_STACK_MAP_SKIP_SIZE(entry);
-                    }
-                    else {
-                        stack--;
-                        *stack = entry;
-                    }
-                }
-            }
+            zjit_write_frame_stack(cfp, jit_frame);
             cfp->jit_return = 0;
         }
         if (end_cfp == cfp) break;
@@ -2957,6 +2968,19 @@ rb_zjit_materialize_frames_for_longjmp(const rb_execution_context_t *ec, rb_cont
     // stack and survives longjmp. Materialize only the frames unwound above it.
     zjit_materialize_frames(ec, cfp, !ec->tag->zjit_frame_active);
 }
+
+// Spill a single ZJIT frame's off-stack Ruby slots (operands and locals) into
+// VM-stack memory, leaving the frame lazy (jit_return and pc/_iseq intact).
+// vm_make_env_each() calls this so an escaping env can MEMCPY live locals out
+// of VM memory.
+void
+rb_zjit_spill_frame(rb_control_frame_t *cfp)
+{
+    if (!rb_zjit_enabled_p) return;
+    if (CFP_ZJIT_FRAME_P(cfp)) {
+        zjit_write_frame_stack(cfp, CFP_ZJIT_FRAME(cfp));
+    }
+}
 #endif
 
 static inline VALUE

zjit.h

@@ -98,6 +98,7 @@ bool rb_zjit_class_allocate_instance_fastpath(VALUE klass, size_t *size_out, sha
 bool rb_zjit_str_resurrect_fastpath(VALUE str, bool chilled, size_t *size_out, VALUE *flags_out, long *len_out, size_t *byte_size_out);
 bool rb_zjit_array_dup_can_fastpath(VALUE ary, size_t *alloc_size_out, VALUE *flags_out, long *len_out);
 void rb_zjit_range_new_fastpath(bool exclude_end, size_t *alloc_size_out, VALUE *flags_out);
+void rb_zjit_spill_frame(rb_control_frame_t *cfp);
 
 // Special value for cfp->jit_return that means "this is a C method frame, use
 // rb_zjit_c_frame as the JITFrame". We don't control the native stack layout
@@ -140,6 +141,7 @@ static inline void rb_zjit_invalidate_root_box(void) {}
 static inline void rb_zjit_jit_frame_update_references(zjit_jit_frame_t *jit_frame) {}
 static inline void rb_zjit_materialize_frames(const rb_execution_context_t *ec, rb_control_frame_t *cfp) {}
 static inline void rb_zjit_materialize_frames_for_longjmp(const rb_execution_context_t *ec, rb_control_frame_t *cfp) {}
+static inline void rb_zjit_spill_frame(rb_control_frame_t *cfp) {}
 static inline const zjit_jit_frame_t *CFP_ZJIT_FRAME(const rb_control_frame_t *cfp) { return NULL; }
 #endif // #if USE_ZJIT
 

zjit/src/codegen.rs

@@ -962,6 +962,9 @@ fn gen_invokebuiltin(jit: &JITState, asm: &mut Assembler, function: &Function, s
     } else {
         // Anything can happen inside builtin functions
         gen_prepare_non_leaf_call(jit, asm, function, state);
+        // cexpr!/cstmt! builtins read this frame's locals (its params) directly
+        // from EP, so make them authoritative in memory.
+        gen_spill_locals(jit, asm, state);
     }
 
     let mut cargs = vec![EC];
@@ -1049,7 +1052,12 @@ fn gen_ccall_with_frame(
     gen_write_jit_frame(asm, state, 0);
     gen_save_sp(asm, caller_stack_size);
     gen_spill_stack(jit, asm, function, state);
-    gen_spill_locals(jit, asm, state);
+    // A passed block can read or write this frame's locals through its EP, so
+    // make them authoritative in memory. Without a block, the stack map is
+    // enough to reconstruct locals lazily on exception/binding.
+    if block.is_some() {
+        gen_spill_locals(jit, asm, state);
+    }
 
     let block_handler_specval = if let Some(BlockHandler::BlockIseq(block_iseq)) = block {
         // Change cfp->block_code in the current frame. See vm_caller_setup_arg_block().
@@ -1144,7 +1152,10 @@ fn gen_ccall_variadic(
     gen_write_jit_frame(asm, state, 0);
     gen_save_sp(asm, caller_stack_size);
     gen_spill_stack(jit, asm, function, state);
-    gen_spill_locals(jit, asm, state);
+    // A passed block can read or write this frame's locals through its EP.
+    if block.is_some() {
+        gen_spill_locals(jit, asm, state);
+    }
 
     let block_handler_specval = if let Some(BlockHandler::BlockIseq(blockiseq)) = block {
         gen_block_handler_specval(asm, blockiseq)
@@ -1505,6 +1516,10 @@ fn gen_send(
     gen_trace_send_fallback(asm, &reason);
 
     gen_prepare_fallback_call(jit, asm, function, state);
+    // A literal block passed here can read or write this frame's locals through its EP.
+    if !blockiseq.is_null() {
+        gen_spill_locals(jit, asm, state);
+    }
     asm_comment!(asm, "call #{} with dynamic dispatch", ruby_call_method_name(cd));
     unsafe extern "C" {
         fn rb_vm_send(ec: EcPtr, cfp: CfpPtr, cd: VALUE, blockiseq: IseqPtr) -> VALUE;
@@ -1530,6 +1545,10 @@ fn gen_send_forward(
     gen_trace_send_fallback(asm, &reason);
 
     gen_prepare_fallback_call(jit, asm, function, state);
+    // A literal block passed here can read or write this frame's locals through its EP.
+    if !blockiseq.is_null() {
+        gen_spill_locals(jit, asm, state);
+    }
 
     asm_comment!(asm, "call #{} with dynamic dispatch", ruby_call_method_name(cd));
     unsafe extern "C" {
@@ -1590,7 +1609,10 @@ fn gen_push_inline_frame(
     gen_write_jit_frame(asm, state, 0);
     gen_save_sp(asm, stack_size);
 
-    gen_spill_locals(jit, asm, state);
+    // A passed block can read or write this frame's locals through its EP.
+    if blockiseq.is_some() {
+        gen_spill_locals(jit, asm, state);
+    }
 
     // This mirrors vm_caller_setup_arg_block() for the `blockiseq != NULL` case.
     // The HIR specialization guards ensure we will only reach here for literal blocks,
@@ -1729,7 +1751,11 @@ fn gen_send_iseq_direct(
     let jit_frame = gen_write_jit_frame(asm, state, stack_map.len());
     gen_save_sp(asm, stack_size);
 
-    gen_spill_locals(jit, asm, state);
+    // A passed block can read or write this frame's locals through its EP.
+    // Without a block, the stack map reconstructs locals lazily on demand.
+    if block.is_some() {
+        gen_spill_locals(jit, asm, state);
+    }
     asm.stack_map(stack_map, jit_frame, state.depth);
 
     // This mirrors vm_caller_setup_arg_block() in for the `blockiseq != NULL` case.
@@ -2018,6 +2044,10 @@ fn gen_invokesuper(
     gen_trace_send_fallback(asm, &reason);
 
     gen_prepare_fallback_call(jit, asm, function, state);
+    // A literal block passed here can read or write this frame's locals through its EP.
+    if !blockiseq.is_null() {
+        gen_spill_locals(jit, asm, state);
+    }
     asm_comment!(asm, "call super with dynamic dispatch");
     unsafe extern "C" {
         fn rb_vm_invokesuper(ec: EcPtr, cfp: CfpPtr, cd: VALUE, blockiseq: IseqPtr) -> VALUE;
@@ -2043,6 +2073,10 @@ fn gen_invokesuperforward(
     gen_trace_send_fallback(asm, &reason);
 
     gen_prepare_fallback_call(jit, asm, function, state);
+    // A literal block passed here can read or write this frame's locals through its EP.
+    if !blockiseq.is_null() {
+        gen_spill_locals(jit, asm, state);
+    }
     asm_comment!(asm, "call super with dynamic dispatch (forwarding)");
     unsafe extern "C" {
         fn rb_vm_invokesuperforward(ec: EcPtr, cfp: CfpPtr, cd: VALUE, blockiseq: IseqPtr) -> VALUE;
@@ -3373,7 +3407,6 @@ fn gen_spill_stack(jit: &JITState, asm: &mut Assembler, function: &Function, sta
 fn gen_prepare_fallback_call(jit: &JITState, asm: &mut Assembler, function: &Function, state: &FrameState) {
     gen_write_jit_frame(asm, state, 0);
     gen_save_sp(asm, state.stack_size());
-    gen_spill_locals(jit, asm, state);
     gen_spill_stack(jit, asm, function, state);
 }
 
@@ -3384,31 +3417,46 @@ fn build_stack_map(jit: &JITState, function: &Function, state: &FrameState) -> V
     let mut stack = Vec::new();
     let mut current_state = state.clone();
     loop {
-        stack.extend(current_state.stack().rev().copied().map(|insn_id| {
+        let to_entry = |insn_id| {
             let opnd = jit.get_opnd(insn_id);
             assert!(
                 matches!(opnd, Opnd::Value(_) | Opnd::VReg { .. }),
                 "FrameState should only reference Opnd::Value or Opnd::VReg, but got: {opnd:?}",
             );
             StackMapEntry::Opnd(opnd)
-        }));
+        };
+
+        // Operand stack, top-down.
+        stack.extend(current_state.stack().rev().copied().map(to_entry));
+        // Frame environment data (me/cref, specval, flags) already lives in memory.
+        stack.push(StackMapEntry::Skip(VM_ENV_DATA_SIZE.to_usize()));
+        // Locals, top-down (local[L-1] .. local[0]). They land at fixed
+        // EP-relative slots because we write down from cfp->sp.
+        //
+        // Locals backed by memory are accessed in mainline code,
+        // so we skip over those slots to not interfere and clobber.
+        // (The stack map is not the 1st write to the slot.)
+        let spilled_locals = current_state.spilled_locals();
+        for (idx, &insn_id) in current_state.locals().enumerate().rev() {
+            if spilled_locals.contains(&(idx as u32)) {
+                stack.push(StackMapEntry::Skip(1));
+            } else {
+                stack.push(to_entry(insn_id));
+            }
+        }
 
         let Some(caller) = current_state.caller() else {
             break;
         };
-        stack.push(StackMapEntry::Skip(inline_frame_stack_gap(current_state.iseq)));
+        // Skip the callee's receiver slot below its local table. We currently
+        // never map out the stack for `invokeblock`, which doesn't put a
+        // receiver on cfp->sp stack.
+        stack.push(StackMapEntry::Skip(1));
         current_state = function.frame_state(caller);
     }
     stack
 }
 
-fn inline_frame_stack_gap(iseq: IseqPtr) -> usize {
-    // The extra slot is for the callee's receiver below its local table.
-    // We currently never map out the stack for `invokeblock`, which doesn't
-    // put a receiver on cfp->sp stack.
-    1 + unsafe { get_iseq_body_local_table_size(iseq) }.to_usize() + VM_ENV_DATA_SIZE.to_usize()
-}
-
 /// Prepare for calling a C function that may call an arbitrary method.
 /// Use gen_prepare_leaf_call_with_gc() if the method is leaf but allocates objects.
 fn gen_prepare_non_leaf_call(jit: &JITState, asm: &mut Assembler, function: &Function, state: &FrameState) {
@@ -3421,9 +3469,6 @@ fn gen_prepare_non_leaf_call(jit: &JITState, asm: &mut Assembler, function: &Fun
     // Remember the stack map in case it raises an exception
     // and the interpreter uses the stack for handling the exception
     asm.stack_map(stack_map, jit_frame, state.depth);
-
-    // Spill locals in case the method looks at caller Bindings
-    gen_spill_locals(jit, asm, state);
 }
 
 /// Frame metadata written by gen_push_frame()

zjit/src/codegen_tests.rs

@@ -6742,9 +6742,9 @@ fn test_polymorphic_getivar_complex_shape() {
 
 /// When a method with keyword defaults contains a block that creates a lambda,
 /// the lambda causes EP escape, which globally patches NoEPEscape PatchPoints.
-/// On subsequent calls the PatchPoint side exit (which uses without_locals())
-/// must not leave stale keyword default values in the frame. We solve this by
-/// invalidating the ISEQ version on EP escape so the interpreter takes over.
+/// On subsequent calls the PatchPoint side exit must not leave stale keyword
+/// default values in the frame. We solve this by invalidating the ISEQ version
+/// on EP escape so the interpreter takes over.
 #[test]
 fn test_ep_escape_preserves_keyword_default() {
     set_call_threshold(1);