9 files changed, 485 insertions, 41 deletions
diff --git a/lib/ruby_vm/mjit/c_pointer.rb b/lib/ruby_vm/mjit/c_pointer.rb
index 0ba9baa7cd..6bdf92b6cf 100644
--- a/lib/ruby_vm/mjit/c_pointer.rb
+++ b/lib/ruby_vm/mjit/c_pointer.rb
@@ -1,4 +1,4 @@
-module RubyVM::MJIT # :nodoc: all
+module RubyVM::MJIT
   # Every class under this namespace is a pointer. Even if the type is
   # immediate, it shouldn't be dereferenced until `*` is called.
   module CPointer
@@ -293,12 +293,12 @@ module RubyVM::MJIT # :nodoc: all
 
       # Dereference
       def *
-        byte = Fiddle::Pointer.new(@addr)[0, Fiddle::SIZEOF_CHAR].unpack1('c')
+        byte = Fiddle::Pointer.new(@addr)[0, Fiddle::SIZEOF_CHAR].unpack('c').first
         if @width == 1
           bit = (1 & (byte >> @offset))
           bit == 1
         elsif @width <= 8 && @offset == 0
-          bitmask = @width.times.sum { |i| 1 << i }
+          bitmask = @width.times.map { |i| 1 << i }.sum
           byte & bitmask
         else
           raise NotImplementedError.new("not-implemented bit field access: width=#{@width} offset=#{@offset}")
diff --git a/lib/ruby_vm/mjit/c_type.rb b/lib/ruby_vm/mjit/c_type.rb
index 9c965ad2fb..9e45d8d41c 100644
--- a/lib/ruby_vm/mjit/c_type.rb
+++ b/lib/ruby_vm/mjit/c_type.rb
@@ -2,7 +2,7 @@ require 'fiddle'
 require 'fiddle/pack'
 require_relative 'c_pointer'
 
-module RubyVM::MJIT # :nodoc: all
+module RubyVM::MJIT
   module CType
     module Struct
       # @param name [String]
diff --git a/lib/ruby_vm/mjit/compiler.rb b/lib/ruby_vm/mjit/compiler.rb
index 3dfea7088e..396e93cb04 100644
--- a/lib/ruby_vm/mjit/compiler.rb
+++ b/lib/ruby_vm/mjit/compiler.rb
@@ -1,8 +1,8 @@
-require 'mjit/context'
-require 'mjit/insn_compiler'
-require 'mjit/instruction'
-require 'mjit/jit_state'
-require 'mjit/x86_assembler'
+require 'ruby_vm/mjit/context'
+require 'ruby_vm/mjit/insn_compiler'
+require 'ruby_vm/mjit/instruction'
+require 'ruby_vm/mjit/jit_state'
+require 'ruby_vm/mjit/x86_assembler'
 
 module RubyVM::MJIT
   # Compilation status
diff --git a/lib/ruby_vm/mjit/context.rb b/lib/ruby_vm/mjit/context.rb
new file mode 100644
index 0000000000..2bc499cd4e
--- /dev/null
+++ b/lib/ruby_vm/mjit/context.rb
@@ -0,0 +1,8 @@
+class RubyVM::MJIT::Context < Struct.new(
+  :stack_size, # @param [Integer]
+)
+  def initialize(*)
+    super
+    self.stack_size ||= 0
+  end
+end
diff --git a/lib/ruby_vm/mjit/hooks.rb b/lib/ruby_vm/mjit/hooks.rb
deleted file mode 100644
index 3fb1004111..0000000000
--- a/lib/ruby_vm/mjit/hooks.rb
+++ /dev/null
@@ -1,32 +0,0 @@
-module RubyVM::MJIT::Hooks # :nodoc: all
-  C = RubyVM::MJIT.const_get(:C, false)
-
-  def self.on_bop_redefined(_redefined_flag, _bop)
-    C.mjit_cancel_all("BOP is redefined")
-  end
-
-  def self.on_cme_invalidate(_cme)
-    # to be used later
-  end
-
-  def self.on_ractor_spawn
-    C.mjit_cancel_all("Ractor is spawned")
-  end
-
-  def self.on_constant_state_changed(_id)
-    # to be used later
-  end
-
-  def self.on_constant_ic_update(_iseq, _ic, _insn_idx)
-    # to be used later
-  end
-
-  def self.on_tracing_invalidate_all(new_iseq_events)
-    # Stop calling all JIT-ed code. We can't rewrite existing JIT-ed code to trace_ insns for now.
-    # :class events are triggered only in ISEQ_TYPE_CLASS, but mjit_target_iseq_p ignores such iseqs.
-    # Thus we don't need to cancel JIT-ed code for :class events.
-    if new_iseq_events != C.RUBY_EVENT_CLASS
-      C.mjit_cancel_all("TracePoint is enabled")
-    end
-  end
-end
diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
new file mode 100644
index 0000000000..9c3e2f2a95
--- /dev/null
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -0,0 +1,51 @@
+module RubyVM::MJIT
+  #  ec: rdi
+  # cfp: rsi
+  #  sp: rbx
+  # scratch regs: rax
+  class InsnCompiler
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::X86Assembler]
+    def putnil(jit, ctx, asm)
+      asm.mov([SP], Qnil)
+      ctx.stack_size += 1
+      KeepCompiling
+    end
+
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::X86Assembler]
+    def leave(jit, ctx, asm)
+      assert_eq!(ctx.stack_size, 1)
+
+      asm.comment("RUBY_VM_CHECK_INTS(ec)")
+      asm.mov(:eax, [EC, C.rb_execution_context_t.offsetof(:interrupt_flag)])
+      asm.test(:eax, :eax)
+      asm.jz(not_interrupted = asm.new_label(:not_interrupted))
+      Compiler.compile_exit(jit, ctx, asm) # TODO: use ocb
+      asm.write_label(not_interrupted)
+
+      asm.comment("pop stack frame")
+      asm.add(CFP, C.rb_control_frame_t.size) # cfp = cfp + 1
+      asm.mov([EC, C.rb_execution_context_t.offsetof(:cfp)], CFP) # ec->cfp = cfp
+
+      # Return a value
+      asm.mov(:rax, [SP])
+
+      # Restore callee-saved registers
+      asm.pop(SP)
+
+      asm.ret
+      EndBlock
+    end
+
+    private
+
+    def assert_eq!(left, right)
+      if left != right
+        raise "'#{left.inspect}' was not '#{right.inspect}'"
+      end
+    end
+  end
+end
diff --git a/lib/ruby_vm/mjit/jit_state.rb b/lib/ruby_vm/mjit/jit_state.rb
new file mode 100644
index 0000000000..819ecc128b
--- /dev/null
+++ b/lib/ruby_vm/mjit/jit_state.rb
@@ -0,0 +1,4 @@
+class RubyVM::MJIT::JITState < Struct.new(
+  :pc, # @param [Integer]
+)
+end
diff --git a/lib/ruby_vm/mjit/stats.rb b/lib/ruby_vm/mjit/stats.rb
new file mode 100644
index 0000000000..263948bc0e
--- /dev/null
+++ b/lib/ruby_vm/mjit/stats.rb
@@ -0,0 +1,68 @@
+# frozen_string_literal: true
+module RubyVM::MJIT
+  def self.runtime_stats
+    stats = {}
+
+    # Insn exits
+    INSNS.each_value do |insn|
+      exits = C.mjit_insn_exits[insn.bin]
+      if exits > 0
+        stats[:"exit_#{insn.name}"] = exits
+      end
+    end
+
+    # Runtime stats
+    C.rb_mjit_runtime_counters.members.each do |member|
+      stats[member] = C.rb_mjit_counters.public_send(member)
+    end
+
+    # Other stats are calculated here
+    stats[:side_exit_count] = stats.select { |name, _count| name.start_with?('exit_') }.sum(&:last)
+    if stats[:vm_insns_count] > 0
+      retired_in_mjit = stats[:mjit_insns_count] - stats[:side_exit_count]
+      stats[:total_insns_count] = retired_in_mjit + stats[:vm_insns_count]
+      stats[:ratio_in_mjit] = 100.0 * retired_in_mjit / stats[:total_insns_count]
+    end
+
+    stats
+  end
+
+  at_exit do
+    if C.mjit_opts.stats
+      print_stats
+    end
+  end
+
+  class << self
+    private
+
+    def print_stats
+      stats = runtime_stats
+      $stderr.puts("***MJIT: Printing MJIT statistics on exit***")
+
+      $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)
+      $stderr.puts "mjit_insns_count:      #{format('%10d', stats[:mjit_insns_count])}"
+      $stderr.puts "ratio_in_yjit:         #{format('%9.1f', stats[:ratio_in_mjit])}%" if stats.key?(:ratio_in_mjit)
+
+      print_exit_counts(stats)
+    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?
+
+      top_exits = exits.sort_by { |_name, count| -count }.first(how_many).to_h
+      total_exits = exits.values.sum
+      $stderr.puts "Top-#{top_exits.size} most frequent exit ops (#{format("%.1f", 100.0 * top_exits.values.sum / total_exits)}% of exits):"
+
+      name_width  = top_exits.map { |name, _count| name.length }.max + padding
+      count_width = top_exits.map { |_name, count| count.to_s.length }.max + padding
+      top_exits.each do |name, count|
+        ratio = 100.0 * count / total_exits
+        $stderr.puts "#{format("%#{name_width}s", name)}: #{format("%#{count_width}d", count)} (#{format('%.1f', ratio)}%)"
+      end
+    end
+  end
+end
diff --git a/lib/ruby_vm/mjit/x86_assembler.rb b/lib/ruby_vm/mjit/x86_assembler.rb
new file mode 100644
index 0000000000..890fa2b80a
--- /dev/null
+++ b/lib/ruby_vm/mjit/x86_assembler.rb
@@ -0,0 +1,345 @@
+# frozen_string_literal: true
+# https://www.intel.com/content/dam/develop/public/us/en/documents/325383-sdm-vol-2abcd.pdf
+module RubyVM::MJIT
+  class X86Assembler
+    class Label < Data.define(:id, :name); end
+
+    ByteWriter = CType::Immediate.parse('char')
+
+    ### prefix ###
+    # REX =   0100WR0B
+    REX_W = 0b01001000
+
+    attr_reader :comments
+
+    def initialize
+      @bytes = []
+      @labels = {}
+      @label_id = 0
+      @comments = Hash.new { |h, k| h[k] = [] }
+    end
+
+    def compile(addr)
+      link_labels
+      writer = ByteWriter.new(addr)
+      # If you pack bytes containing \x00, Ruby fails to recognize bytes after \x00.
+      # So writing byte by byte to avoid hitting that situation.
+      @bytes.each_with_index do |byte, index|
+        writer[index] = byte
+      end
+      @bytes.size
+    ensure
+      @bytes.clear
+    end
+
+    def add(dst, src)
+      case [dst, src]
+      # ADD r/m64, imm8 (Mod 11)
+      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: mod_rm(mod: 0b11, rm: reg_code(dst_reg)),
+          imm: imm8(src_imm),
+        )
+      # ADD r/m64, imm8 (Mod 00)
+      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: mod_rm(mod: 0b00, rm: reg_code(dst_reg)), # Mod 00: [reg]
+          imm: imm8(src_imm),
+        )
+      else
+        raise NotImplementedError, "add: not-implemented operands: #{dst.inspect}, #{src.inspect}"
+      end
+    end
+
+    # JZ rel8
+    # @param [RubyVM::MJIT::X86Assembler::Label] label
+    def jz(label)
+      # 74 cb
+      insn(opcode: 0x74)
+      @bytes.push(label)
+    end
+
+    def mov(dst, src)
+      case [dst, src]
+      # MOV r32 r/m32 (Mod 01)
+      in [Symbol => dst_reg, [Symbol => src_reg, Integer => src_disp]] if r32?(dst_reg) && imm8?(src_disp)
+        # 8B /r
+        # RM: Operand 1: ModRM:reg (w), Operand 2: ModRM:r/m (r)
+        insn(
+          opcode: 0x8b,
+          mod_rm: mod_rm(mod: 0b01, reg: reg_code(dst_reg), rm: reg_code(src_reg)), # Mod 01: [reg]+disp8
+          disp: src_disp,
+        )
+      # MOV r/m64, imm32 (Mod 00)
+      in [[Symbol => dst_reg], Integer => src_imm] if r64?(dst_reg)
+        # REX.W + C7 /0 id
+        # MI: Operand 1: ModRM:r/m (w), Operand 2: imm8/16/32/64
+        insn(
+          prefix: REX_W,
+          opcode: 0xc7,
+          mod_rm: mod_rm(mod: 0b00, rm: reg_code(dst_reg)), # Mod 00: [reg]
+          imm: imm32(src_imm),
+        )
+      # MOV r/m64, imm32 (Mod 11)
+      in [Symbol => dst_reg, Integer => src_imm] if r64?(dst_reg) && imm32?(src_imm)
+        # REX.W + C7 /0 id
+        # MI: Operand 1: ModRM:r/m (w), Operand 2: imm8/16/32/64
+        insn(
+          prefix: REX_W,
+          opcode: 0xc7,
+          mod_rm: mod_rm(mod: 0b11, rm: reg_code(dst_reg)), # Mod 11: reg
+          imm: imm32(src_imm),
+        )
+      # MOV r64, imm64
+      in [Symbol => dst_reg, Integer => src_imm] if r64?(dst_reg) && imm64?(src_imm)
+        # REX.W + B8+ rd io
+        # OI: Operand 1: opcode + rd (w), Operand 2: imm8/16/32/64
+        insn(
+          prefix: REX_W,
+          opcode: 0xb8 + reg_code(dst_reg),
+          imm: imm64(src_imm),
+        )
+      # MOV r/m64, r64
+      in [[Symbol => dst_reg, Integer => dst_disp], Symbol => src_reg] if r64?(dst_reg) && r64?(src_reg) && imm8?(dst_disp)
+        # REX.W + 89 /r
+        # MR: Operand 1: ModRM:r/m (w), Operand 2: ModRM:reg (r)
+        insn(
+          prefix: REX_W,
+          opcode: 0x89,
+          mod_rm: mod_rm(mod: 0b01, reg: reg_code(src_reg), rm: reg_code(dst_reg)), # Mod 01: [reg]+disp8
+          disp: dst_disp,
+        )
+      # MOV r64, r/m64 (Mod 00)
+      in [Symbol => dst_reg, [Symbol => src_reg]] if r64?(dst_reg) && r64?(src_reg)
+        # REX.W + 8B /r
+        # RM: Operand 1: ModRM:reg (w), Operand 2: ModRM:r/m (r)
+        insn(
+          prefix: REX_W,
+          opcode: 0x8b,
+          mod_rm: mod_rm(mod: 0b00, reg: reg_code(dst_reg), rm: reg_code(src_reg)), # Mod 00: [reg]
+        )
+      # MOV r64, r/m64 (Mod 01)
+      in [Symbol => dst_reg, [Symbol => src_reg, Integer => src_offset]] if r64?(dst_reg) && r64?(src_reg) && imm8?(src_offset)
+        # REX.W + 8B /r
+        # RM: Operand 1: ModRM:reg (w), Operand 2: ModRM:r/m (r)
+        insn(
+          prefix: REX_W,
+          opcode: 0x8b,
+          mod_rm: mod_rm(mod: 0b01, reg: reg_code(dst_reg), rm: reg_code(src_reg)), # Mod 01: [reg]+disp8
+          disp: src_offset,
+        )
+      else
+        raise NotImplementedError, "mov: not-implemented operands: #{dst.inspect}, #{src.inspect}"
+      end
+    end
+
+    def push(src)
+      case src
+      # PUSH r64
+      in Symbol => src_reg if r64?(src_reg)
+        # 50+rd
+        # O: Operand 1: opcode + rd (r)
+        insn(opcode: 0x50 + reg_code(src_reg))
+      else
+        raise NotImplementedError, "push: not-implemented operands: #{src.inspect}"
+      end
+    end
+
+    def pop(dst)
+      case dst
+      # POP r64
+      in Symbol => dst_reg if r64?(dst_reg)
+        # 58+ rd
+        # O: Operand 1: opcode + rd (r)
+        insn(opcode: 0x58 + reg_code(dst_reg))
+      else
+        raise NotImplementedError, "pop: not-implemented operands: #{dst.inspect}"
+      end
+    end
+
+    # RET
+    def ret
+      # Near return: A return to a procedure within the current code segment
+      insn(opcode: 0xc3)
+    end
+
+    def test(left, right)
+      case [left, right]
+      # TEST r/m32, r32 (Mod 11)
+      in [Symbol => left_reg, Symbol => right_reg] if r32?(left_reg) && r32?(right_reg)
+        # 85 /r
+        # MR: Operand 1: ModRM:r/m (r), Operand 2: ModRM:reg (r)
+        insn(
+          opcode: 0x85,
+          mod_rm: mod_rm(mod: 0b11, reg: reg_code(right_reg), rm: reg_code(left_reg)), # Mod 11: reg
+        )
+      else
+        raise NotImplementedError, "pop: not-implemented operands: #{dst.inspect}"
+      end
+    end
+
+    def comment(message)
+      @comments[@bytes.size] << message
+    end
+
+    def new_label(name)
+      Label.new(id: @label_id += 1, name:)
+    end
+
+    # @param [RubyVM::MJIT::X86Assembler::Label] label
+    def write_label(label)
+      @labels[label] = @bytes.size
+    end
+
+    def incr_counter(name)
+      if C.mjit_opts.stats
+        comment("increment counter #{name}")
+        mov(:rax, C.rb_mjit_counters[name].to_i)
+        add([:rax], 1) # TODO: lock
+      end
+    end
+
+    private
+
+    def insn(prefix: nil, opcode:, mod_rm: nil, disp: nil, imm: nil)
+      if prefix
+        @bytes.push(prefix)
+      end
+      @bytes.push(opcode)
+      if mod_rm
+        @bytes.push(mod_rm)
+      end
+      if disp
+        if disp < 0 || disp > 0xff # TODO: support displacement in 2 or 4 bytes as well
+          raise NotImplementedError, "not-implemented disp: #{disp}"
+        end
+        @bytes.push(disp)
+      end
+      if imm
+        @bytes.push(*imm)
+      end
+    end
+
+    def reg_code(reg)
+      case reg
+      when :al, :ax, :eax, :rax then 0
+      when :cl, :cx, :ecx, :rcx then 1
+      when :dl, :dx, :edx, :rdx then 2
+      when :bl, :bx, :ebx, :rbx then 3
+      when :ah, :sp, :esp, :rsp then 4
+      when :ch, :bp, :ebp, :rbp then 5
+      when :dh, :si, :esi, :rsi then 6
+      when :bh, :di, :edi, :rdi then 7
+      else raise ArgumentError, "unexpected reg: #{reg.inspect}"
+      end
+    end
+
+    # Table 2-2. 32-Bit Addressing Forms with the ModR/M Byte
+    #
+    #  7  6  5  4  3  2  1  0
+    # +--+--+--+--+--+--+--+--+
+    # | Mod | Reg/   | R/M    |
+    # |     | Opcode |        |
+    # +--+--+--+--+--+--+--+--+
+    #
+    # The r/m field can specify a register as an operand or it can be combined
+    # with the mod field to encode an addressing mode.
+    #
+    # /0: R/M is 0 (not used)
+    # /r: R/M is a register
+    def mod_rm(mod:, reg: 0, rm: 0)
+      if mod > 0b11
+        raise ArgumentError, "too large Mod: #{mod}"
+      end
+      if reg > 0b111
+        raise ArgumentError, "too large Reg/Opcode: #{reg}"
+      end
+      if rm > 0b111
+        raise ArgumentError, "too large R/M: #{rm}"
+      end
+      (mod << 6) + (reg << 3) + rm
+    end
+
+    # ib: 1 byte
+    def imm8(imm)
+      unless imm8?(imm)
+        raise ArgumentError, "unexpected imm8: #{imm}"
+      end
+      imm_bytes(imm, 1)
+    end
+
+    # id: 4 bytes
+    def imm32(imm)
+      unless imm32?(imm)
+        raise ArgumentError, "unexpected imm32: #{imm}"
+      end
+      imm_bytes(imm, 4)
+    end
+
+    # io: 8 bytes
+    def imm64(imm)
+      unless imm64?(imm)
+        raise ArgumentError, "unexpected imm64: #{imm}"
+      end
+      imm_bytes(imm, 8)
+    end
+
+    def imm_bytes(imm, num_bytes)
+      bytes = []
+      bits = imm
+      num_bytes.times do
+        bytes << (bits & 0xff)
+        bits >>= 8
+      end
+      if bits != 0
+        raise ArgumentError, "unexpected imm with #{num_bytes} bytes: #{imm}"
+      end
+      bytes
+    end
+
+    def imm8?(imm)
+      raise "negative imm not supported: #{imm}" if imm.negative? # TODO: support this
+      imm <= 0x7f # TODO: consider uimm
+    end
+
+    def imm32?(imm)
+      raise "negative imm not supported: #{imm}" if imm.negative? # TODO: support this
+      # TODO: consider rejecting small values
+      imm <= 0x7fff_ffff # TODO: consider uimm
+    end
+
+    def imm64?(imm)
+      raise "negative imm not supported: #{imm}" if imm.negative? # TODO: support this
+      # TODO: consider rejecting small values
+      imm <= 0x7fff_ffff_ffff_ffff # TODO: consider uimm
+    end
+
+    def r32?(reg)
+      reg.start_with?('e')
+    end
+
+    def r64?(reg)
+      reg.start_with?('r')
+    end
+
+    def link_labels
+      @bytes.each_with_index do |byte, index|
+        if byte.is_a?(Label)
+          src_index = index + 1 # offset 1 byte for rel8 itself
+          dst_index = @labels.fetch(byte)
+          rel8 = dst_index - src_index
+          raise "unexpected offset: #{rel8}" unless imm8?(rel8)
+          @bytes[index] = rel8
+        end
+      end
+    end
+  end
+end