A "MJIT Custom Compiler" JIT in ~100 lines of Ruby

hawthjit.rb

# $ ruby -v
# ruby 3.2.0dev (2022-09-11T14:08:14Z master 684353fc03) [x86_64-linux]
# $ ruby --mjit=pause --mjit-wait --mjit-min-calls=5 hawthjit.rb
# 246
# 246
# 246
# 246
# attempting to compile block in <main>
# can't compile putself
# attempting to compile double
# successfully compiled double!
# 246
# 246
# 246
# 246
# 246
# 246

# https://github.com/jhawthorn/asmjit-ruby
require "asmjit"

class Compiler
  include AsmJIT
  INSNS = RubyVM::MJIT.const_get(:INSNS)
  C = RubyVM::MJIT.const_get(:C)
  CantCompile = Class.new(StandardError)

  attr_reader :asm, :iseq
  def initialize(iseq)
    @iseq = iseq
    @code = CodeHolder.new
    @asm = X86::Assembler.new(@code)

    # Some registers to use as VM stack storage and a sentinel
    @stack = [:r8, :r9, nil]
  end

  def stack_push(from)
    @asm.mov(@stack[0], from)
    @stack.rotate!(1)
  end

  def stack_pop(into)
    @stack.rotate!(-1)
    @asm.mov(into, @stack[0])
  end

  def compile
    # The code we generate will be called as a function with the following
    # signature:
    # jit_func(rb_execution_context_t *ec, rb_control_frame_t *cfp)
    #          RDI                       RSI
    ec_reg = :rdi
    cfp_reg = :rsi

    pos = 0
    while pos < iseq.body.iseq_size
      insn = INSNS.fetch(C.rb_vm_insn_decode(iseq.body.iseq_encoded[pos]))

      case insn.name
      when :getlocal_WC_0
        # CRuby stores local variables relative to an EP pointer. We can get that
        # from the rb_control_frame_t struct we were passed (it's at offset 32)
        cfp_ep_ptr = X86.qword_ptr(cfp_reg, 32)
        asm.mov(:rax, cfp_ep_ptr) # RAX = cfp->ep

        # We read that local variable relative to the EP pointer we fetched
        # with the offset encoded in the iseq
        local0_offset = -iseq.body.iseq_encoded[pos + 1] * 8
        asm.mov(:rax, X86.qword_ptr(:rax, local0_offset))

        # "push" the value onto our "stack"
        stack_push(:rax)
      when :opt_plus
        # pop two numbers off of our stack into temporary registers
        stack_pop(:rcx)
        stack_pop(:rax)

        # Next we're going to add two numbers
        # We're going to assume that they're both "fixnum" and that adding them
        # together won't overflow (both potentially unsafe and incorrect
        # assumptions).
        # A Fixnum is CRuby's way of storing a 63-bit integer. The number is
        # shifted left by one bit, and the lowest bit is set to 1 (essentially
        # `x * 2 + 1`)
        # To deal with this we remove the tag bit from one of the numbers and
        # leave it on the other, so that the result is a fixnum

        asm.sub(:rcx, 1)
        asm.add(:rax, :rcx)

        stack_push(:rax)
      when :leave
        # Finally we need to pop the VM frame that was pushed in order to call this
        # method. This would normally be done by the "leave" instruction.
        # ec->cfp++
        ec_cfp_ptr = X86.qword_ptr(ec_reg, 0x10)
        asm.add(ec_cfp_ptr, 0x40)

        # Move our result into RAX, which is where values are returned
        # according to the calling convention.
        stack_pop(:rax)
        asm.ret
      else
        puts "can't compile #{insn.name}"
        raise CantCompile
      end

      pos += insn.len
    end

    puts "successfully compiled #{iseq.body.location.label}!"

    # AsmJIT will put this code into a piece of executable memory, and return us
    # the pointer to it as an Integer.
    @code.to_ptr
  end
end

class << RubyVM::MJIT
  def compile(iseq)
    puts "attempting to compile #{iseq.body.location.label}"
    Compiler.new(iseq).compile
  rescue Compiler::CantCompile
    0 # Don't jit
  end
end

def double(n)
  n + n
end

RubyVM::MJIT.resume

10.times do
  p double(123)
end