From db4a8afa5e55591aa1e167ffde5924fc8f1f4807 Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Thu, 2 Mar 2023 22:23:47 -0800
Subject: [PATCH] Optimize BasicObject#!

---
 lib/ruby_vm/mjit/insn_compiler.rb | 52 +++++++++++++++++++++++++++++++
 mjit_c.rb                         | 20 ++++++++++++
 tool/mjit/bindgen.rb              |  1 +
 3 files changed, 73 insertions(+)

diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index e337eed106..6d241b7553 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -11,6 +11,8 @@ module RubyVM::MJIT
         @ocb.write(asm)
       end
 
+      @cfunc_codegen_table = {}
+      register_cfunc_codegen_funcs
       # freeze # workaround a binding.irb issue. TODO: resurrect this
     end
 
@@ -1675,10 +1677,49 @@ module RubyVM::MJIT
       putobject(jit, ctx, asm, val: C.to_value(1))
     end
 
+    #
+    # C func
+    #
+
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
+    def jit_rb_obj_not(jit, ctx, asm)
+      recv = ctx.stack_pop
+      # This `test` sets ZF only for Qnil and Qfalse, which let cmovz set.
+      asm.test(recv, ~Qnil)
+      asm.mov(:rax, Qfalse)
+      asm.mov(:rcx, Qtrue)
+      asm.cmovz(:rax, :rcx)
+      asm.mov(ctx.stack_push, :rax)
+
+      true
+    end
+
     #
     # Helpers
     #
 
+    def register_cfunc_codegen_funcs
+      register_cfunc_method(BasicObject, '!', :jit_rb_obj_not)
+    end
+
+    def register_cfunc_method(klass, mid_str, func)
+      mid = C.rb_intern(mid_str)
+      me = C.rb_method_entry_at(klass, mid)
+
+      assert_equal(false, me.nil?)
+
+      # Only cfuncs are supported
+      method_serial = me.def.method_serial
+
+      @cfunc_codegen_table[method_serial] = method(func)
+    end
+
+    def lookup_cfunc_codegen(cme_def)
+      @cfunc_codegen_table[cme_def.method_serial]
+    end
+
     def jit_getlocal_generic(jit, ctx, asm, idx:, level:)
       # Load environment pointer EP at level
       ep_reg = :rax
@@ -2545,6 +2586,17 @@ module RubyVM::MJIT
         return CantCompile
       end
 
+      # Delegate to codegen for C methods if we have it.
+      if flags & C.VM_CALL_KWARG == 0 && flags & C.VM_CALL_OPT_SEND == 0
+        known_cfunc_codegen = lookup_cfunc_codegen(cme.def)
+        if known_cfunc_codegen&.call(jit, ctx, asm)
+          # cfunc codegen generated code. Terminate the block so
+          # there isn't multiple calls in the same block.
+          jump_to_next_insn(jit, ctx, asm)
+          return EndBlock
+        end
+      end
+
       # We will not have side exits from here. Adjust the stack.
       if flags & C.VM_CALL_OPT_SEND != 0
         jit_call_opt_send_shift_stack(ctx, asm, argc, send_shift:)
diff --git a/mjit_c.rb b/mjit_c.rb
index bfdf6481a3..94ba211b78 100644
--- a/mjit_c.rb
+++ b/mjit_c.rb
@@ -269,6 +269,15 @@ module RubyVM::MJIT # :nodoc: all
       Primitive.cexpr! 'rb_obj_frozen_p(obj)'
     end
 
+    def rb_intern(str)
+      Primitive.cexpr! 'SIZET2NUM((size_t)rb_intern(RSTRING_PTR(str)))'
+    end
+
+    def rb_method_entry_at(klass, mid)
+      me_addr = Primitive.cexpr! 'SIZET2NUM((size_t)rb_method_entry_at(klass, (ID)NUM2SIZET(mid)))'
+      me_addr == 0 ? nil : rb_method_entry_t.new(me_addr)
+    end
+
     #========================================================================================
     #
     # Old stuff
@@ -1235,6 +1244,17 @@ module RubyVM::MJIT # :nodoc: all
     )
   end
 
+  def C.rb_method_entry_t
+    @rb_method_entry_t ||= CType::Struct.new(
+      "rb_method_entry_struct", Primitive.cexpr!("SIZEOF(struct rb_method_entry_struct)"),
+      flags: [self.VALUE, Primitive.cexpr!("OFFSETOF((*((struct rb_method_entry_struct *)NULL)), flags)")],
+      defined_class: [self.VALUE, Primitive.cexpr!("OFFSETOF((*((struct rb_method_entry_struct *)NULL)), defined_class)")],
+      def: [CType::Pointer.new { self.rb_method_definition_struct }, Primitive.cexpr!("OFFSETOF((*((struct rb_method_entry_struct *)NULL)), def)")],
+      called_id: [self.ID, Primitive.cexpr!("OFFSETOF((*((struct rb_method_entry_struct *)NULL)), called_id)")],
+      owner: [self.VALUE, Primitive.cexpr!("OFFSETOF((*((struct rb_method_entry_struct *)NULL)), owner)")],
+    )
+  end
+
   def C.rb_method_iseq_t
     @rb_method_iseq_t ||= CType::Struct.new(
       "rb_method_iseq_struct", Primitive.cexpr!("SIZEOF(struct rb_method_iseq_struct)"),
diff --git a/tool/mjit/bindgen.rb b/tool/mjit/bindgen.rb
index 571b94d9a2..105f878cdc 100755
--- a/tool/mjit/bindgen.rb
+++ b/tool/mjit/bindgen.rb
@@ -471,6 +471,7 @@ generator = BindingGenerator.new(
     rb_call_data
     rb_callable_method_entry_struct
     rb_callable_method_entry_t
+    rb_method_entry_t
     rb_callcache
     rb_callinfo
     rb_control_frame_t