Check interrupts on leave

4 files changed, 121 insertions, 40 deletions

diff --git a/lib/mjit/codegen.rb b/lib/mjit/insn_compiler.rb
index 757e08c1a0..8fd8dfcef8 100644
--- a/lib/mjit/codegen.rb
+++ b/lib/mjit/insn_compiler.rb
@@ -3,20 +3,28 @@ module RubyVM::MJIT
   # cfp: rsi
   #  sp: rbx
   # scratch regs: rax
-  class Codegen
+  class InsnCompiler
+    # @param jit [RubyVM::MJIT::JITState]
     # @param ctx [RubyVM::MJIT::Context]
     # @param asm [RubyVM::MJIT::X86Assembler]
-    def putnil(ctx, asm)
+    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(ctx, asm)
+    def leave(jit, ctx, asm)
       assert_eq!(ctx.stack_size, 1)
-      # TODO: Check interrupts
+
+      # Check interrupts
+      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)
 
       # Pop the current frame (ec->cfp++)
       asm.add(CFP, C.rb_control_frame_t.size) # cfp = cfp + 1
diff --git a/lib/mjit/jit_state.rb b/lib/mjit/jit_state.rb
new file mode 100644
index 0000000000..819ecc128b
--- /dev/null
+++ b/lib/mjit/jit_state.rb
@@ -0,0 +1,4 @@
+class RubyVM::MJIT::JITState < Struct.new(
+  :pc, # @param [Integer]
+)
+end
diff --git a/lib/mjit/x86_assembler.rb b/lib/mjit/x86_assembler.rb
index c6c0cc2e19..1d010411db 100644
--- a/lib/mjit/x86_assembler.rb
+++ b/lib/mjit/x86_assembler.rb
@@ -2,6 +2,8 @@
 # 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 ###
@@ -10,9 +12,12 @@ module RubyVM::MJIT
 
     def initialize
       @bytes = []
+      @label_id = 0
+      @labels = {}
     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.
@@ -41,8 +46,25 @@ module RubyVM::MJIT
       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
@@ -73,14 +95,14 @@ module RubyVM::MJIT
           imm: imm64(src_imm),
         )
       # MOV r/m64, r64
-      in [[Symbol => dst_reg, Integer => dst_offset], Symbol => src_reg] if r64?(dst_reg) && r64?(src_reg) && imm8?(dst_offset)
+      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_offset,
+          disp: dst_disp,
         )
       # MOV r64, r/m64 (Mod 00)
       in [Symbol => dst_reg, [Symbol => src_reg]] if r64?(dst_reg) && r64?(src_reg)
@@ -136,6 +158,30 @@ module RubyVM::MJIT
       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 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
+
     private
 
     def insn(prefix: nil, opcode:, mod_rm: nil, disp: nil, imm: nil)
@@ -235,24 +281,40 @@ module RubyVM::MJIT
     end
 
     def imm8?(imm)
-      # TODO: consider negative values
-      imm <= 0xff
+      raise "negative imm not supported: #{imm}" if imm.negative? # TODO: support this
+      imm <= 0x7f # TODO: consider uimm
     end
 
     def imm32?(imm)
-      # TODO: consider negative values
+      raise "negative imm not supported: #{imm}" if imm.negative? # TODO: support this
       # TODO: consider rejecting small values
-      imm <= 0xffff_ffff
+      imm <= 0x7fff_ffff # TODO: consider uimm
     end
 
     def imm64?(imm)
-      # TODO: consider negative values
+      raise "negative imm not supported: #{imm}" if imm.negative? # TODO: support this
       # TODO: consider rejecting small values
-      imm <= 0xffff_ffff_ffff_ffff
+      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
diff --git a/lib/ruby_vm/mjit/compiler.rb b/lib/ruby_vm/mjit/compiler.rb
index c744c533a0..6641c749b6 100644
--- a/lib/ruby_vm/mjit/compiler.rb
+++ b/lib/ruby_vm/mjit/compiler.rb
@@ -1,6 +1,7 @@
-require 'mjit/codegen'
 require 'mjit/context'
+require 'mjit/insn_compiler'
 require 'mjit/instruction'
+require 'mjit/jit_state'
 require 'mjit/x86_assembler'
 
 module RubyVM::MJIT
@@ -21,11 +22,32 @@ module RubyVM::MJIT
   class Compiler
     attr_accessor :write_pos
 
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::X86Assembler]
+    def self.compile_exit(jit, ctx, asm)
+      # update pc
+      asm.mov(:rax, jit.pc) # rax = jit.pc
+      asm.mov([CFP, C.rb_control_frame_t.offsetof(:pc)], :rax) # cfp->pc = rax
+
+      # update sp
+      if ctx.stack_size > 0
+        asm.add(SP, C.VALUE.size * ctx.stack_size) # rbx += stack_size
+        asm.mov([CFP, C.rb_control_frame_t.offsetof(:sp)], SP) # cfp->sp = rbx
+      end
+
+      # Restore callee-saved registers
+      asm.pop(SP)
+
+      asm.mov(:rax, Qundef)
+      asm.ret
+    end
+
     # @param mem_block [Integer] JIT buffer address
     def initialize(mem_block)
       @mem_block = mem_block
       @write_pos = 0
-      @codegen = Codegen.new
+      @insn_compiler = InsnCompiler.new
     end
 
     # @param iseq [RubyVM::MJIT::CPointer::Struct]
@@ -79,58 +101,43 @@ module RubyVM::MJIT
 
     # @param asm [RubyVM::MJIT::X86Assembler]
     def compile_block(asm, iseq)
+      jit = JITState.new
       ctx = Context.new
+
       index = 0
       while index < iseq.body.iseq_size
         insn = decode_insn(iseq.body.iseq_encoded[index])
-        case compile_insn(ctx, asm, insn)
+        jit.pc = (iseq.body.iseq_encoded + index).to_i
+
+        case compile_insn(jit, ctx, asm, insn)
         when EndBlock
           break
         when CantCompile
-          compile_exit(ctx, asm, (iseq.body.iseq_encoded + index).to_i)
+          self.class.compile_exit(jit, ctx, asm)
           break
         end
         index += insn.len
       end
     end
 
+    # @param jit [RubyVM::MJIT::JITState]
     # @param ctx [RubyVM::MJIT::Context]
     # @param asm [RubyVM::MJIT::X86Assembler]
-    def compile_insn(ctx, asm, insn)
+    def compile_insn(jit, ctx, asm, insn)
       case insn.name
-      when :putnil then @codegen.putnil(ctx, asm)
-      when :leave  then @codegen.leave(ctx, asm)
+      when :putnil then @insn_compiler.putnil(jit, ctx, asm)
+      when :leave  then @insn_compiler.leave(jit, ctx, asm)
       else CantCompile
       end
     end
 
-    # @param ctx [RubyVM::MJIT::Context]
-    # @param asm [RubyVM::MJIT::X86Assembler]
-    def compile_exit(ctx, asm, exit_pc)
-      # update pc
-      asm.mov(:rax, exit_pc) # rax = exit_pc
-      asm.mov([CFP, C.rb_control_frame_t.offsetof(:pc)], :rax) # cfp->pc = rax
-
-      # update sp
-      if ctx.stack_size > 0
-        asm.add(SP, C.VALUE.size * ctx.stack_size) # rbx += stack_size
-        asm.mov([CFP, C.rb_control_frame_t.offsetof(:sp)], SP) # cfp->sp = rbx
-      end
-
-      # Restore callee-saved registers
-      asm.pop(SP)
-
-      asm.mov(:rax, Qundef)
-      asm.ret
-    end
-
     def decode_insn(encoded)
       INSNS.fetch(C.rb_vm_insn_decode(encoded))
     end
 
     def dump_disasm(from, to)
       C.dump_disasm(from, to).each do |address, mnemonic, op_str|
-        puts "  0x#{"%p" % address}: #{mnemonic} #{op_str}"
+        puts "  0x#{"%x" % address}: #{mnemonic} #{op_str}"
       end
       puts
     end