From 1a65ab20cb6519ab3d4e58141cfd812eaea5f7e0 Mon Sep 17 00:00:00 2001
From: Jimmy Miller <jimmy.miller@shopify.com>
Date: Tue, 8 Nov 2022 15:28:28 -0500
Subject: [PATCH] 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.
---
 yjit.c                         |  9 +++++++
 yjit/bindgen/src/main.rs       |  2 +-
 yjit/src/codegen.rs            | 49 ++++++++++++++++++++++++++++++++--
 yjit/src/cruby_bindings.inc.rs | 10 +++++++
 yjit/src/stats.rs              |  3 +++
 5 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/yjit.c b/yjit.c
index b943277d61..aa49b3cfdc 100644
--- a/yjit.c
+++ b/yjit.c
@@ -716,6 +716,15 @@ rb_get_iseq_body_param_opt_table(const rb_iseq_t *iseq)
     return iseq->body->param.opt_table;
 }
 
+VALUE
+rb_optimized_call(VALUE *recv, rb_execution_context_t *ec, int argc, VALUE *argv, int kw_splat, VALUE block_handler)
+{
+    rb_proc_t *proc;
+    GetProcPtr(recv, proc);
+    return rb_vm_invoke_proc(ec, proc, argc, argv, kw_splat, block_handler);
+}
+
+
 // If true, the iseq is leaf and it can be replaced by a single C call.
 bool
 rb_leaf_invokebuiltin_iseq_p(const rb_iseq_t *iseq)
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
@@ -1480,6 +1480,16 @@ extern "C" {
 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;
 }
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,