RJIT: Optimize Kernel#respond_to?

5 files changed, 128 insertions, 1 deletions

diff --git a/lib/ruby_vm/rjit/insn_compiler.rb b/lib/ruby_vm/rjit/insn_compiler.rb
index 58028d51ab..b05578b6c8 100644
--- a/lib/ruby_vm/rjit/insn_compiler.rb
+++ b/lib/ruby_vm/rjit/insn_compiler.rb
@@ -2947,6 +2947,95 @@ module RubyVM::RJIT
     # @param jit [RubyVM::RJIT::JITState]
     # @param ctx [RubyVM::RJIT::Context]
     # @param asm [RubyVM::RJIT::Assembler]
+    def jit_obj_respond_to(jit, ctx, asm, argc, known_recv_class)
+      # respond_to(:sym) or respond_to(:sym, true)
+      if argc != 1 && argc != 2
+        return false
+      end
+
+      if known_recv_class.nil?
+        return false
+      end
+
+      recv_class = known_recv_class
+
+      # Get the method_id from compile time. We will later add a guard against it.
+      mid_sym = jit.peek_at_stack(argc - 1)
+      unless static_symbol?(mid_sym)
+        return false
+      end
+      mid = C.rb_sym2id(mid_sym)
+
+      target_cme = C.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_equal(false, target_cme.nil?)
+
+      cme_def_type =
+        if C.UNDEFINED_METHOD_ENTRY_P(target_cme)
+          C::VM_METHOD_TYPE_UNDEF
+        else
+          target_cme.def.type
+        end
+
+      if cme_def_type == C::VM_METHOD_TYPE_REFINED
+        return false
+      end
+
+      visibility = if cme_def_type == C::VM_METHOD_TYPE_UNDEF
+        C::METHOD_VISI_UNDEF
+      else
+        C.METHOD_ENTRY_VISI(target_cme)
+      end
+
+      result =
+        case visibility
+        in C::METHOD_VISI_UNDEF
+          Qfalse # No method => false
+        in C::METHOD_VISI_PUBLIC
+          Qtrue # Public method => true regardless of include_all
+        else
+          return false # not public and include_all not known, can't compile
+        end
+
+      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?
+        unless Invariants.assume_method_basic_definition(jit, recv_class, C.idRespond_to_missing)
+          return false
+        end
+      end
+
+      # 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.
+      Invariants.assume_method_lookup_stable(jit, target_cme)
+
+      # Generate a side exit
+      side_exit = side_exit(jit, ctx)
+
+      if argc == 2
+        # pop include_all argument (we only use its type info)
+        ctx.stack_pop(1)
+      end
+
+      sym_opnd = ctx.stack_pop(1)
+      _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.mov(:rax, to_value(mid_sym))
+      asm.cmp(sym_opnd, :rax)
+      asm.jne(side_exit)
+
+      putobject(jit, ctx, asm, val: result)
+
+      true
+    end
+
+    # @param jit [RubyVM::RJIT::JITState]
+    # @param ctx [RubyVM::RJIT::Context]
+    # @param asm [RubyVM::RJIT::Assembler]
     def jit_thread_s_current(jit, ctx, asm, argc, _known_recv_class)
       return false if argc != 0
       asm.comment('Thread.current')
@@ -2999,7 +3088,7 @@ module RubyVM::RJIT
       # rb_ary_empty_p() method in array.c
       register_cfunc_method(Array, :empty?, :jit_rb_ary_empty_p)
 
-      #register_cfunc_method(Kernel, :respond_to?, :jit_obj_respond_to)
+      register_cfunc_method(Kernel, :respond_to?, :jit_obj_respond_to)
       #register_cfunc_method(Kernel, :block_given?, :jit_rb_f_block_given_p)
 
       # Thread.current
diff --git a/lib/ruby_vm/rjit/invariants.rb b/lib/ruby_vm/rjit/invariants.rb
index 8a3fcb4874..db1068b5c2 100644
--- a/lib/ruby_vm/rjit/invariants.rb
+++ b/lib/ruby_vm/rjit/invariants.rb
@@ -38,6 +38,17 @@ module RubyVM::RJIT
         @cme_blocks[cme.to_i] << jit.block
       end
 
+      # @param jit [RubyVM::RJIT::JITState]
+      def assume_method_basic_definition(jit, klass, mid)
+        if C.rb_method_basic_definition_p(klass, mid)
+          cme = C.rb_callable_method_entry(klass, mid)
+          assume_method_lookup_stable(jit, cme)
+          true
+        else
+          false
+        end
+      end
+
       def assume_stable_constant_names(jit, idlist)
         (0..).each do |i|
           break if (id = idlist[i]) == 0
diff --git a/rjit_c.c b/rjit_c.c
index bf0bf6f410..a72ae4eec8 100644
--- a/rjit_c.c
+++ b/rjit_c.c
@@ -504,6 +504,7 @@ extern VALUE rb_vm_throw(const rb_execution_context_t *ec, rb_control_frame_t *r
 extern VALUE rb_reg_new_ary(VALUE ary, int opt);
 extern void rb_vm_setclassvariable(const rb_iseq_t *iseq, const rb_control_frame_t *cfp, ID id, VALUE val, ICVARC ic);
 extern VALUE rb_str_bytesize(VALUE str);
+extern const rb_callable_method_entry_t *rb_callable_method_entry_or_negative(VALUE klass, ID mid);
 
 #include "rjit_c.rbinc"
 
diff --git a/rjit_c.rb b/rjit_c.rb
index 0c6cfd3944..a9522b7454 100644
--- a/rjit_c.rb
+++ b/rjit_c.rb
@@ -310,6 +310,25 @@ module RubyVM::RJIT # :nodoc: all
     def rb_obj_class(obj)
       Primitive.cexpr! 'rb_obj_class(obj)'
     end
+
+    def rb_sym2id(sym)
+      Primitive.cexpr! 'SIZET2NUM((size_t)rb_sym2id(sym))'
+    end
+
+    def rb_callable_method_entry_or_negative(klass, mid)
+      cme_addr = Primitive.cexpr! 'SIZET2NUM((size_t)rb_callable_method_entry_or_negative(klass, (ID)NUM2SIZET(mid)))'
+      return nil if cme_addr == 0
+      rb_callable_method_entry_t.new(cme_addr)
+    end
+
+    def rb_method_basic_definition_p(klass, mid)
+      Primitive.cexpr! 'RBOOL(rb_method_basic_definition_p(klass, (ID)NUM2SIZET(mid)))'
+    end
+
+    def UNDEFINED_METHOD_ENTRY_P(cme)
+      _cme_addr = cme.to_i
+      Primitive.cexpr! 'RBOOL(UNDEFINED_METHOD_ENTRY_P((const rb_callable_method_entry_t *)NUM2SIZET(_cme_addr)))'
+    end
   end
 
   #
@@ -352,6 +371,7 @@ module RubyVM::RJIT # :nodoc: all
   C::METHOD_VISI_PRIVATE = Primitive.cexpr! %q{ SIZET2NUM(METHOD_VISI_PRIVATE) }
   C::METHOD_VISI_PROTECTED = Primitive.cexpr! %q{ SIZET2NUM(METHOD_VISI_PROTECTED) }
   C::METHOD_VISI_PUBLIC = Primitive.cexpr! %q{ SIZET2NUM(METHOD_VISI_PUBLIC) }
+  C::METHOD_VISI_UNDEF = Primitive.cexpr! %q{ SIZET2NUM(METHOD_VISI_UNDEF) }
   C::OBJ_TOO_COMPLEX_SHAPE_ID = Primitive.cexpr! %q{ SIZET2NUM(OBJ_TOO_COMPLEX_SHAPE_ID) }
   C::OPTIMIZED_METHOD_TYPE_BLOCK_CALL = Primitive.cexpr! %q{ SIZET2NUM(OPTIMIZED_METHOD_TYPE_BLOCK_CALL) }
   C::OPTIMIZED_METHOD_TYPE_CALL = Primitive.cexpr! %q{ SIZET2NUM(OPTIMIZED_METHOD_TYPE_CALL) }
@@ -433,6 +453,10 @@ module RubyVM::RJIT # :nodoc: all
     Primitive.cexpr! %q{ SIZET2NUM(block_type_iseq) }
   end
 
+  def C.idRespond_to_missing
+    Primitive.cexpr! %q{ SIZET2NUM(idRespond_to_missing) }
+  end
+
   def C.imemo_iseq
     Primitive.cexpr! %q{ SIZET2NUM(imemo_iseq) }
   end
diff --git a/tool/rjit/bindgen.rb b/tool/rjit/bindgen.rb
index bbf174d6fa..9916ebe0d0 100755
--- a/tool/rjit/bindgen.rb
+++ b/tool/rjit/bindgen.rb
@@ -401,6 +401,7 @@ generator = BindingGenerator.new(
       METHOD_VISI_PRIVATE
       METHOD_VISI_PROTECTED
       METHOD_VISI_PUBLIC
+      METHOD_VISI_UNDEF
       OBJ_TOO_COMPLEX_SHAPE_ID
       OPTIMIZED_METHOD_TYPE_BLOCK_CALL
       OPTIMIZED_METHOD_TYPE_CALL
@@ -492,6 +493,7 @@ generator = BindingGenerator.new(
       rb_cTrueClass
       rb_rjit_global_events
       rb_mRubyVMFrozenCore
+      idRespond_to_missing
     ],
   },
   funcs: %w[