7 files changed, 112 insertions, 21 deletions
diff --git a/lib/ruby_vm/mjit/assembler.rb b/lib/ruby_vm/mjit/assembler.rb
index 87701a7102..c55df0235d 100644
--- a/lib/ruby_vm/mjit/assembler.rb
+++ b/lib/ruby_vm/mjit/assembler.rb
@@ -31,12 +31,13 @@ module RubyVM::MJIT
       @labels = {}
       @label_id = 0
       @comments = Hash.new { |h, k| h[k] = [] }
+      @blocks = Hash.new { |h, k| h[k] = [] }
       @stub_starts = Hash.new { |h, k| h[k] = [] }
       @stub_ends = Hash.new { |h, k| h[k] = [] }
     end
 
     def assemble(addr)
-      set_stub_addrs(addr)
+      set_code_addrs(addr)
       resolve_rel32(addr)
       resolve_labels
 
@@ -307,6 +308,12 @@ module RubyVM::MJIT
       @comments[@bytes.size] << message
     end
 
+    # Mark the starting address of a block
+    def block(block)
+      @blocks[@bytes.size] << block
+    end
+
+    # Mark the starting/ending addresses of a stub
     def stub(stub)
       @stub_starts[@bytes.size] << stub
       yield
@@ -495,8 +502,11 @@ module RubyVM::MJIT
       [Rel32.new(addr), Rel32Pad, Rel32Pad, Rel32Pad]
     end
 
-    def set_stub_addrs(write_addr)
+    def set_code_addrs(write_addr)
       (@bytes.size + 1).times do |index|
+        @blocks.fetch(index, []).each do |block|
+          block.start_addr = write_addr + index
+        end
         @stub_starts.fetch(index, []).each do |stub|
           stub.start_addr = write_addr + index
         end
diff --git a/lib/ruby_vm/mjit/block.rb b/lib/ruby_vm/mjit/block.rb
new file mode 100644
index 0000000000..58c21c9c2a
--- /dev/null
+++ b/lib/ruby_vm/mjit/block.rb
@@ -0,0 +1,6 @@
+class RubyVM::MJIT::Block < Struct.new(
+  :pc,         # @param [Integer] Starting PC
+  :start_addr, # @param [Integer] Starting address of this block's JIT code
+  :entry_exit, # @param [Integer] Address of entry exit (optional)
+)
+end
diff --git a/lib/ruby_vm/mjit/compiler.rb b/lib/ruby_vm/mjit/compiler.rb
index 3302b02ffa..56e23a1f77 100644
--- a/lib/ruby_vm/mjit/compiler.rb
+++ b/lib/ruby_vm/mjit/compiler.rb
@@ -1,10 +1,12 @@
 require 'ruby_vm/mjit/assembler'
+require 'ruby_vm/mjit/block'
 require 'ruby_vm/mjit/block_stub'
 require 'ruby_vm/mjit/code_block'
 require 'ruby_vm/mjit/context'
 require 'ruby_vm/mjit/exit_compiler'
 require 'ruby_vm/mjit/insn_compiler'
 require 'ruby_vm/mjit/instruction'
+require 'ruby_vm/mjit/invariants'
 require 'ruby_vm/mjit/jit_state'
 
 module RubyVM::MJIT
@@ -36,7 +38,7 @@ module RubyVM::MJIT
       @cb = CodeBlock.new(mem_block: mem_block, mem_size: mem_size / 2)
       @ocb = CodeBlock.new(mem_block: mem_block + mem_size / 2, mem_size: mem_size / 2, outlined: true)
       @exit_compiler = ExitCompiler.new
-      @insn_compiler = InsnCompiler.new(@ocb)
+      @insn_compiler = InsnCompiler.new(@ocb, @exit_compiler)
     end
 
     # Compile an ISEQ from its entry point.
@@ -66,8 +68,7 @@ module RubyVM::MJIT
       # Prepare the jump target
       new_asm = Assembler.new.tap do |asm|
         jit = JITState.new(iseq: stub.iseq, cfp:)
-        index = (stub.pc - stub.iseq.body.iseq_encoded.to_i) / C.VALUE.size
-        compile_block(asm, jit:, index:, ctx: stub.ctx)
+        compile_block(asm, jit:, pc: stub.pc, ctx: stub.ctx)
       end
 
       # Rewrite the stub
@@ -110,17 +111,24 @@ module RubyVM::MJIT
     end
 
     # @param asm [RubyVM::MJIT::Assembler]
-    def compile_block(asm, jit:, index: 0, ctx: Context.new)
+    def compile_block(asm, jit:, pc: jit.iseq.body.iseq_encoded.to_i, ctx: Context.new)
+      # Mark the block start address and prepare an exit code storage
+      jit.block = Block.new(pc:)
+      asm.block(jit.block)
+
+      # Compile each insn
       iseq = jit.iseq
+      index = (pc - iseq.body.iseq_encoded.to_i) / C.VALUE.size
       while index < iseq.body.iseq_size
         insn = self.class.decode_insn(iseq.body.iseq_encoded[index])
         jit.pc = (iseq.body.iseq_encoded + index).to_i
 
         case @insn_compiler.compile(jit, ctx, asm, insn)
         when EndBlock
+          # TODO: pad nops if entry exit exists
           break
         when CantCompile
-          @exit_compiler.compile_exit(jit, ctx, asm)
+          @exit_compiler.compile_side_exit(jit, ctx, asm)
           break
         end
         index += insn.len
diff --git a/lib/ruby_vm/mjit/exit_compiler.rb b/lib/ruby_vm/mjit/exit_compiler.rb
index bec2eab514..4577f65e20 100644
--- a/lib/ruby_vm/mjit/exit_compiler.rb
+++ b/lib/ruby_vm/mjit/exit_compiler.rb
@@ -5,16 +5,31 @@ module RubyVM::MJIT
       @gc_refs = []
     end
 
+    # Used for invalidating a block on entry.
+    # @param pc [Integer]
+    # @param asm [RubyVM::MJIT::Assembler]
+    def compile_entry_exit(pc, asm, cause:)
+      # Increment per-insn exit counter
+      incr_insn_exit(pc)
+
+      # TODO: Saving pc and sp may be needed later
+
+      # Restore callee-saved registers
+      asm.comment("#{cause}: entry exit")
+      asm.pop(SP)
+      asm.pop(EC)
+      asm.pop(CFP)
+
+      asm.mov(:rax, Qundef)
+      asm.ret
+    end
+
     # @param jit [RubyVM::MJIT::JITState]
     # @param ctx [RubyVM::MJIT::Context]
     # @param asm [RubyVM::MJIT::Assembler]
-    def compile_exit(jit, ctx, asm)
-      if C.mjit_opts.stats
-        insn = decode_insn(C.VALUE.new(jit.pc).*)
-        asm.comment("increment insn exit: #{insn.name}")
-        asm.mov(:rax, (C.mjit_insn_exits + insn.bin).to_i)
-        asm.add([:rax], 1) # TODO: lock
-      end
+    def compile_side_exit(jit, ctx, asm)
+      # Increment per-insn exit counter
+      incr_insn_exit(jit.pc)
 
       # Fix pc/sp offsets for the interpreter
       save_pc_and_sp(jit, ctx, asm)
@@ -50,11 +65,21 @@ module RubyVM::MJIT
 
     private
 
+    # @param pc [Integer]
+    def incr_insn_exit(pc)
+      if C.mjit_opts.stats
+        insn = decode_insn(C.VALUE.new(pc).*)
+        asm.comment("increment insn exit: #{insn.name}")
+        asm.mov(:rax, (C.mjit_insn_exits + insn.bin).to_i)
+        asm.add([:rax], 1) # TODO: lock
+      end
+    end
+
     # @param jit [RubyVM::MJIT::JITState]
     # @param ctx [RubyVM::MJIT::Context]
     # @param asm [RubyVM::MJIT::Assembler]
     def save_pc_and_sp(jit, ctx, asm)
-      # Update pc
+      # Update pc (TODO: manage PC offset?)
       asm.comment("save pc #{'and sp' if ctx.sp_offset != 0}")
       asm.mov(:rax, jit.pc) # rax = jit.pc
       asm.mov([CFP, C.rb_control_frame_t.offsetof(:pc)], :rax) # cfp->pc = rax
diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index debd8361fb..262d33387a 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -4,9 +4,11 @@ module RubyVM::MJIT
   # 5/101
   class InsnCompiler
     # @param ocb [CodeBlock]
-    def initialize(ocb)
+    # @param exit_compiler [RubyVM::MJIT::ExitCompiler]
+    def initialize(ocb, exit_compiler)
       @ocb = ocb
-      @exit_compiler = ExitCompiler.new
+      @exit_compiler = exit_compiler
+      @invariants = Invariants.new(ocb, exit_compiler)
       freeze
     end
 
@@ -265,6 +267,10 @@ module RubyVM::MJIT
         defer_compilation(jit, ctx, asm)
         return EndBlock
       end
+
+      unless @invariants.assume_bop_not_redefined(jit, C.INTEGER_REDEFINED_OP_FLAG, C.BOP_LT)
+        return CantCompile
+      end
       CantCompile
     end
 
@@ -350,7 +356,7 @@ module RubyVM::MJIT
     # @param ctx [RubyVM::MJIT::Context]
     def compile_side_exit(jit, ctx)
       asm = Assembler.new
-      @exit_compiler.compile_exit(jit, ctx, asm)
+      @exit_compiler.compile_side_exit(jit, ctx, asm)
       @ocb.write(asm)
     end
   end
diff --git a/lib/ruby_vm/mjit/invariants.rb b/lib/ruby_vm/mjit/invariants.rb
new file mode 100644
index 0000000000..c9d9b9fb6a
--- /dev/null
+++ b/lib/ruby_vm/mjit/invariants.rb
@@ -0,0 +1,35 @@
+require 'set'
+
+module RubyVM::MJIT
+  class Invariants
+    # @param ocb [CodeBlock]
+    # @param exit_compiler [RubyVM::MJIT::ExitCompiler]
+    def initialize(ocb, exit_compiler)
+      @ocb = ocb
+      @exit_compiler = exit_compiler
+      @bop_blocks = Set.new # TODO: actually invalidate this
+    end
+
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param klass [Integer]
+    # @param op [Integer]
+    def assume_bop_not_redefined(jit, klass, op)
+      return false unless C.BASIC_OP_UNREDEFINED_P(klass, op)
+
+      ensure_block_entry_exit(jit.block, cause: 'assume_bop_not_redefined')
+      @bop_blocks << jit.block
+      true
+    end
+
+    private
+
+    # @param block [RubyVM::MJIT::Block]
+    def ensure_block_entry_exit(block, cause:)
+      if block.entry_exit.nil?
+        asm = Assembler.new
+        @exit_compiler.compile_entry_exit(block.pc, asm, cause:)
+        block.entry_exit = @ocb.write(asm)
+      end
+    end
+  end
+end
diff --git a/lib/ruby_vm/mjit/jit_state.rb b/lib/ruby_vm/mjit/jit_state.rb
index 1a3e78b3da..97e5b7ffdd 100644
--- a/lib/ruby_vm/mjit/jit_state.rb
+++ b/lib/ruby_vm/mjit/jit_state.rb
@@ -1,8 +1,9 @@
 module RubyVM::MJIT
   class JITState < Struct.new(
-    :iseq,
-    :pc,  # @param [Integer] The JIT target PC
-    :cfp, # @param `RubyVM::MJIT::CPointer::Struct_rb_control_frame_t` The JIT source CFP (before MJIT is called)
+    :iseq,  # @param `RubyVM::MJIT::CPointer::Struct_rb_iseq_t`
+    :pc,    # @param [Integer] The JIT target PC
+    :cfp,   # @param `RubyVM::MJIT::CPointer::Struct_rb_control_frame_t` The JIT source CFP (before MJIT is called)
+    :block, # @param [RubyVM::MJIT::Block]
   )
     def operand(index)
       C.VALUE.new(pc)[index + 1]