Implement optimize call (#6691)

This dispatches to a c func for doing the dynamic lookup. I experimented with chain on the proc but wasn't able to detect which call sites would be monomorphic vs polymorphic. There is definitely room for optimization here, but it does reduce exits.

4 files changed, 61 insertions, 3 deletions

diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index 167ab2a74f..acbbaa613b 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -288,7 +288,7 @@ fn main() {
         .allowlist_function("rb_yjit_get_proc_ptr")
         .allowlist_function("rb_yjit_exit_locations_dict")
         .allowlist_function("rb_yjit_icache_invalidate")
-
+        .allowlist_function("rb_optimized_call")
         // from vm_sync.h
         .allowlist_function("rb_vm_barrier")
 
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 46f5ed64d3..7a1673c5cb 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -5557,8 +5557,53 @@ fn gen_send_general(
 
                     }
                     OPTIMIZED_METHOD_TYPE_CALL => {
-                        gen_counter_incr!(asm, send_optimized_method_call);
-                        return CantCompile;
+
+                        if block.is_some() {
+                            gen_counter_incr!(asm, send_call_block);
+                            return CantCompile;
+                        }
+
+                        if flags & VM_CALL_KWARG != 0 {
+                            gen_counter_incr!(asm, send_call_kwarg);
+                            return CantCompile;
+                        }
+
+                        // Optimize for single ractor mode and avoid runtime check for
+                        // "defined with an un-shareable Proc in a different Ractor"
+                        if !assume_single_ractor_mode(jit, ocb) {
+                            gen_counter_incr!(asm, send_call_multi_ractor);
+                            return CantCompile;
+                        }
+
+                        // About to reset the SP, need to load this here
+                        let recv_load = asm.load(recv);
+
+                        let sp = asm.lea(ctx.sp_opnd(0));
+
+                        // Write interpreter SP into CFP.
+                        // Needed in case the callee yields to the block.
+                        jit_save_pc(jit, asm);
+                        // Store incremented PC into current control frame in case callee raises.
+                        gen_save_sp(jit, asm, ctx);
+
+                        let kw_splat = flags & VM_CALL_KW_SPLAT;
+                        let stack_argument_pointer = asm.lea(Opnd::mem(64, sp, -(argc) * SIZEOF_VALUE_I32));
+
+                        let ret = asm.ccall(rb_optimized_call as *const u8, vec![
+                            recv_load,
+                            EC,
+                            argc.into(),
+                            stack_argument_pointer,
+                            kw_splat.into(),
+                            VM_BLOCK_HANDLER_NONE.into(),
+                        ]);
+
+                        ctx.stack_pop(argc as usize + 1);
+
+                        let stack_ret = ctx.stack_push(Type::Unknown);
+                        asm.mov(stack_ret, ret);
+                        return KeepCompiling;
+
                     }
                     OPTIMIZED_METHOD_TYPE_BLOCK_CALL => {
                         gen_counter_incr!(asm, send_optimized_method_block_call);
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index 3373c6d76e..9ede3030ff 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -1481,6 +1481,16 @@ extern "C" {
     pub fn rb_get_iseq_body_param_opt_table(iseq: *const rb_iseq_t) -> *const VALUE;
 }
 extern "C" {
+    pub fn rb_optimized_call(
+        recv: *mut VALUE,
+        ec: *mut rb_execution_context_t,
+        argc: ::std::os::raw::c_int,
+        argv: *mut VALUE,
+        kw_splat: ::std::os::raw::c_int,
+        block_handler: VALUE,
+    ) -> VALUE;
+}
+extern "C" {
     pub fn rb_leaf_invokebuiltin_iseq_p(iseq: *const rb_iseq_t) -> bool;
 }
 extern "C" {
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index a60fcaf836..b7bbb4ae3e 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -173,6 +173,9 @@ make_counters! {
     send_optimized_method,
     send_optimized_method_call,
     send_optimized_method_block_call,
+    send_call_block,
+    send_call_kwarg,
+    send_call_multi_ractor,
     send_missing_method,
     send_refined_method,
     send_cfunc_ruby_array_varg,