ZJIT: Fix stack map clobbering VM stack due to wrong starting point

Problem: Any C function called after gen_prepare_non_leaf_call() have
the freedom to push values through `cfp->sp`, moving it, and then ask
for a zjit_materialize_frames(). (The C function has to restore
`cfp->sp` before it returns, but it can move it temporarily.) Since
zjit_materialize_frames() used `cfp->sp` as the starting point of
the cursor that writes out the stack, it wrote to the wrong slots.
In practice this happened when the non-leaf C function fails to make
a call and raises ArgumentError.

Solution: Stash `SP`, the VM stack base pointer in
gen_prepare_non_leaf_call() and add a new stack map opcode for
zjit_materialize_frames() to find it. Materialization during non-leaf
calls now reliably starts from `base_ptr+stack_size` and we don't look
at `cfp->sp`.

Author: XrXr

vm.c

@@ -2943,6 +2943,11 @@ zjit_materialize_frames(const rb_execution_context_t *ec, rb_control_frame_t *cf
                     else if (ZJIT_STACK_MAP_SKIP_P(entry)) {
                         stack -= ZJIT_STACK_MAP_SKIP_SIZE(entry);
                     }
+                    else if (ZJIT_STACK_MAP_BASE_PTR_P(entry)) {
+                        RUBY_ASESRT_ALWAYS(0 == i, "base_ptr stack map code only makes sense at 0");
+                        VALUE *base_ptr = (VALUE *)((VALUE *)cfp->jit_return)[-(ssize_t)ZJIT_STACK_MAP_BASE_PTR_SLOT_INDEX(entry)];
+                        stack = base_ptr + ZJIT_STACK_MAP_BASE_PTR_STACK_SIZE(entry);
+                    }
                     else {
                         stack--;
                         *stack = entry;

zjit.h

@@ -11,15 +11,24 @@
 # define ZJIT_STATS (USE_ZJIT && RUBY_DEBUG)
 #endif
 
-// Stack map entries are either immediate Ruby VALUEs, tagged native-stack
-// locations, or tagged skip counts. Stack maps never contain heap VALUEs, so
-// these tags are available: they are not Qfalse (0), and their low 3 bits are
-// zero, so RB_SPECIAL_CONST_P is false.
+// Stack map entries are opcodes for zjit_materialize_frames(), which walks them
+// in order while moving a cursor down the VM stack. An untagged entry is an
+// immediate Ruby VALUE to store; the tagged forms below copy from the native
+// stack, skip slots, or move the cursor. Stack maps never contain heap VALUEs,
+// so these tags are available: they are not Qfalse (0), and their low 3 bits
+// are zero, so RB_SPECIAL_CONST_P is false. Tags must stay non-zero multiples
+// of 8 for that to hold.
 #define ZJIT_STACK_MAP_VREG_TAG 0x08
 #define ZJIT_STACK_MAP_SKIP_TAG 0x10
+#define ZJIT_STACK_MAP_BASE_PTR_TAG 0x18
 #define ZJIT_STACK_MAP_TAG_MASK 0xff
 #define ZJIT_STACK_MAP_SHIFT 8
 
+// The BASE_PTR payload packs two fields above the tag byte: the slot index in
+// bits 8..31 and the operand stack size in bits 32..63.
+#define ZJIT_STACK_MAP_BASE_PTR_SIZE_SHIFT 32
+#define ZJIT_STACK_MAP_BASE_PTR_INDEX_MASK 0xffffff
+
 static inline bool
 ZJIT_STACK_MAP_VREG_P(VALUE entry)
 {
@@ -44,6 +53,36 @@ ZJIT_STACK_MAP_SKIP_SIZE(VALUE entry)
     return entry >> ZJIT_STACK_MAP_SHIFT;
 }
 
+// Anchor the write cursor using the SP register the JIT saved on its native
+// stack, instead of cfp->sp. cfp->sp is not a reliable starting point for a
+// frame that is in the middle of a non-leaf C call, as e.g. raising
+// ArgumentError can push through and move cfp->sp, then use the stack map.
+// gen_prepare_non_leaf_call() emits this opcode, always as stack[0], so the
+// entries after decode to the right place.
+static inline bool
+ZJIT_STACK_MAP_BASE_PTR_P(VALUE entry)
+{
+    return (entry & ZJIT_STACK_MAP_TAG_MASK) == ZJIT_STACK_MAP_BASE_PTR_TAG;
+}
+
+// VALUE index from cfp->jit_return down to the native stack slot holding the
+// saved SP register, i.e. base_ptr is `((VALUE **)cfp->jit_return)[-index]`.
+// There is one such slot per compiled function, so the index depends on the
+// frame's inlining depth.
+static inline size_t
+ZJIT_STACK_MAP_BASE_PTR_SLOT_INDEX(VALUE entry)
+{
+    return (entry >> ZJIT_STACK_MAP_SHIFT) & ZJIT_STACK_MAP_BASE_PTR_INDEX_MASK;
+}
+
+// Number of VM stack slots above base_ptr, i.e. the frame's operand stack
+// depth. The cursor is set to `base_ptr + this`.
+static inline size_t
+ZJIT_STACK_MAP_BASE_PTR_STACK_SIZE(VALUE entry)
+{
+    return entry >> ZJIT_STACK_MAP_BASE_PTR_SIZE_SHIFT;
+}
+
 // JITFrame is defined here as the single source of truth and imported into
 // Rust via bindgen. C code reads fields directly; Rust uses an impl block.
 typedef struct zjit_jit_frame {
@@ -61,9 +100,8 @@ typedef struct zjit_jit_frame {
 
     // Number of stack map entries in stack[].
     uint32_t stack_size;
-    // Flexible array of stack map entries. Each entry is either an immediate
-    // VALUE, a tagged native-stack index from cfp->jit_return for a value
-    // kept by the JIT, or a tagged count of VM stack slots to skip.
+    // Flexible array of stack map entries, executed in order by
+    // zjit_materialize_frames(). See the ZJIT_STACK_MAP_* opcodes above.
     VALUE stack[];
 } zjit_jit_frame_t;
 

zjit/bindgen/src/main.rs

@@ -344,6 +344,9 @@ fn main() {
         .allowlist_var("ZJIT_STACK_MAP_SHIFT")
         .allowlist_var("ZJIT_STACK_MAP_VREG_TAG")
         .allowlist_var("ZJIT_STACK_MAP_SKIP_TAG")
+        .allowlist_var("ZJIT_STACK_MAP_BASE_PTR_TAG")
+        .allowlist_var("ZJIT_STACK_MAP_BASE_PTR_SIZE_SHIFT")
+        .allowlist_var("ZJIT_STACK_MAP_BASE_PTR_INDEX_MASK")
         .allowlist_var("ZJIT_JIT_RETURN_C_FRAME")
         .allowlist_function("rb_assert_holding_vm_lock")
         .allowlist_function("rb_jit_shape_complex_p")

zjit/src/backend/lir.rs

@@ -5,7 +5,7 @@ use std::mem::take;
 use std::rc::Rc;
 use crate::bitset::BitSet;
 use crate::codegen::{perf_symbol_range_start, perf_symbol_range_end, register_with_perf};
-use crate::cruby::{IseqPtr, RUBY_OFFSET_CFP_ISEQ, RUBY_OFFSET_CFP_JIT_RETURN, RUBY_OFFSET_CFP_PC, RUBY_OFFSET_CFP_SP, SIZEOF_VALUE_I32, VALUE, ZJIT_STACK_MAP_SHIFT, ZJIT_STACK_MAP_SKIP_TAG, ZJIT_STACK_MAP_VREG_TAG, vm_stack_canary, YarvInsnIdx, zjit_jit_frame, local_size_and_idx_to_ep_offset};
+use crate::cruby::{IseqPtr, RUBY_OFFSET_CFP_ISEQ, RUBY_OFFSET_CFP_JIT_RETURN, RUBY_OFFSET_CFP_PC, RUBY_OFFSET_CFP_SP, SIZEOF_VALUE_I32, VALUE, ZJIT_STACK_MAP_BASE_PTR_INDEX_MASK, ZJIT_STACK_MAP_BASE_PTR_SIZE_SHIFT, ZJIT_STACK_MAP_BASE_PTR_TAG, ZJIT_STACK_MAP_SHIFT, ZJIT_STACK_MAP_SKIP_TAG, ZJIT_STACK_MAP_VREG_TAG, vm_stack_canary, YarvInsnIdx, zjit_jit_frame, local_size_and_idx_to_ep_offset};
 use crate::hir::{Invariant, SideExitReason};
 use crate::hir;
 use crate::options::{TraceExits, PerfMap, get_option};
@@ -1529,7 +1529,10 @@ const JIT_FRAME_OFFSET_FROM_JIT_RETURN: usize = 1;
 ///                                | +-------------------------+ |
 ///                                | |          ...            | | JITState::jit_frame_size
 ///                 stack_base_idx | +-------------------------+ |
-///                                | | JITFrame slot depth X   | v
+///                                | | JITFrame slot depth X   | |
+///                                | +-------------------------+ |
+///                                | | saved SP (stack map     | | <-- one per function, not per depth
+///                                | | anchor, base_ptr)       | v     see base_ptr_slot_offset()
 ///                                | +-------------------------+
 ///                                | | opnds.last()            | ^
 ///                                | +-------------------------+ |
@@ -1656,13 +1659,52 @@ impl StackMap {
     }
 }
 
-/// Entry in a JITFrame stack map.
+/// Entry in a JITFrame stack map. These are opcodes for zjit_materialize_frames(),
+/// which walks them in order moving a write cursor down the VM stack.
 #[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
 pub enum StackMapEntry {
     /// Immediate Ruby VALUE or VReg to materialize.
     Opnd(Opnd),
     /// Number of VM stack slots to skip when materializing across inlined frames.
     Skip(usize),
+    /// Anchor the write cursor on the SP register saved on the native stack
+    /// at `cfp->jit_return[-slot_index]`, plus `stack_size` VM slots. Emitted as
+    /// the first entry by gen_prepare_non_leaf_call(); see zjit.h.
+    BasePtr { slot_index: u32, stack_size: u32 },
+}
+
+/// The base_ptr payload splits in two above the tag byte, so the slot index has
+/// to occupy exactly the bits between the tag and the stack size field.
+const _: () = assert!(
+    ZJIT_STACK_MAP_BASE_PTR_INDEX_MASK as u64
+        == (1u64 << (ZJIT_STACK_MAP_BASE_PTR_SIZE_SHIFT - ZJIT_STACK_MAP_SHIFT)) - 1,
+    "BASE_PTR index mask must cover the bits between the tag byte and the stack size field",
+);
+
+/// The stack size field runs from its shift to the top of the VALUE, so a u32
+/// always fits and StackMapEntry::BasePtr needs no runtime check for it.
+const _: () = assert!(
+    ZJIT_STACK_MAP_BASE_PTR_SIZE_SHIFT + u32::BITS <= usize::BITS,
+    "BASE_PTR stack size field must be at least 32 bits wide",
+);
+
+impl StackMapEntry {
+    /// Encode a [`StackMapEntry::BasePtr`] into its tagged VALUE form.
+    /// `stack_size` is bound by [`CompileError::IseqStackTooLarge`]
+    /// `slot_index` by the max inline iteration.
+    // TODO(alan): bounds check max inline iteration as it's user-controlled.
+    fn encode_base_ptr(slot_index: u32, stack_size: u32) -> VALUE {
+        const INDEX_BITS: u32 = ZJIT_STACK_MAP_BASE_PTR_SIZE_SHIFT - ZJIT_STACK_MAP_SHIFT;
+        assert!(
+            slot_index <= ZJIT_STACK_MAP_BASE_PTR_INDEX_MASK,
+            "StackMap base_ptr slot index {slot_index} does not fit in {INDEX_BITS} bits",
+        );
+        let encoded = ((stack_size as usize) << ZJIT_STACK_MAP_BASE_PTR_SIZE_SHIFT)
+            | ((slot_index as usize) << ZJIT_STACK_MAP_SHIFT)
+            | ZJIT_STACK_MAP_BASE_PTR_TAG as usize;
+        debug_assert!(!VALUE(encoded).special_const_p(), "encoded StackMap base_ptr should not look like an immediate VALUE");
+        VALUE(encoded)
+    }
 }
 
 /// Initial capacity for asm.insns vector
@@ -2530,6 +2572,10 @@ impl Assembler
                                     debug_assert!(!VALUE(encoded).special_const_p(), "encoded StackMap skip should not look like an immediate VALUE");
                                     VALUE(encoded)
                                 }
+                                StackMapEntry::BasePtr { slot_index, stack_size } => {
+                                    debug_assert_eq!(idx, 0, "base_ptr must be the first StackMap entry so later entries decode from it");
+                                    StackMapEntry::encode_base_ptr(slot_index, stack_size)
+                                }
                                 StackMapEntry::Opnd(Opnd::VReg { idx: vreg, .. }) => {
                                     let vreg_stack_index = match assignments[vreg].expect("StackMap VReg should have an allocation") {
                                         Allocation::Reg(_) | Allocation::Fixed(_) => {

zjit/src/codegen.rs

@@ -61,9 +61,10 @@ struct JITState {
     iseq_calls: Vec<IseqCallRef>,
 
     /// The number of native stack slots reserved for JITFrame, one per
-    /// simultaneously live frame (`inlining_depth() + 1`). gen_write_jit_frame()
-    /// and the inlined frame push write a JITFrame into the slot selected by the
-    /// current frame's depth.
+    /// simultaneously live frame (`inlining_depth() + 1`), plus one shared slot
+    /// for the saved SP register at the bottom. gen_write_jit_frame() and the
+    /// inlined frame push write a JITFrame into the slot selected by the current
+    /// frame's depth; gen_prepare_non_leaf_call() writes the SP slot.
     jit_frame_size: usize,
 }
 
@@ -107,6 +108,18 @@ impl JITState {
         }
     }
 
+    /// Byte offset from [NATIVE_BASE_PTR] to the slot holding the SP VM stack base pointer.
+    fn base_ptr_slot_native_base_ptr_offset(&self) -> i32 {
+        -(i32::try_from(SIZEOF_VALUE * self.jit_frame_size).expect("base_ptr_slot_index overflow"))
+    }
+
+    /// The VALUE index a frame at `depth` uses to read the saved SP register as
+    /// `((VALUE **)cfp->jit_return)[-index]`. Distance between this inline frame's
+    /// `jit_return` and first slot past all `(jit_return, jit_frame)` tuples.
+    /// Encoded into [`StackMapEntry::BasePtr`].
+    fn base_ptr_slot_index(&self, depth: InlineDepth) -> u32 {
+        (self.jit_frame_size - depth).try_into().expect("base_ptr slot index overflow")
+    }
 }
 
 impl Assembler {
@@ -406,7 +419,10 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, version: IseqVersionRef, func
         // frame push select among these slots by the frame's depth, keeping each
         // frame's `cfp->jit_return` pointed at its own slot rather than a shared
         // one.
-        let jit_frame_size = function.inlining_depth() + 1;
+        //
+        // One more slot below those holds the saved SP register that stack maps
+        // are anchored on (see base_ptr_slot_offset()).
+        let jit_frame_size = function.inlining_depth() + 2;
         let mut jit = JITState::new(version, function.num_insns(), function.num_blocks(), jit_frame_size);
         let mut asm = Assembler::new_with_stack_slots(jit_frame_size);
 
@@ -3243,8 +3259,9 @@ pub(crate) fn block_iseq_may_throw(iseq: IseqPtr) -> bool {
 /// Byte offset from NATIVE_BASE_PTR of the JITFrame storage slot for a frame at
 /// the given inlining depth. Depth 0 (the top-level frame) lives at
 /// `[NATIVE_BASE_PTR - 8]`; each deeper inlined frame gets the next slot below.
-/// gen_function() reserves `inlining_depth() + 1` slots, so every live frame's
-/// depth maps to a distinct slot inside that reserved region.
+/// gen_function() reserves `inlining_depth() + 1` of these, so every live
+/// frame's depth maps to a distinct slot, followed by the single saved-SP slot
+/// of base_ptr_slot_offset().
 fn jit_frame_slot_offset(depth: InlineDepth) -> i32 {
     -(SIZEOF_VALUE_I32 * (depth as i32 + 1))
 }
@@ -3365,6 +3382,8 @@ fn gen_spill_stack(jit: &JITState, asm: &mut Assembler, function: &Function, sta
             StackMapEntry::Skip(skip) => {
                 offset -= skip as i32;
             }
+            // Only gen_prepare_non_leaf_call() prepends this, and it doesn't spill.
+            StackMapEntry::BasePtr { .. } => unreachable!("build_stack_map() does not emit BasePtr"),
         }
     }
 }
@@ -3416,12 +3435,21 @@ fn inline_frame_stack_gap(iseq: IseqPtr) -> 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) {
-    // TODO: Lazily materialize caller frames when needed
-    // Save PC for backtraces and allocation tracing
-    // and SP to avoid marking uninitialized stack slots
-    let stack_map = build_stack_map(jit, function, state);
+    // Anchor the stack map on a private copy of SP rather than on cfp->sp. The callee is free to
+    // use the stack map after pushing through and moving cfp->sp (e.g. rb_funcall() + a raise in
+    // vm_callee_setup_arg()).
+    let mut stack_map = vec![StackMapEntry::BasePtr {
+        slot_index: jit.base_ptr_slot_index(state.depth),
+        stack_size: state.stack_size().try_into().expect("stack size overflow"),
+    }];
+    stack_map.extend(build_stack_map(jit, function, state));
     let jit_frame = gen_prepare_call_with_gc(asm, state, false, stack_map.len());
 
+    // NOTE(alan): This store can be done once on function entry, but analysis is required
+    //             to avoid the store in functions that make no non-leaf call.
+    asm_comment!(asm, "save SP as the stack map anchor");
+    asm.mov(Opnd::mem(64, NATIVE_BASE_PTR, jit.base_ptr_slot_native_base_ptr_offset()), SP);
+
     // 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);

zjit/src/cruby_bindings.inc.rs

@@ -292,6 +292,55 @@ mod tests {
         "#), @r#""no bar""#);
     }
 
+    // A C function that calls back into Ruby (rb_const_missing here) pushes
+    // recv + args onto *this* frame's VM stack via vm_call0_body(). A wrong
+    // arity makes vm_callee_setup_arg() raise before vm_call_iseq_setup_normal()
+    // restores cfp->sp, so cfp->sp is left 1 + argc slots high while this frame
+    // is materialized for the rescue. The stack map has to be decoded from the
+    // JIT's saved SP instead, or the live `1` operand lands above its slot and
+    // the array comes back as [false, 2].
+    #[test]
+    fn test_stack_map_anchor_after_callee_arity_error() {
+        assert_snapshot!(inspect(r#"
+            class Holder
+              def self.const_missing(a, b) = nil # wrong arity: called with 1 arg
+            end
+            def jit_entry
+              [1, (begin      # the 1 is live across the const_missing call
+                     Holder::NOPE
+                   rescue ArgumentError
+                     2
+                   end)]
+            end
+            jit_entry
+            jit_entry
+        "#), @"[1, 2]");
+    }
+
+    // Same displaced cfp->sp, but reaching across an inlined frame: `defined?`
+    // calls respond_to_missing? with 2 args, which raises on arity. The map for
+    // the inlined `inner` frame keeps going past its frame gap into `outer`'s
+    // operand stack, so a displaced anchor writes `outer`'s live 1 above its
+    // slot and leaves the real one holding garbage. Giving `inner` a local
+    // pushes that garbage into env data instead, which crashes rather than
+    // returning a wrong answer.
+    #[test]
+    fn test_stack_map_anchor_with_inlined_frame() {
+        assert_snapshot!(inspect(r#"
+            class BadResponder
+              def respond_to_missing?(name) = true # wrong arity: called with 2
+            end
+            class Test
+              def initialize = @o = BadResponder.new
+              def inner = defined?(@o.nope) # must have 0 locals
+              def outer = [1, inner] # the 1 is live across the inlined call
+            end
+            test = Test.new
+            test.outer
+            test.outer
+        "#), @"[1, nil]");
+    }
+
     // Proc.new inside a block passed via invokeblock captures the caller's
     // block_code. When the JIT compiles the caller, block_code must be
     // correctly available for the proc to work.