3 files changed, 56 insertions, 11 deletions
diff --git a/lib/ruby_vm/mjit/assembler.rb b/lib/ruby_vm/mjit/assembler.rb
index 6fd5d12bd8..ca663fa575 100644
--- a/lib/ruby_vm/mjit/assembler.rb
+++ b/lib/ruby_vm/mjit/assembler.rb
@@ -54,26 +54,36 @@ module RubyVM::MJIT
 
     def add(dst, src)
       case [dst, src]
-      # ADD r/m64, imm8 (Mod 11: reg)
-      in [Symbol => dst_reg, Integer => src_imm] if r64?(dst_reg) && imm8?(src_imm)
+      # ADD r/m64, imm8 (Mod 00: [reg])
+      in [[Symbol => dst_reg], Integer => src_imm] if r64?(dst_reg) && imm8?(src_imm)
         # REX.W + 83 /0 ib
         # MI: Operand 1: ModRM:r/m (r, w), Operand 2: imm8/16/32
         insn(
           prefix: REX_W,
           opcode: 0x83,
-          mod_rm: ModRM[mod: Mod11, reg: 0, rm: dst_reg],
+          mod_rm: ModRM[mod: Mod00, reg: 0, rm: dst_reg],
           imm: imm8(src_imm),
         )
-      # ADD r/m64, imm8 (Mod 00: [reg])
-      in [[Symbol => dst_reg], Integer => src_imm] if r64?(dst_reg) && imm8?(src_imm)
+      # ADD r/m64, imm8 (Mod 11: reg)
+      in [Symbol => dst_reg, Integer => src_imm] if r64?(dst_reg) && imm8?(src_imm)
         # REX.W + 83 /0 ib
         # MI: Operand 1: ModRM:r/m (r, w), Operand 2: imm8/16/32
         insn(
           prefix: REX_W,
           opcode: 0x83,
-          mod_rm: ModRM[mod: Mod00, reg: 0, rm: dst_reg],
+          mod_rm: ModRM[mod: Mod11, reg: 0, rm: dst_reg],
           imm: imm8(src_imm),
         )
+      # ADD r/m64 imm32 (Mod 11: reg)
+      in [Symbol => dst_reg, Integer => src_imm] if r64?(dst_reg) && imm32?(src_imm)
+        # REX.W + 81 /0 id
+        # MI: Operand 1: ModRM:r/m (r, w), Operand 2: imm8/16/32
+        insn(
+          prefix: REX_W,
+          opcode: 0x81,
+          mod_rm: ModRM[mod: Mod11, reg: 0, rm: dst_reg],
+          imm: imm32(src_imm),
+        )
       else
         raise NotImplementedError, "add: not-implemented operands: #{dst.inspect}, #{src.inspect}"
       end
diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index e5a8cb58d1..686e9554d6 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -508,6 +508,7 @@ module RubyVM::MJIT
 
       if flags & C.VM_CALL_KW_SPLAT != 0
         # recv_index calculation may not work for this
+        asm.incr_counter(:send_kw_splat)
         return CantCompile
       end
       assert_equal(ctx.sp_offset, ctx.stack_size) # TODO: support SP motion
@@ -525,12 +526,14 @@ module RubyVM::MJIT
         asm.jne(side_exit(jit, ctx))
       else
         # TODO: support more classes
+        asm.incr_counter(:send_guard_known_object)
         return CantCompile
       end
 
       # Do method lookup (vm_cc_cme(cc) != NULL)
       cme = C.rb_callable_method_entry(comptime_recv_klass, mid)
       if cme.nil?
+        asm.incr_counter(:send_missing_cme)
         return CantCompile # We don't support vm_call_method_name
       end
 
@@ -541,9 +544,11 @@ module RubyVM::MJIT
       when C.METHOD_VISI_PRIVATE
         # Allow only callsites without a receiver
         if flags & C.VM_CALL_FCALL == 0
+          asm.incr_counter(:send_private)
           return CantCompile
         end
       when C.METHOD_VISI_PROTECTED
+        asm.incr_counter(:send_protected)
         return CantCompile # TODO: support this
       else
         raise 'unreachable'
@@ -564,6 +569,7 @@ module RubyVM::MJIT
       when C.VM_METHOD_TYPE_ISEQ
         jit_call_iseq_setup(jit, ctx, asm, ci, cme, flags, argc)
       else
+        asm.incr_counter(:send_not_iseq)
         return CantCompile
       end
     end
@@ -582,6 +588,7 @@ module RubyVM::MJIT
 
       if flags & C.VM_CALL_TAILCALL != 0
         # We don't support vm_call_iseq_setup_tailcall
+        asm.incr_counter(:send_tailcall)
         return CantCompile
       end
       jit_call_iseq_setup_normal(jit, ctx, asm, ci, cme, flags, argc, iseq)
@@ -612,13 +619,14 @@ module RubyVM::MJIT
     def jit_push_frame(jit, ctx, asm, ci, cme, flags, argc, iseq, frame_type, next_pc)
       # TODO: stack overflow check
 
-      local_size = iseq.body.local_table_size
-      if local_size > 0
-        # TODO: support local variables
-        return CantCompile
+      local_size = iseq.body.local_table_size - iseq.body.param.size
+      local_size.times do |i|
+        asm.comment('set local variables') if i == 0
+        assert_equal(ctx.sp_offset, ctx.stack_size) # TODO: support SP motion
+        local_index = ctx.stack_size + i
+        asm.mov([SP, C.VALUE.size * local_index], Qnil)
       end
 
-      asm.comment('move SP register to callee SP')
       assert_equal(ctx.sp_offset, ctx.stack_size) # TODO: support SP motion
       sp_offset = ctx.stack_size + local_size + 3
       asm.add(SP, C.VALUE.size * sp_offset)
@@ -656,6 +664,7 @@ module RubyVM::MJIT
 
       # Stub cfp->jit_return
       return_ctx = ctx.dup
+      return_ctx.stack_size -= argc + ((flags & C.VM_CALL_ARGS_BLOCKARG == 0) ? 0 : 1) # Pop args
       return_ctx.sp_offset = 1 # SP is in the position after popping a receiver and arguments
       jit_return_stub = BlockStub.new(iseq: jit.iseq, pc: next_pc, ctx: return_ctx)
       jit_return = Assembler.new.then do |ocb_asm|
@@ -718,10 +727,12 @@ module RubyVM::MJIT
     def jit_caller_setup_arg(jit, ctx, asm, flags)
       if flags & C.VM_CALL_ARGS_SPLAT != 0
         # We don't support vm_caller_setup_arg_splat
+        asm.incr_counter(:send_args_splat)
         return CantCompile
       end
       if flags & (C.VM_CALL_KWARG | C.VM_CALL_KW_SPLAT) != 0
         # We don't support keyword args either
+        asm.incr_counter(:send_kwarg)
         return CantCompile
       end
     end
diff --git a/lib/ruby_vm/mjit/stats.rb b/lib/ruby_vm/mjit/stats.rb
index b372ffa07e..7fa9236257 100644
--- a/lib/ruby_vm/mjit/stats.rb
+++ b/lib/ruby_vm/mjit/stats.rb
@@ -34,6 +34,8 @@ module RubyVM::MJIT
       stats = runtime_stats
       $stderr.puts("***MJIT: Printing MJIT statistics on exit***")
 
+      print_counters(stats, prefix: 'send_', prompt: 'method call exit reasons')
+
       $stderr.puts "side_exit_count:       #{format('%10d', stats[:side_exit_count])}"
       $stderr.puts "total_insns_count:     #{format('%10d', stats[:total_insns_count])}" if stats.key?(:total_insns_count)
       $stderr.puts "vm_insns_count:        #{format('%10d', stats[:vm_insns_count])}" if stats.key?(:vm_insns_count)
@@ -43,6 +45,28 @@ module RubyVM::MJIT
       print_exit_counts(stats)
     end
 
+    def print_counters(stats, prefix:, prompt:)
+      $stderr.puts("#{prompt}: ")
+      counters = stats.filter { |key, _| key.start_with?(prefix) }
+      counters.filter! { |_, value| value != 0 }
+      counters.transform_keys! { |key| key.to_s.delete_prefix(prefix) }
+
+      if counters.empty?
+        $stderr.puts("    (all relevant counters are zero)")
+        return
+      end
+
+      counters = counters.to_a
+      counters.sort_by! { |(_, counter_value)| counter_value }
+      longest_name_length = counters.max_by { |(name, _)| name.length }.first.length
+      total = counters.sum { |(_, counter_value)| counter_value }
+
+      counters.reverse_each do |(name, value)|
+        percentage = value.fdiv(total) * 100
+        $stderr.printf("    %*s %10d (%4.1f%%)\n", longest_name_length, name, value, percentage)
+      end
+    end
+
     def print_exit_counts(stats, how_many: 20, padding: 2)
       exits = stats.filter_map { |name, count| [name.to_s.delete_prefix('exit_'), count] if name.start_with?('exit_') }.to_h
       return if exits.empty?