From e4d3e652ff41943c7f50d4552c1c98dfaa736cf9 Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Mon, 12 Feb 2024 09:56:44 -0800
Subject: [PATCH] YJIT: Prefer an overloaded cme if available (#9913)

YJIT: Prefer an overloaded cme if applicable
---
 vm_insnhelper.c     | 4 ++--
 vm_method.c         | 4 ++--
 yjit/src/codegen.rs | 4 ++++
 yjit/src/cruby.rs   | 4 ++++
 4 files changed, 12 insertions(+), 4 deletions(-)

diff --git a/vm_insnhelper.c b/vm_insnhelper.c
index a74218427868eb..aef71970f69cdb 100644
--- a/vm_insnhelper.c
+++ b/vm_insnhelper.c
@@ -2032,7 +2032,7 @@ vm_ccs_verify(struct rb_class_cc_entries *ccs, ID mid, VALUE klass)
 }
 #endif
 
-static const rb_callable_method_entry_t *check_overloaded_cme(const rb_callable_method_entry_t *cme, const struct rb_callinfo * const ci);
+const rb_callable_method_entry_t *rb_check_overloaded_cme(const rb_callable_method_entry_t *cme, const struct rb_callinfo * const ci);
 
 static const struct rb_callcache *
 vm_search_cc(const VALUE klass, const struct rb_callinfo * const ci)
@@ -2118,7 +2118,7 @@ vm_search_cc(const VALUE klass, const struct rb_callinfo * const ci)
         }
     }
 
-    cme = check_overloaded_cme(cme, ci);
+    cme = rb_check_overloaded_cme(cme, ci);
 
     const struct rb_callcache *cc = vm_cc_new(klass, cme, vm_call_general, cc_type_normal);
     vm_ccs_push(klass, ccs, ci, cc);
diff --git a/vm_method.c b/vm_method.c
index a4d0a9f4112e9d..05c99b9905db64 100644
--- a/vm_method.c
+++ b/vm_method.c
@@ -1063,8 +1063,8 @@ get_overloaded_cme(const rb_callable_method_entry_t *cme)
     }
 }
 
-static const rb_callable_method_entry_t *
-check_overloaded_cme(const rb_callable_method_entry_t *cme, const struct rb_callinfo * const ci)
+const rb_callable_method_entry_t *
+rb_check_overloaded_cme(const rb_callable_method_entry_t *cme, const struct rb_callinfo * const ci)
 {
     if (UNLIKELY(cme->def->iseq_overload) &&
         (vm_ci_flag(ci) & (VM_CALL_ARGS_SIMPLE)) &&
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index ac8670fe407a5c..421fa989e23e22 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -7679,6 +7679,10 @@ fn gen_send_general(
         return None;
     }
 
+    // Load an overloaded cme if applicable. See vm_search_cc().
+    // It allows you to use a faster ISEQ if possible.
+    cme = unsafe { rb_check_overloaded_cme(cme, ci) };
+
     let visi = unsafe { METHOD_ENTRY_VISI(cme) };
     match visi {
         METHOD_VISI_PUBLIC => {
diff --git a/yjit/src/cruby.rs b/yjit/src/cruby.rs
index 78cb2d65ea91c3..fb98e3474bb166 100644
--- a/yjit/src/cruby.rs
+++ b/yjit/src/cruby.rs
@@ -112,6 +112,10 @@ pub use autogened::*;
 // Use bindgen for functions that are defined in headers or in yjit.c.
 #[cfg_attr(test, allow(unused))] // We don't link against C code when testing
 extern "C" {
+    pub fn rb_check_overloaded_cme(
+        me: *const rb_callable_method_entry_t,
+        ci: *const rb_callinfo,
+    ) -> *const rb_callable_method_entry_t;
     pub fn rb_hash_empty_p(hash: VALUE) -> VALUE;
     pub fn rb_vm_splat_array(flag: VALUE, ary: VALUE) -> VALUE;
     pub fn rb_vm_concat_array(ary1: VALUE, ary2st: VALUE) -> VALUE;