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>

bootstraptest/test_yjit.rb

@@ -3255,3 +3255,57 @@ assert_equal '[1, 2]', %q{
   foo
   foo
 }
+
+# respond_to? with changing symbol
+assert_equal 'false', %q{
+  def foo(name)
+    :sym.respond_to?(name)
+  end
+  foo(:to_s)
+  foo(:to_s)
+  foo(:not_exist)
+}
+
+# respond_to? with method being defined
+assert_equal 'true', %q{
+  def foo
+    :sym.respond_to?(:not_yet_defined)
+  end
+  foo
+  foo
+  module Kernel
+    def not_yet_defined = true
+  end
+  foo
+}
+
+# respond_to? with undef method
+assert_equal 'false', %q{
+  module Kernel
+    def to_be_removed = true
+  end
+  def foo
+    :sym.respond_to?(:to_be_removed)
+  end
+  foo
+  foo
+  class Object
+    undef_method :to_be_removed
+  end
+  foo
+}
+
+# respond_to? with respond_to_missing?
+assert_equal 'true', %q{
+  class Foo
+  end
+  def foo(x)
+    x.respond_to?(:bar)
+  end
+  foo(Foo.new)
+  foo(Foo.new)
+  class Foo
+    def respond_to_missing?(*) = true
+  end
+  foo(Foo.new)
+}

method.h

@@ -226,6 +226,7 @@ const rb_method_entry_t *rb_resolve_me_location(const rb_method_entry_t *, VALUE
 RUBY_SYMBOL_EXPORT_END
 
 const rb_callable_method_entry_t *rb_callable_method_entry(VALUE klass, ID id);
+const rb_callable_method_entry_t *rb_callable_method_entry_or_negative(VALUE klass, ID id);
 const rb_callable_method_entry_t *rb_callable_method_entry_with_refinements(VALUE klass, ID id, VALUE *defined_class);
 const rb_callable_method_entry_t *rb_callable_method_entry_without_refinements(VALUE klass, ID id, VALUE *defined_class);
 

vm_method.c

@@ -1347,7 +1347,7 @@ negative_cme(ID mid)
 }
 
 static const rb_callable_method_entry_t *
-callable_method_entry(VALUE klass, ID mid, VALUE *defined_class_ptr)
+callable_method_entry_or_negative(VALUE klass, ID mid, VALUE *defined_class_ptr)
 {
     const rb_callable_method_entry_t *cme;
 
@@ -1376,6 +1376,20 @@ callable_method_entry(VALUE klass, ID mid, VALUE *defined_class_ptr)
     }
     RB_VM_LOCK_LEAVE();
 
+    return cme;
+}
+
+// This is exposed for YJIT so that we can make assumptions that methods are
+// not defined.
+const rb_callable_method_entry_t *
+rb_callable_method_entry_or_negative(VALUE klass, ID mid) {
+    return callable_method_entry_or_negative(klass, mid, NULL);
+}
+
+static const rb_callable_method_entry_t *
+callable_method_entry(VALUE klass, ID mid, VALUE *defined_class_ptr) {
+    const rb_callable_method_entry_t *cme;
+    cme = callable_method_entry_or_negative(klass, mid, defined_class_ptr);
     return !UNDEFINED_METHOD_ENTRY_P(cme) ? cme : NULL;
 }
 

yjit.c

@@ -486,7 +486,11 @@ rb_METHOD_ENTRY_VISI(const rb_callable_method_entry_t *me)
 rb_method_type_t
 rb_get_cme_def_type(const rb_callable_method_entry_t *cme)
 {
+    if (UNDEFINED_METHOD_ENTRY_P(cme)) {
+        return VM_METHOD_TYPE_UNDEF;
+    } else {
         return cme->def->type;
+    }
 }
 
 ID

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")

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),

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 {

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 {