ZJIT: Split Send into Lookup+Call

tekknolagi · tekknolagi · commit 99741491661c · 2025-06-06T13:09:41.000-04:00
diff --git a/jit.c b/jit.c
@@ -421,3 +421,16 @@ rb_assert_cme_handle(VALUE handle)
     RUBY_ASSERT_ALWAYS(!rb_objspace_garbage_object_p(handle));
     RUBY_ASSERT_ALWAYS(IMEMO_TYPE_P(handle, imemo_ment));
 }
+
+extern
+VALUE
+rb_vm_call0(rb_execution_context_t *ec, VALUE recv, ID id, int argc, const VALUE *argv, const rb_callable_method_entry_t *cme, int kw_splat);
+
+VALUE
+rb_zjit_vm_call0_no_splat(rb_execution_context_t *ec, VALUE recv, ID id, int argc, const VALUE *argv, const rb_callable_method_entry_t *cme) {
+    return rb_vm_call0(ec, recv, id, argc, argv, cme, /*kw_splat=*/0);
+}
+
+extern
+VALUE
+rb_vm_send(rb_execution_context_t *ec, rb_control_frame_t *reg_cfp, CALL_DATA cd, ISEQ blockiseq);
diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs
@@ -347,6 +347,8 @@ fn main() {
         .allowlist_function("rb_optimized_call")
         .allowlist_function("rb_zjit_icache_invalidate")
         .allowlist_function("rb_zjit_print_exception")
+        .allowlist_function("rb_zjit_vm_call0_no_splat")
+        .allowlist_function("rb_vm_send")
         .allowlist_type("robject_offsets")
         .allowlist_type("rstring_offsets")
 
diff --git a/zjit/src/backend/arm64/mod.rs b/zjit/src/backend/arm64/mod.rs
@@ -28,6 +28,7 @@ pub const _C_ARG_OPNDS: [Opnd; 6] = [
 // C return value register on this platform
 pub const C_RET_REG: Reg = X0_REG;
 pub const _C_RET_OPND: Opnd = Opnd::Reg(X0_REG);
+pub const _NATIVE_STACK_PTR: Opnd = Opnd::Reg(X13_REG);
 
 // These constants define the way we work with Arm64's stack pointer. The stack
 // pointer always needs to be aligned to a 16-byte boundary.
@@ -679,6 +680,10 @@ impl Assembler
                     *opnd = opnd0;
                     asm.push_insn(insn);
                 },
+                Insn::CPush(opnd) => {
+                    let opnd = split_load_operand(asm, *opnd);
+                    asm.push_insn(Insn::CPush(opnd));
+                }
                 Insn::Store { dest, src } => {
                     // The value being stored must be in a register, so if it's
                     // not already one we'll load it first.
diff --git a/zjit/src/backend/lir.rs b/zjit/src/backend/lir.rs
@@ -16,6 +16,7 @@ pub const SP: Opnd = _SP;
 
 pub const C_ARG_OPNDS: [Opnd; 6] = _C_ARG_OPNDS;
 pub const C_RET_OPND: Opnd = _C_RET_OPND;
+pub const NATIVE_STACK_PTR: Opnd = _NATIVE_STACK_PTR;
 pub use crate::backend::current::{Reg, C_RET_REG};
 
 // Memory operand base
@@ -2225,7 +2226,7 @@ impl Assembler {
 /// when not dumping disassembly.
 macro_rules! asm_comment {
     ($asm:expr, $($fmt:tt)*) => {
-        if $crate::options::get_option!(dump_disasm) {
+        if $crate::options::get_option!(dump_disasm) || $crate::options::get_option!(dump_lir) {
             $asm.push_insn(crate::backend::lir::Insn::Comment(format!($($fmt)*)));
         }
     };
diff --git a/zjit/src/backend/x86_64/mod.rs b/zjit/src/backend/x86_64/mod.rs
@@ -28,6 +28,7 @@ pub const _C_ARG_OPNDS: [Opnd; 6] = [
 // C return value register on this platform
 pub const C_RET_REG: Reg = RAX_REG;
 pub const _C_RET_OPND: Opnd = Opnd::Reg(RAX_REG);
+pub const _NATIVE_STACK_PTR: Opnd = Opnd::Reg(RSP_REG);
 
 impl CodeBlock {
     // The number of bytes that are generated by jmp_ptr
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
@@ -6,8 +6,8 @@ use crate::profile::get_or_create_iseq_payload;
 use crate::state::ZJITState;
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::invariants::{iseq_escapes_ep, track_no_ep_escape_assumption};
-use crate::backend::lir::{self, asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, SP};
-use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, CallInfo, RangeType, SELF_PARAM_IDX};
+use crate::backend::lir::{self, asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, SP, NATIVE_STACK_PTR};
+use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, RangeType, SELF_PARAM_IDX};
 use crate::hir::{Const, FrameState, Function, Insn, InsnId};
 use crate::hir_type::{types::Fixnum, Type};
 use crate::options::get_option;
@@ -257,8 +257,10 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::Jump(branch) => return gen_jump(jit, asm, branch),
         Insn::IfTrue { val, target } => return gen_if_true(jit, asm, opnd!(val), target),
         Insn::IfFalse { val, target } => return gen_if_false(jit, asm, opnd!(val), target),
-        Insn::SendWithoutBlock { call_info, cd, state, self_val, args, .. } => gen_send_without_block(jit, asm, call_info, *cd, &function.frame_state(*state), self_val, args)?,
-        Insn::SendWithoutBlockDirect { iseq, self_val, args, .. } => gen_send_without_block_direct(cb, jit, asm, *iseq, opnd!(self_val), args)?,
+        Insn::LookupMethod { self_val, method_id, state } => gen_lookup_method(jit, asm, opnd!(self_val), *method_id, &function.frame_state(*state))?,
+        Insn::CallMethod { callable, cd, self_val, args, state } => gen_call_method(jit, asm, opnd!(callable), *cd, opnd!(self_val), args, &function.frame_state(*state))?,
+        Insn::CallCFunc { cfunc, self_val, args, .. } => gen_call_cfunc(jit, asm, *cfunc, opnd!(self_val), args)?,
+        Insn::CallIseq { iseq, self_val, args, .. } => gen_send_without_block_direct(cb, jit, asm, *iseq, opnd!(self_val), args)?,
         Insn::Return { val } => return Some(gen_return(asm, opnd!(val))?),
         Insn::FixnumAdd { left, right, state } => gen_fixnum_add(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
         Insn::FixnumSub { left, right, state } => gen_fixnum_sub(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
@@ -455,38 +457,109 @@ fn gen_if_false(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, branch:
 }
 
 /// Compile a dynamic dispatch without block
-fn gen_send_without_block(
+fn gen_lookup_method(
     jit: &mut JITState,
     asm: &mut Assembler,
-    call_info: &CallInfo,
-    cd: *const rb_call_data,
+    recv: Opnd,
+    method_id: ID,
+    state: &FrameState,
+) -> Option<lir::Opnd> {
+    asm_comment!(asm, "get the class of the receiver with rb_obj_class");
+    let class = asm.ccall(rb_obj_class as *const u8, vec![recv]);
+    // TODO(max): Figure out if we need to do anything here to save state to CFP
+    let method_opnd = Opnd::UImm(method_id.0.into());
+    asm_comment!(asm, "call rb_callable_method_entry");
+    let result = asm.ccall(
+        rb_callable_method_entry as *const u8,
+        vec![class, method_opnd],
+    );
+    asm.test(result, result);
+    asm.jz(side_exit(jit, state)?);
+    Some(result)
+}
+
+fn gen_call_method(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    callable: Opnd,
+    cd: CallDataPtr,
+    recv: Opnd,
+    args: &Vec<InsnId>,
     state: &FrameState,
-    self_val: &InsnId,
+) -> Option<Opnd> {
+    // Don't push recv; it is passed in separately.
+    asm_comment!(asm, "make stack-allocated array of {} args", args.len());
+    for &arg in args.iter().rev() {
+        asm.cpush(jit.get_opnd(arg)?);
+    }
+    // Save PC for GC
+    gen_save_pc(asm, state);
+    // Call rb_zjit_vm_call0_no_splat, which will push a frame
+    // TODO(max): Figure out if we need to manually handle stack alignment and how to do it
+    let call_info = unsafe { rb_get_call_data_ci(cd) };
+    let method_id = unsafe { rb_vm_ci_mid(call_info) };
+    asm_comment!(asm, "get stack pointer");
+    let sp = asm.lea(Opnd::mem(VALUE_BITS, NATIVE_STACK_PTR, 0));
+    asm_comment!(asm, "call rb_zjit_vm_call0_no_splat");
+    let result = asm.ccall(
+        rb_zjit_vm_call0_no_splat as *const u8,
+        vec![EC, recv, Opnd::UImm(method_id.0), Opnd::UImm(args.len().try_into().unwrap()), sp, callable],
+    );
+    // Pop all the args off the stack
+    asm_comment!(asm, "clear stack-allocated array of {} args", args.len());
+    let new_sp = asm.add(NATIVE_STACK_PTR, (args.len()*SIZEOF_VALUE).into());
+    asm.mov(NATIVE_STACK_PTR, new_sp);
+    Some(result)
+}
+
+/// Compile an interpreter frame
+fn gen_push_frame(asm: &mut Assembler, recv: Opnd) {
+    // Write to a callee CFP
+    fn cfp_opnd(offset: i32) -> Opnd {
+        Opnd::mem(64, CFP, offset - (RUBY_SIZEOF_CONTROL_FRAME as i32))
+    }
+
+    asm_comment!(asm, "push callee control frame");
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SELF), recv);
+    // TODO: Write more fields as needed
+}
+
+fn gen_call_cfunc(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    cfunc: CFuncPtr,
+    recv: Opnd,
     args: &Vec<InsnId>,
 ) -> Option<lir::Opnd> {
-    // Spill the receiver and the arguments onto the stack. They need to be marked by GC and may be caller-saved registers.
-    // TODO: Avoid spilling operands that have been spilled before.
-    for (idx, &insn_id) in [*self_val].iter().chain(args.iter()).enumerate() {
-        // Currently, we don't move the SP register. So it's equal to the base pointer.
-        let stack_opnd = Opnd::mem(64, SP, idx as i32  * SIZEOF_VALUE_I32);
-        asm.mov(stack_opnd, jit.get_opnd(insn_id)?);
+    let cfunc_argc = unsafe { get_mct_argc(cfunc) };
+    // NB: The presence of self is assumed (no need for +1).
+    if args.len() != cfunc_argc as usize {
+        // TODO(max): We should check this at compile-time. If we have an arity mismatch at this
+        // point, we should side-exit (we're definitely going to raise) and if we don't, we should
+        // not check anything.
+        todo!("Arity mismatch");
     }
 
-    // Save PC and SP
-    gen_save_pc(asm, state);
-    gen_save_sp(asm, 1 + args.len()); // +1 for receiver
+    // Set up the new frame
+    gen_push_frame(asm, recv);
+
+    asm_comment!(asm, "switch to new CFP");
+    let new_cfp = asm.sub(CFP, RUBY_SIZEOF_CONTROL_FRAME.into());
+    asm.mov(CFP, new_cfp);
+    asm.store(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), CFP);
 
-    asm_comment!(asm, "call #{} with dynamic dispatch", call_info.method_name);
-    unsafe extern "C" {
-        fn rb_vm_opt_send_without_block(ec: EcPtr, cfp: CfpPtr, cd: VALUE) -> VALUE;
+    // Set up arguments
+    let mut c_args: Vec<Opnd> = vec![recv];
+    for &arg in args.iter() {
+        c_args.push(jit.get_opnd(arg)?);
     }
-    let ret = asm.ccall(
-        rb_vm_opt_send_without_block as *const u8,
-        vec![EC, CFP, (cd as usize).into()],
-    );
+
+    // Make a method call. The target address will be rewritten once compiled.
+    let cfun = unsafe { get_mct_func(cfunc) }.cast();
     // TODO(max): Add a PatchPoint here that can side-exit the function if the callee messed with
     // the frame's locals
-
+    let ret = asm.ccall(cfun, c_args);
+    gen_pop_frame(asm);
     Some(ret)
 }
 
@@ -590,14 +663,18 @@ fn gen_new_range(
     new_range
 }
 
-/// Compile code that exits from JIT code with a return value
-fn gen_return(asm: &mut Assembler, val: lir::Opnd) -> Option<()> {
+fn gen_pop_frame(asm: &mut Assembler) {
     // Pop the current frame (ec->cfp++)
     // Note: the return PC is already in the previous CFP
     asm_comment!(asm, "pop stack frame");
     let incr_cfp = asm.add(CFP, RUBY_SIZEOF_CONTROL_FRAME.into());
     asm.mov(CFP, incr_cfp);
     asm.mov(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), CFP);
+}
+
+/// Compile code that exits from JIT code with a return value
+fn gen_return(asm: &mut Assembler, val: lir::Opnd) -> Option<()> {
+    gen_pop_frame(asm);
 
     asm.frame_teardown();
 
diff --git a/zjit/src/cruby.rs b/zjit/src/cruby.rs
@@ -120,9 +120,9 @@ pub use autogened::*;
 #[cfg_attr(test, allow(unused))] // We don't link against C code when testing
 unsafe extern "C" {
     pub fn rb_check_overloaded_cme(
-        me: *const rb_callable_method_entry_t,
+        me: CmePtr,
         ci: *const rb_callinfo,
-    ) -> *const rb_callable_method_entry_t;
+    ) -> CmePtr;
 
     // Floats within range will be encoded without creating objects in the heap.
     // (Range is 0x3000000000000001 to 0x4fffffffffffffff (1.7272337110188893E-77 to 2.3158417847463237E+77).
@@ -144,8 +144,8 @@ unsafe extern "C" {
     pub fn rb_vm_set_ivar_id(obj: VALUE, idx: u32, val: VALUE) -> VALUE;
     pub fn rb_vm_setinstancevariable(iseq: IseqPtr, obj: VALUE, id: ID, val: VALUE, ic: IVC);
     pub fn rb_aliased_callable_method_entry(
-        me: *const rb_callable_method_entry_t,
-    ) -> *const rb_callable_method_entry_t;
+        me: CmePtr,
+    ) -> CmePtr;
     pub fn rb_vm_getclassvariable(iseq: IseqPtr, cfp: CfpPtr, id: ID, ic: ICVARC) -> VALUE;
     pub fn rb_vm_setclassvariable(
         iseq: IseqPtr,
@@ -156,7 +156,7 @@ unsafe extern "C" {
     ) -> VALUE;
     pub fn rb_vm_ic_hit_p(ic: IC, reg_ep: *const VALUE) -> bool;
     pub fn rb_vm_stack_canary() -> VALUE;
-    pub fn rb_vm_push_cfunc_frame(cme: *const rb_callable_method_entry_t, recv_idx: c_int);
+    pub fn rb_vm_push_cfunc_frame(cme: CmePtr, recv_idx: c_int);
 }
 
 // Renames
@@ -259,6 +259,12 @@ pub struct ID(pub ::std::os::raw::c_ulong);
 /// Pointer to an ISEQ
 pub type IseqPtr = *const rb_iseq_t;
 
+pub type CallDataPtr = *const rb_call_data;
+
+pub type CmePtr = *const rb_callable_method_entry_t;
+
+pub type CFuncPtr = *const rb_method_cfunc_t;
+
 // Given an ISEQ pointer, convert PC to insn_idx
 pub fn iseq_pc_to_insn_idx(iseq: IseqPtr, pc: *mut VALUE) -> Option<u16> {
     let pc_zero = unsafe { rb_iseq_pc_at_idx(iseq, 0) };
@@ -579,8 +585,8 @@ impl VALUE {
     }
 
     /// Assert that `self` is a method entry in debug builds
-    pub fn as_cme(self) -> *const rb_callable_method_entry_t {
-        let ptr: *const rb_callable_method_entry_t = self.as_ptr();
+    pub fn as_cme(self) -> CmePtr {
+        let ptr: CmePtr = self.as_ptr();
 
         #[cfg(debug_assertions)]
         if !ptr.is_null() {
@@ -611,9 +617,9 @@ impl From<IseqPtr> for VALUE {
     }
 }
 
-impl From<*const rb_callable_method_entry_t> for VALUE {
+impl From<CmePtr> for VALUE {
     /// For `.into()` convenience
-    fn from(cme: *const rb_callable_method_entry_t) -> Self {
+    fn from(cme: CmePtr) -> Self {
         VALUE(cme as usize)
     }
 }
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
diff --git a/zjit/src/cruby_methods.rs b/zjit/src/cruby_methods.rs
@@ -37,7 +37,7 @@ impl Annotations {
             if VM_METHOD_TYPE_CFUNC != get_cme_def_type(method) {
                 return None;
             }
-            get_mct_func(get_cme_def_body_cfunc(method.cast()))
+            rb_get_mct_func(rb_get_cme_def_body_cfunc(method.cast()))
         };
         self.cfuncs.get(&fn_ptr).copied()
     }
@@ -65,11 +65,12 @@ pub fn init() -> Annotations {
     let cfuncs = &mut HashMap::new();
 
     macro_rules! annotate {
-        ($module:ident, $method_name:literal, $return_type:expr, $($properties:ident),+) => {
+        ($module:ident, $method_name:literal, $return_type:expr, $($properties:ident),*) => {
+            #[allow(unused_mut)]
             let mut props = FnProperties { no_gc: false, leaf: false, elidable: false, return_type: $return_type };
             $(
                 props.$properties = true;
-            )+
+            )*
             annotate_c_method(cfuncs, unsafe { $module }, $method_name, props);
         }
     }
@@ -80,6 +81,7 @@ pub fn init() -> Annotations {
     annotate!(rb_cModule, "===", types::BoolExact, no_gc, leaf);
     annotate!(rb_cArray, "length", types::Fixnum, no_gc, leaf, elidable);
     annotate!(rb_cArray, "size", types::Fixnum, no_gc, leaf, elidable);
+    annotate!(rb_cInteger, "+", types::IntegerExact,);
 
     Annotations {
         cfuncs: std::mem::take(cfuncs)
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
