YJIT: `Kernel#{is_a?,instance_of?}` fast paths (GH-7297)

Co-authored-by: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

bootstraptest/test_yjit.rb

@@ -3552,3 +3552,30 @@ assert_equal 'threw', %q{
 
   bar([Hash.ruby2_keywords_hash({})])
 }
+
+# Test instance_of? and is_a?
+assert_equal 'true', %q{
+  1.instance_of?(Integer) && 1.is_a?(Integer)
+}
+
+# Test instance_of? and is_a? for singleton classes
+assert_equal 'true', %q{
+  a = []
+  def a.test = :test
+  a.instance_of?(Array) && a.is_a?(Array)
+}
+
+# Test instance_of? for singleton_class
+# Yes this does really return false
+assert_equal 'false', %q{
+  a = []
+  def a.test = :test
+  a.instance_of?(a.singleton_class)
+}
+
+# Test is_a? for singleton_class
+assert_equal 'true', %q{
+  a = []
+  def a.test = :test
+  a.is_a?(a.singleton_class)
+}

yjit/bindgen/src/main.rs

@@ -419,6 +419,7 @@ fn main() {
         .allowlist_function("rb_RCLASS_ORIGIN")
         .allowlist_function("rb_method_basic_definition_p")
         .allowlist_function("rb_yjit_array_len")
+        .allowlist_function("rb_obj_class")
 
         // We define VALUE manually, don't import it
         .blocklist_type("VALUE")

yjit/src/codegen.rs

@@ -4020,6 +4020,122 @@ fn jit_rb_false(
     true
 }
 
+/// Codegen for Kernel#is_a?
+fn jit_rb_kernel_is_a(
+    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 {
+    if argc != 1 {
+        return false;
+    }
+
+    // If this is a super call we might not know the class
+    if known_recv_class.is_null() {
+        return false;
+    }
+
+    // Important note: The output code will simply `return true/false`.
+    // Correctness follows from:
+    //  - `known_recv_class` implies there is a guard scheduled before here
+    //    for a particular `CLASS_OF(lhs)`.
+    //  - We guard that rhs is identical to the compile-time sample
+    //  - In general, for any two Class instances A, B, `A < B` does not change at runtime.
+    //    Class#superclass is stable.
+
+    let sample_rhs = jit_peek_at_stack(jit, ctx, 0);
+    let sample_lhs = jit_peek_at_stack(jit, ctx, 1);
+
+    // We are not allowing module here because the module hierachy can change at runtime.
+    if !unsafe { RB_TYPE_P(sample_rhs, RUBY_T_CLASS) } {
+        return false;
+    }
+    let sample_is_a = unsafe { rb_obj_is_kind_of(sample_lhs, sample_rhs) == Qtrue };
+
+    let side_exit = get_side_exit(jit, ocb, ctx);
+    asm.comment("Kernel#is_a?");
+    asm.cmp(ctx.stack_opnd(0), sample_rhs.into());
+    asm.jne(counted_exit!(ocb, side_exit, send_is_a_class_mismatch));
+
+    ctx.stack_pop(2);
+
+    if sample_is_a {
+        let stack_ret = ctx.stack_push(Type::True);
+        asm.mov(stack_ret, Qtrue.into());
+    } else {
+        let stack_ret = ctx.stack_push(Type::False);
+        asm.mov(stack_ret, Qfalse.into());
+    }
+    return true;
+}
+
+/// Codegen for Kernel#instance_of?
+fn jit_rb_kernel_instance_of(
+    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 {
+    if argc != 1 {
+        return false;
+    }
+
+    // If this is a super call we might not know the class
+    if known_recv_class.is_null() {
+        return false;
+    }
+
+    // Important note: The output code will simply `return true/false`.
+    // Correctness follows from:
+    //  - `known_recv_class` implies there is a guard scheduled before here
+    //    for a particular `CLASS_OF(lhs)`.
+    //  - We guard that rhs is identical to the compile-time sample
+    //  - For a particular `CLASS_OF(lhs)`, `rb_obj_class(lhs)` does not change.
+    //    (because for any singleton class `s`, `s.superclass.equal?(s.attached_object.class)`)
+
+    let sample_rhs = jit_peek_at_stack(jit, ctx, 0);
+    let sample_lhs = jit_peek_at_stack(jit, ctx, 1);
+
+    // Filters out cases where the C implementation raises
+    if unsafe { !(RB_TYPE_P(sample_rhs, RUBY_T_CLASS) || RB_TYPE_P(sample_rhs, RUBY_T_MODULE)) } {
+        return false;
+    }
+
+    // We need to grab the class here to deal with singleton classes.
+    // Instance of grabs the "real class" of the object rather than the
+    // singleton class.
+    let sample_lhs_real_class = unsafe { rb_obj_class(sample_lhs) };
+
+    let sample_instance_of = sample_lhs_real_class == sample_rhs;
+
+    let side_exit = get_side_exit(jit, ocb, ctx);
+    asm.comment("Kernel#instance_of?");
+    asm.cmp(ctx.stack_opnd(0), sample_rhs.into());
+    asm.jne(counted_exit!(ocb, side_exit, send_instance_of_class_mismatch));
+
+    ctx.stack_pop(2);
+
+    if sample_instance_of {
+        let stack_ret = ctx.stack_push(Type::True);
+        asm.mov(stack_ret, Qtrue.into());
+    } else {
+        let stack_ret = ctx.stack_push(Type::False);
+        asm.mov(stack_ret, Qfalse.into());
+    }
+    return true;
+}
+
 // Codegen for rb_obj_equal()
 // object identity comparison
 fn jit_rb_obj_equal(
@@ -7727,6 +7843,9 @@ impl CodegenGlobals {
 
             self.yjit_reg_method(rb_cNilClass, "nil?", jit_rb_true);
             self.yjit_reg_method(rb_mKernel, "nil?", jit_rb_false);
+            self.yjit_reg_method(rb_mKernel, "is_a?", jit_rb_kernel_is_a);
+            self.yjit_reg_method(rb_mKernel, "kind_of?", jit_rb_kernel_is_a);
+            self.yjit_reg_method(rb_mKernel, "instance_of?", jit_rb_kernel_instance_of);
 
             self.yjit_reg_method(rb_cBasicObject, "==", jit_rb_obj_equal);
             self.yjit_reg_method(rb_cBasicObject, "equal?", jit_rb_obj_equal);

yjit/src/cruby_bindings.inc.rs

@@ -1090,6 +1090,7 @@ extern "C" {
     pub static mut rb_cSymbol: VALUE;
     pub static mut rb_cThread: VALUE;
     pub static mut rb_cTrueClass: VALUE;
+    pub fn rb_obj_class(obj: VALUE) -> VALUE;
     pub fn rb_ary_new_capa(capa: ::std::os::raw::c_long) -> VALUE;
     pub fn rb_ary_store(ary: VALUE, key: ::std::os::raw::c_long, val: VALUE);
     pub fn rb_ary_resurrect(ary: VALUE) -> VALUE;

yjit/src/stats.rs

@@ -237,6 +237,8 @@ make_counters! {
     send_send_chain_not_string_or_sym,
     send_send_getter,
     send_send_builtin,
+    send_is_a_class_mismatch,
+    send_instance_of_class_mismatch,
 
     send_bmethod_ractor,
     send_bmethod_block_arg,