[pull] master from ruby:master

ext/objspace/objspace_dump.c

@@ -801,7 +801,7 @@ shape_id_i(shape_id_t shape_id, void *data)
 
     if (RSHAPE_TYPE(shape_id) != SHAPE_ROOT) {
         dump_append(dc, ", \"parent_id\":");
-        dump_append_lu(dc, RSHAPE_PARENT(shape_id));
+        dump_append_lu(dc, RSHAPE_PARENT_RAW_ID(shape_id));
     }
 
     dump_append(dc, ", \"depth\":");

jit.c

@@ -442,3 +442,15 @@ rb_yarv_ary_entry_internal(VALUE ary, long offset)
 {
     return rb_ary_entry_internal(ary, offset);
 }
+
+void
+rb_set_cfp_pc(struct rb_control_frame_struct *cfp, const VALUE *pc)
+{
+    cfp->pc = pc;
+}
+
+void
+rb_set_cfp_sp(struct rb_control_frame_struct *cfp, VALUE *sp)
+{
+    cfp->sp = sp;
+}

shape.h

@@ -271,11 +271,18 @@ rb_shape_root(size_t heap_id)
 }
 
 static inline shape_id_t
-RSHAPE_PARENT(shape_id_t shape_id)
+RSHAPE_PARENT_RAW_ID(shape_id_t shape_id)
 {
     return RSHAPE(shape_id)->parent_id;
 }
 
+static inline bool
+RSHAPE_DIRECT_CHILD_P(shape_id_t parent_id, shape_id_t child_id)
+{
+    return (parent_id & SHAPE_ID_FLAGS_MASK) == (child_id & SHAPE_ID_FLAGS_MASK) &&
+        RSHAPE(child_id)->parent_id == (parent_id & SHAPE_ID_OFFSET_MASK);
+}
+
 static inline enum shape_type
 RSHAPE_TYPE(shape_id_t shape_id)
 {

variable.c

@@ -1774,7 +1774,9 @@ general_ivar_set(VALUE obj, ID id, VALUE val, void *data,
             shape_resize_fields_func(obj, RSHAPE_CAPACITY(current_shape_id), RSHAPE_CAPACITY(next_shape_id), data);
         }
 
-        RUBY_ASSERT(RSHAPE_TYPE_P(next_shape_id, SHAPE_IVAR));
+        RUBY_ASSERT(RSHAPE_TYPE_P(next_shape_id, SHAPE_IVAR),
+                    "next_shape_id: 0x%" PRIx32 " RSHAPE_TYPE(next_shape_id): %d",
+                    next_shape_id, (int)RSHAPE_TYPE(next_shape_id));
         RUBY_ASSERT(index == (RSHAPE_INDEX(next_shape_id)));
         set_shape_id_func(obj, next_shape_id, data);
     }

vm_insnhelper.c

@@ -1473,7 +1473,7 @@ vm_setivar_default(VALUE obj, ID id, VALUE val, shape_id_t dest_shape_id, attr_i
         RUBY_ASSERT(dest_shape_id != INVALID_SHAPE_ID && shape_id != INVALID_SHAPE_ID);
     }
     else if (dest_shape_id != INVALID_SHAPE_ID) {
-        if (shape_id == RSHAPE_PARENT(dest_shape_id) && RSHAPE_EDGE_NAME(dest_shape_id) == id && RSHAPE_CAPACITY(shape_id) == RSHAPE_CAPACITY(dest_shape_id)) {
+        if (RSHAPE_DIRECT_CHILD_P(shape_id, dest_shape_id) && RSHAPE_EDGE_NAME(dest_shape_id) == id && RSHAPE_CAPACITY(shape_id) == RSHAPE_CAPACITY(dest_shape_id)) {
             RUBY_ASSERT(index < RSHAPE_CAPACITY(dest_shape_id));
         }
         else {
@@ -1514,14 +1514,11 @@ vm_setivar(VALUE obj, ID id, VALUE val, shape_id_t dest_shape_id, attr_index_t i
                 VM_ASSERT(!rb_ractor_shareable_p(obj));
             }
             else if (dest_shape_id != INVALID_SHAPE_ID) {
-                shape_id_t source_shape_id = RSHAPE_PARENT(dest_shape_id);
-
-                if (shape_id == source_shape_id && RSHAPE_EDGE_NAME(dest_shape_id) == id && RSHAPE_CAPACITY(shape_id) == RSHAPE_CAPACITY(dest_shape_id)) {
+                if (RSHAPE_DIRECT_CHILD_P(shape_id, dest_shape_id) && RSHAPE_EDGE_NAME(dest_shape_id) == id && RSHAPE_CAPACITY(shape_id) == RSHAPE_CAPACITY(dest_shape_id)) {
                     RUBY_ASSERT(dest_shape_id != INVALID_SHAPE_ID && shape_id != INVALID_SHAPE_ID);
 
                     RBASIC_SET_SHAPE_ID(obj, dest_shape_id);
 
-                    RUBY_ASSERT(rb_shape_get_next_iv_shape(source_shape_id, id) == dest_shape_id);
                     RUBY_ASSERT(index < RSHAPE_CAPACITY(dest_shape_id));
                 }
                 else {

yjit.c

@@ -499,18 +499,6 @@ rb_yjit_str_simple_append(VALUE str1, VALUE str2)
     return rb_str_cat(str1, RSTRING_PTR(str2), RSTRING_LEN(str2));
 }
 
-void
-rb_set_cfp_pc(struct rb_control_frame_struct *cfp, const VALUE *pc)
-{
-    cfp->pc = pc;
-}
-
-void
-rb_set_cfp_sp(struct rb_control_frame_struct *cfp, VALUE *sp)
-{
-    cfp->sp = sp;
-}
-
 extern VALUE *rb_vm_base_ptr(struct rb_control_frame_struct *cfp);
 
 // YJIT needs this function to never allocate and never raise

zjit/src/codegen.rs

@@ -1,6 +1,6 @@
 use std::cell::Cell;
 use std::rc::Rc;
-use std::ffi::{c_int};
+use std::ffi::{c_int, c_void};
 
 use crate::asm::Label;
 use crate::backend::current::{Reg, ALLOC_REGS};
@@ -105,53 +105,33 @@ fn gen_iseq_entry_point(iseq: IseqPtr) -> *const u8 {
     let code_ptr = gen_iseq_entry_point_body(cb, iseq);
 
     // Always mark the code region executable if asm.compile() has been used.
-    // We need to do this even if code_ptr is null because, whether gen_entry()
-    // or gen_iseq() fails or not, gen_function() has already used asm.compile().
+    // We need to do this even if code_ptr is null because, whether gen_entry() or
+    // gen_function_stub() fails or not, gen_function() has already used asm.compile().
     cb.mark_all_executable();
 
-    code_ptr
+    code_ptr.map_or(std::ptr::null(), |ptr| ptr.raw_ptr(cb))
 }
 
 /// Compile an entry point for a given ISEQ
-fn gen_iseq_entry_point_body(cb: &mut CodeBlock, iseq: IseqPtr) -> *const u8 {
+fn gen_iseq_entry_point_body(cb: &mut CodeBlock, iseq: IseqPtr) -> Option<CodePtr> {
     // Compile ISEQ into High-level IR
-    let function = match compile_iseq(iseq) {
-        Some(function) => function,
-        None => return std::ptr::null(),
-    };
+    let function = compile_iseq(iseq)?;
 
     // Compile the High-level IR
     let Some((start_ptr, gc_offsets, jit)) = gen_function(cb, iseq, &function) else {
         debug!("Failed to compile iseq: gen_function failed: {}", iseq_get_location(iseq, 0));
-        return std::ptr::null();
+        return None;
     };
 
     // Compile an entry point to the JIT code
     let Some(entry_ptr) = gen_entry(cb, iseq, &function, start_ptr) else {
         debug!("Failed to compile iseq: gen_entry failed: {}", iseq_get_location(iseq, 0));
-        return std::ptr::null();
+        return None;
     };
 
-    let mut branch_iseqs = jit.branch_iseqs;
-
-    // Recursively compile callee ISEQs
-    let caller_iseq = iseq;
-    while let Some((branch, iseq)) = branch_iseqs.pop() {
-        // Disable profiling. This will be the last use of the profiling information for the ISEQ.
-        unsafe { rb_zjit_profile_disable(iseq); }
-
-        // Compile the ISEQ
-        let Some((callee_ptr, callee_branch_iseqs)) = gen_iseq(cb, iseq) else {
-            // Failed to compile the callee. Bail out of compiling this graph of ISEQs.
-            debug!("Failed to compile iseq: could not compile callee: {} -> {}",
-                   iseq_get_location(caller_iseq, 0), iseq_get_location(iseq, 0));
-            return std::ptr::null();
-        };
-        let callee_addr = callee_ptr.raw_ptr(cb);
-        branch.regenerate(cb, |asm| {
-            asm.ccall(callee_addr, vec![]);
-        });
-        branch_iseqs.extend(callee_branch_iseqs);
+    // Stub callee ISEQs for JIT-to-JIT calls
+    for (branch, callee_iseq) in jit.branch_iseqs.into_iter() {
+        gen_iseq_branch(cb, callee_iseq, iseq, branch)?;
     }
 
     // Remember the block address to reuse it later
@@ -160,7 +140,27 @@ fn gen_iseq_entry_point_body(cb: &mut CodeBlock, iseq: IseqPtr) -> *const u8 {
     append_gc_offsets(iseq, &gc_offsets);
 
     // Return a JIT code address
-    entry_ptr.raw_ptr(cb)
+    Some(entry_ptr)
+}
+
+/// Stub a branch for a JIT-to-JIT call
+fn gen_iseq_branch(cb: &mut CodeBlock, iseq: IseqPtr, caller_iseq: IseqPtr, branch: Rc<Branch>) -> Option<()> {
+    // Compile a function stub
+    let Some((stub_ptr, gc_offsets)) = gen_function_stub(cb, iseq, branch.clone()) else {
+        // Failed to compile the stub. Bail out of compiling the caller ISEQ.
+        debug!("Failed to compile iseq: could not compile stub: {} -> {}",
+               iseq_get_location(caller_iseq, 0), iseq_get_location(iseq, 0));
+        return None;
+    };
+    append_gc_offsets(iseq, &gc_offsets);
+
+    // Update the JIT-to-JIT call to call the stub
+    let stub_addr = stub_ptr.raw_ptr(cb);
+    branch.regenerate(cb, |asm| {
+        asm_comment!(asm, "call function stub: {}", iseq_get_location(iseq, 0));
+        asm.ccall(stub_addr, vec![]);
+    });
+    Some(())
 }
 
 /// Write an entry to the perf map in /tmp
@@ -244,7 +244,11 @@ fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function) -> Optio
     let reverse_post_order = function.rpo();
     for &block_id in reverse_post_order.iter() {
         let block = function.block(block_id);
-        asm_comment!(asm, "Block: {block_id}({})", block.params().map(|param| format!("{param}")).collect::<Vec<_>>().join(", "));
+        asm_comment!(
+            asm, "{block_id}({}): {}",
+            block.params().map(|param| format!("{param}")).collect::<Vec<_>>().join(", "),
+            iseq_get_location(iseq, block.insn_idx),
+        );
 
         // Write a label to jump to the basic block
         let label = jit.get_label(&mut asm, block_id);
@@ -1263,6 +1267,127 @@ fn max_num_params(function: &Function) -> usize {
     }).max().unwrap_or(0)
 }
 
+#[cfg(target_arch = "x86_64")]
+macro_rules! c_callable {
+    ($(#[$outer:meta])*
+    fn $f:ident $args:tt $(-> $ret:ty)? $body:block) => {
+        $(#[$outer])*
+        extern "sysv64" fn $f $args $(-> $ret)? $body
+    };
+}
+#[cfg(target_arch = "aarch64")]
+macro_rules! c_callable {
+    ($(#[$outer:meta])*
+    fn $f:ident $args:tt $(-> $ret:ty)? $body:block) => {
+        $(#[$outer])*
+        extern "C" fn $f $args $(-> $ret)? $body
+    };
+}
+pub(crate) use c_callable;
+
+c_callable! {
+    /// Generated code calls this function with the SysV calling convention.
+    /// See [gen_function_stub].
+    fn function_stub_hit(iseq: IseqPtr, branch_ptr: *const c_void, ec: EcPtr, sp: *mut VALUE) -> *const u8 {
+        with_vm_lock(src_loc!(), || {
+            // Get a pointer to compiled code or the side-exit trampoline
+            let cb = ZJITState::get_code_block();
+            let code_ptr = if let Some(code_ptr) = function_stub_hit_body(cb, iseq, branch_ptr) {
+                code_ptr
+            } else {
+                // gen_push_frame() doesn't set PC and SP, so we need to set them for side-exit
+                // TODO: We could generate code that sets PC/SP. Note that we'd still need to handle OOM.
+                let cfp = unsafe { get_ec_cfp(ec) };
+                let pc = unsafe { rb_iseq_pc_at_idx(iseq, 0) }; // TODO: handle opt_pc once supported
+                unsafe { rb_set_cfp_pc(cfp, pc) };
+                unsafe { rb_set_cfp_sp(cfp, sp) };
+
+                // Exit to the interpreter
+                ZJITState::get_stub_exit()
+            };
+
+            cb.mark_all_executable();
+            code_ptr.raw_ptr(cb)
+        })
+    }
+}
+
+/// Compile an ISEQ for a function stub
+fn function_stub_hit_body(cb: &mut CodeBlock, iseq: IseqPtr, branch_ptr: *const c_void) -> Option<CodePtr> {
+    // Compile the stubbed ISEQ
+    let Some((code_ptr, branch_iseqs)) = gen_iseq(cb, iseq) else {
+        debug!("Failed to compile iseq: gen_iseq failed: {}", iseq_get_location(iseq, 0));
+        return None;
+    };
+
+    // Stub callee ISEQs for JIT-to-JIT calls
+    for (branch, callee_iseq) in branch_iseqs.into_iter() {
+        gen_iseq_branch(cb, callee_iseq, iseq, branch)?;
+    }
+
+    // Update the stub to call the code pointer
+    let branch = unsafe { Rc::from_raw(branch_ptr as *const Branch) };
+    let code_addr = code_ptr.raw_ptr(cb);
+    branch.regenerate(cb, |asm| {
+        asm_comment!(asm, "call compiled function: {}", iseq_get_location(iseq, 0));
+        asm.ccall(code_addr, vec![]);
+    });
+
+    Some(code_ptr)
+}
+
+/// Compile a stub for an ISEQ called by SendWithoutBlockDirect
+/// TODO: Consider creating a trampoline to share some of the code among function stubs
+fn gen_function_stub(cb: &mut CodeBlock, iseq: IseqPtr, branch: Rc<Branch>) -> Option<(CodePtr, Vec<CodePtr>)> {
+    let mut asm = Assembler::new();
+    asm_comment!(asm, "Stub: {}", iseq_get_location(iseq, 0));
+
+    // Maintain alignment for x86_64, and set up a frame for arm64 properly
+    asm.frame_setup(&[], 0);
+
+    asm_comment!(asm, "preserve argument registers");
+    for &reg in ALLOC_REGS.iter() {
+        asm.cpush(Opnd::Reg(reg));
+    }
+    const { assert!(ALLOC_REGS.len() % 2 == 0, "x86_64 would need to push one more if we push an odd number of regs"); }
+
+    // Compile the stubbed ISEQ
+    let branch_addr = Rc::into_raw(branch);
+    let jump_addr = asm_ccall!(asm, function_stub_hit,
+        Opnd::Value(iseq.into()),
+        Opnd::const_ptr(branch_addr as *const u8),
+        EC,
+        SP
+    );
+    asm.mov(Opnd::Reg(Assembler::SCRATCH_REG), jump_addr);
+
+    asm_comment!(asm, "restore argument registers");
+    for &reg in ALLOC_REGS.iter().rev() {
+        asm.cpop_into(Opnd::Reg(reg));
+    }
+
+    // Discard the current frame since the JIT function will set it up again
+    asm.frame_teardown(&[]);
+
+    // Jump to SCRATCH_REG so that cpop_all() doesn't clobber it
+    asm.jmp_opnd(Opnd::Reg(Assembler::SCRATCH_REG));
+    asm.compile(cb)
+}
+
+/// Generate a trampoline that is used when a function stub fails to compile the ISEQ
+pub fn gen_stub_exit(cb: &mut CodeBlock) -> Option<CodePtr> {
+    let mut asm = Assembler::new();
+
+    asm_comment!(asm, "exit from function stub");
+    asm.frame_teardown(lir::JIT_PRESERVED_REGS);
+    asm.cret(Qundef.into());
+
+    asm.compile(cb).map(|(code_ptr, gc_offsets)| {
+        assert_eq!(gc_offsets.len(), 0);
+        code_ptr
+    })
+}
+
 impl Assembler {
     /// Make a C call while marking the start and end positions of it
     fn ccall_with_branch(&mut self, fptr: *const u8, opnds: Vec<Opnd>, branch: &Rc<Branch>) -> Opnd {

zjit/src/cruby.rs

@@ -715,7 +715,7 @@ pub fn iseq_name(iseq: IseqPtr) -> String {
 // Location is the file defining the method, colon, method name.
 // Filenames are sometimes internal strings supplied to eval,
 // so be careful with them.
-pub fn iseq_get_location(iseq: IseqPtr, pos: u16) -> String {
+pub fn iseq_get_location(iseq: IseqPtr, pos: u32) -> String {
     let iseq_path = unsafe { rb_iseq_path(iseq) };
     let iseq_lineno = unsafe { rb_iseq_line_no(iseq, pos as usize) };