YJIT: Implement specialized respond_to? (#6363)

* Add rb_callable_method_entry_or_negative

* YJIT: Implement specialized respond_to?

This implements a specialized respond_to? in YJIT.

* Update yjit/src/codegen.rs

Co-authored-by: Maxime Chevalier-Boisvert <maximechevalierb@gmail.com>

4 files changed, 128 insertions, 0 deletions

diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index 294da21378..c3d4a39a2b 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -225,6 +225,7 @@ fn main() {
         .allowlist_var(".*_REDEFINED_OP_FLAG")
         .allowlist_type("rb_num_t")
         .allowlist_function("rb_callable_method_entry")
+        .allowlist_function("rb_callable_method_entry_or_negative")
         .allowlist_function("rb_vm_frame_method_entry")
         .allowlist_type("IVC") // pointer to iseq_inline_iv_cache_entry
         .allowlist_type("IC")  // pointer to iseq_inline_constant_cache
@@ -367,6 +368,7 @@ fn main() {
         .allowlist_function("rb_vm_ci_kwarg")
         .allowlist_function("rb_METHOD_ENTRY_VISI")
         .allowlist_function("rb_RCLASS_ORIGIN")
+        .allowlist_function("rb_method_basic_definition_p")
 
         // We define VALUE manually, don't import it
         .blocklist_type("VALUE")
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index f13505365e..11f7085635 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -3841,6 +3841,102 @@ fn jit_rb_str_concat(
     true
 }
 
+fn jit_obj_respond_to(
+    jit: &mut JITState,
+    ctx: &mut Context,
+    asm: &mut Assembler,
+    ocb: &mut OutlinedCb,
+    _ci: *const rb_callinfo,
+    _cme: *const rb_callable_method_entry_t,
+    _block: Option<IseqPtr>,
+    argc: i32,
+    known_recv_class: *const VALUE,
+) -> bool {
+    // respond_to(:sym) or respond_to(:sym, true)
+    if argc != 1 && argc != 2 {
+        return false;
+    }
+
+    if known_recv_class.is_null() {
+        return false;
+    }
+
+    let recv_class = unsafe { *known_recv_class };
+
+    // Get the method_id from compile time. We will later add a guard against it.
+    let mid_sym = jit_peek_at_stack(jit, ctx, (argc - 1) as isize);
+    if !mid_sym.static_sym_p() {
+        return false
+    }
+    let mid = unsafe { rb_sym2id(mid_sym) };
+
+    // Option<bool> representing the value of the "include_all" argument and whether it's known
+    let allow_priv = if argc == 1 {
+        // Default is false
+        Some(false)
+    } else {
+        // Get value from type information (may or may not be known)
+        ctx.get_opnd_type(StackOpnd(0)).known_truthy()
+    };
+
+    let mut target_cme = unsafe { rb_callable_method_entry_or_negative(recv_class, mid) };
+
+    // Should never be null, as in that case we will be returned a "negative CME"
+    assert!(!target_cme.is_null());
+
+    let cme_def_type = unsafe { get_cme_def_type(target_cme) };
+
+    if cme_def_type == VM_METHOD_TYPE_REFINED {
+        return false;
+    }
+
+    let visibility = if cme_def_type == VM_METHOD_TYPE_UNDEF {
+        METHOD_VISI_UNDEF
+    } else {
+        unsafe { METHOD_ENTRY_VISI(target_cme) }
+    };
+
+    let result = match (visibility, allow_priv) {
+        (METHOD_VISI_UNDEF, _) => Qfalse, // No method => false
+        (METHOD_VISI_PUBLIC, _) => Qtrue, // Public method => true regardless of include_all
+        (_, Some(true)) => Qtrue, // include_all => always true
+        (_, _) => return false // not public and include_all not known, can't compile
+    };
+
+    if result != Qtrue {
+        // Only if respond_to_missing? hasn't been overridden
+        // In the future, we might want to jit the call to respond_to_missing?
+        if !assume_method_basic_definition(jit, ocb, recv_class, idRespond_to_missing.into()) {
+            return false;
+        }
+    }
+
+    // Invalidate this block if method lookup changes for the method being queried. This works
+    // both for the case where a method does or does not exist, as for the latter we asked for a
+    // "negative CME" earlier.
+    assume_method_lookup_stable(jit, ocb, recv_class, target_cme);
+
+    // Generate a side exit
+    let side_exit = get_side_exit(jit, ocb, ctx);
+
+    if argc == 2 {
+        // pop include_all argument (we only use its type info)
+        ctx.stack_pop(1);
+    }
+
+    let sym_opnd = ctx.stack_pop(1);
+    let recv_opnd = ctx.stack_pop(1);
+
+    // This is necessary because we have no guarantee that sym_opnd is a constant
+    asm.comment("guard known mid");
+    asm.cmp(sym_opnd, mid_sym.into());
+    asm.jne(side_exit.into());
+
+    jit_putobject(jit, ctx, asm, result);
+
+    true
+}
+
 fn jit_thread_s_current(
     _jit: &mut JITState,
     ctx: &mut Context,
@@ -6292,6 +6388,8 @@ impl CodegenGlobals {
             self.yjit_reg_method(rb_cString, "<<", jit_rb_str_concat);
             self.yjit_reg_method(rb_cString, "+@", jit_rb_str_uplus);
 
+            self.yjit_reg_method(rb_mKernel, "respond_to?", jit_obj_respond_to);
+
             // Thread.current
             self.yjit_reg_method(
                 rb_singleton_class(rb_cThread),
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index e3dbdc0d4b..b391a6cda5 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -26,6 +26,9 @@ pub type rb_alloc_func_t = ::std::option::Option<unsafe extern "C" fn(klass: VAL
 extern "C" {
     pub fn rb_get_alloc_func(klass: VALUE) -> rb_alloc_func_t;
 }
+extern "C" {
+    pub fn rb_method_basic_definition_p(klass: VALUE, mid: ID) -> ::std::os::raw::c_int;
+}
 #[repr(C)]
 pub struct RBasic {
     pub flags: VALUE,
@@ -575,6 +578,12 @@ extern "C" {
 extern "C" {
     pub fn rb_callable_method_entry(klass: VALUE, id: ID) -> *const rb_callable_method_entry_t;
 }
+extern "C" {
+    pub fn rb_callable_method_entry_or_negative(
+        klass: VALUE,
+        id: ID,
+    ) -> *const rb_callable_method_entry_t;
+}
 pub type rb_num_t = ::std::os::raw::c_ulong;
 #[repr(C)]
 pub struct iseq_inline_constant_cache_entry {
diff --git a/yjit/src/invariants.rs b/yjit/src/invariants.rs
index 4ed21118cc..ee79b2938a 100644
--- a/yjit/src/invariants.rs
+++ b/yjit/src/invariants.rs
@@ -154,6 +154,25 @@ pub fn assume_method_lookup_stable(
         .insert(block);
 }
 
+// Checks rb_method_basic_definition_p and registers the current block for invalidation if method
+// lookup changes.
+// A "basic method" is one defined during VM boot, so we can use this to check assumptions based on
+// default behavior.
+pub fn assume_method_basic_definition(
+    jit: &mut JITState,
+    ocb: &mut OutlinedCb,
+    klass: VALUE,
+    mid: ID
+    ) -> bool {
+    if unsafe { rb_method_basic_definition_p(klass, mid) } != 0 {
+        let mut cme = unsafe { rb_callable_method_entry(klass, mid) };
+        assume_method_lookup_stable(jit, ocb, klass, cme);
+        true
+    } else {
+        false
+    }
+}
+
 /// Tracks that a block is assuming it is operating in single-ractor mode.
 #[must_use]
 pub fn assume_single_ractor_mode(jit: &mut JITState, ocb: &mut OutlinedCb) -> bool {