Implement optimized call

3 files changed, 85 insertions, 8 deletions

diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index bb81fc87a3..ecdd00a617 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -2694,7 +2694,7 @@ module RubyVM::MJIT
       end
     end
 
-    # Note: This clobbers :rax
+    # NOTE: This clobbers :rax
     # @param jit [RubyVM::MJIT::JITState]
     # @param ctx [RubyVM::MJIT::Context]
     # @param asm [RubyVM::MJIT::Assembler]
@@ -3377,8 +3377,7 @@ module RubyVM::MJIT
       when C.OPTIMIZED_METHOD_TYPE_SEND
         jit_call_opt_send(jit, ctx, asm, cme, flags, argc, block_handler, known_recv_class, send_shift:)
       when C.OPTIMIZED_METHOD_TYPE_CALL
-        asm.incr_counter(:send_optimized_call)
-        return CantCompile
+        jit_call_opt_call(jit, ctx, asm, cme, flags, argc, block_handler, known_recv_class, send_shift:)
       when C.OPTIMIZED_METHOD_TYPE_BLOCK_CALL
         asm.incr_counter(:send_optimized_block_call)
         return CantCompile
@@ -3422,6 +3421,64 @@ module RubyVM::MJIT
       jit_call_symbol(jit, ctx, asm, cme, C.VM_CALL_FCALL, argc, kw_splat, block_handler, known_recv_class, send_shift:)
     end
 
+    # vm_call_opt_call
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
+    def jit_call_opt_call(jit, ctx, asm, cme, flags, argc, block_handler, known_recv_class, send_shift:)
+      if block_handler != C.VM_BLOCK_HANDLER_NONE
+        asm.incr_counter(:send_optimized_call_block)
+        return CantCompile
+      end
+
+      if flags & C.VM_CALL_KWARG != 0
+        asm.incr_counter(:send_optimized_call_kwarg)
+        return CantCompile
+      end
+
+      if flags & C.VM_CALL_ARGS_SPLAT != 0
+        asm.incr_counter(:send_optimized_call_splat)
+        return CantCompile
+      end
+
+      # TODO: implement this
+      # Optimize for single ractor mode and avoid runtime check for
+      # "defined with an un-shareable Proc in a different Ractor"
+      # if !assume_single_ractor_mode(jit, ocb)
+      #   return CantCompile
+      # end
+
+      # If this is a .send call we need to adjust the stack
+      if flags & C.VM_CALL_OPT_SEND != 0
+        jit_call_opt_send_shift_stack(ctx, asm, argc, send_shift:)
+      end
+
+      # About to reset the SP, need to load this here
+      recv_idx = argc # blockarg is not supported. send_shift is already handled.
+      asm.mov(:rcx, ctx.stack_opnd(recv_idx)) # recv
+
+      # Save the PC and SP because the callee can make Ruby calls
+      jit_prepare_routine_call(jit, ctx, asm) # NOTE: clobbers rax
+
+      asm.lea(:rax, ctx.sp_opnd(0)) # sp
+
+      kw_splat = flags & C.VM_CALL_KW_SPLAT
+
+      asm.mov(C_ARGS[0], :rcx)
+      asm.mov(C_ARGS[1], EC)
+      asm.mov(C_ARGS[2], argc)
+      asm.lea(C_ARGS[3], [:rax, -argc * C.VALUE.size]) # stack_argument_pointer. NOTE: C_ARGS[3] is rcx
+      asm.mov(C_ARGS[4], kw_splat)
+      asm.mov(C_ARGS[5], C.VM_BLOCK_HANDLER_NONE)
+      asm.call(C.rb_optimized_call)
+
+      ctx.stack_pop(argc + 1)
+
+      stack_ret = ctx.stack_push
+      asm.mov(stack_ret, C_RET)
+      return KeepCompiling
+    end
+
     # @param ctx [RubyVM::MJIT::Context]
     # @param asm [RubyVM::MJIT::Assembler]
     def jit_call_opt_send_shift_stack(ctx, asm, argc, send_shift:)
@@ -3635,9 +3692,8 @@ module RubyVM::MJIT
           # Qnil push is handled in jit_push_frame
 
           return iseq.body.param.opt_table[opt]
-        else
-          # We don't support the remaining `else if`s yet.
-          asm.incr_counter(:send_iseq_not_only_optparam)
+        elsif C.rb_iseq_only_kwparam_p(iseq) && (flags & C.VM_CALL_ARGS_SPLAT) == 0
+          asm.incr_counter(:send_iseq_kwparam)
           return CantCompile
         end
       end
diff --git a/mjit_c.h b/mjit_c.h
index 40336ca0ca..2e233d41de 100644
--- a/mjit_c.h
+++ b/mjit_c.h
@@ -140,7 +140,7 @@ MJIT_RUNTIME_COUNTERS(
     send_block_not_nil,
     send_block_not_proxy,
 
-    send_iseq_not_only_optparam,
+    send_iseq_kwparam,
     send_iseq_kw_splat,
 
     send_cfunc_variadic,
@@ -164,6 +164,9 @@ MJIT_RUNTIME_COUNTERS(
     send_optimized_send_mid_id_changed,
     send_optimized_send_null_mid,
     send_optimized_send_send,
+    send_optimized_call_block,
+    send_optimized_call_kwarg,
+    send_optimized_call_splat,
     send_optimized_blockarg,
 
     invokesuper_me_changed,
diff --git a/mjit_c.rb b/mjit_c.rb
index 6dca2e2441..9bdf2831f1 100644
--- a/mjit_c.rb
+++ b/mjit_c.rb
@@ -272,6 +272,14 @@ module RubyVM::MJIT # :nodoc: all
       }
     end
 
+    def rb_iseq_only_kwparam_p(iseq)
+      _iseq = iseq.to_i
+      Primitive.cstmt! %{
+        extern bool rb_iseq_only_kwparam_p(const rb_iseq_t *iseq);
+        return RBOOL(rb_iseq_only_kwparam_p((rb_iseq_t *)NUM2SIZET(_iseq)));
+      }
+    end
+
     def rb_vm_opt_newarray_min
       Primitive.cstmt! %{
         extern VALUE rb_vm_opt_newarray_min(rb_execution_context_t *ec, rb_num_t num, const VALUE *ptr);
@@ -388,6 +396,13 @@ module RubyVM::MJIT # :nodoc: all
       Primitive.cexpr! 'rb_singleton_class(obj)'
     end
 
+    def rb_optimized_call
+      Primitive.cstmt! %{
+        extern VALUE rb_optimized_call(VALUE *recv, rb_execution_context_t *ec, int argc, VALUE *argv, int kw_splat, VALUE block_handler);
+        return SIZET2NUM((size_t)rb_optimized_call);
+      }
+    end
+
     #========================================================================================
     #
     # Old stuff
@@ -1448,7 +1463,7 @@ module RubyVM::MJIT # :nodoc: all
       send_block_setup: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_block_setup)")],
       send_block_not_nil: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_block_not_nil)")],
       send_block_not_proxy: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_block_not_proxy)")],
-      send_iseq_not_only_optparam: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_iseq_not_only_optparam)")],
+      send_iseq_kwparam: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_iseq_kwparam)")],
       send_iseq_kw_splat: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_iseq_kw_splat)")],
       send_cfunc_variadic: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_cfunc_variadic)")],
       send_cfunc_too_many_args: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_cfunc_too_many_args)")],
@@ -1468,6 +1483,9 @@ module RubyVM::MJIT # :nodoc: all
       send_optimized_send_mid_id_changed: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_send_mid_id_changed)")],
       send_optimized_send_null_mid: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_send_null_mid)")],
       send_optimized_send_send: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_send_send)")],
+      send_optimized_call_block: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_call_block)")],
+      send_optimized_call_kwarg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_call_kwarg)")],
+      send_optimized_call_splat: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_call_splat)")],
       send_optimized_blockarg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_optimized_blockarg)")],
       invokesuper_me_changed: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), invokesuper_me_changed)")],
       invokesuper_same_me: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), invokesuper_same_me)")],