Rewrite the stub if it's the last code

lib/ruby_vm/mjit/assembler.rb

@@ -31,7 +31,8 @@ module RubyVM::MJIT
       @labels = {}
       @label_id = 0
       @comments = Hash.new { |h, k| h[k] = [] }
-      @stubs = 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)
@@ -307,7 +308,10 @@ module RubyVM::MJIT
     end
 
     def stub(stub)
-      @stubs[@bytes.size] << stub
+      @stub_starts[@bytes.size] << stub
+      yield
+    ensure
+      @stub_ends[@bytes.size] << stub
     end
 
     def new_label(name)
@@ -492,9 +496,12 @@ module RubyVM::MJIT
     end
 
     def set_stub_addrs(write_addr)
-      @bytes.each_with_index do |byte, index|
-        @stubs.fetch(index, []).each do |stub|
-          stub.addr = write_addr + index
+      (@bytes.size + 1).times do |index|
+        @stub_starts.fetch(index, []).each do |stub|
+          stub.start_addr = write_addr + index
+        end
+        @stub_ends.fetch(index, []).each do |stub|
+          stub.end_addr = write_addr + index
           stub.freeze
         end
       end

lib/ruby_vm/mjit/block_stub.rb

@@ -1,7 +1,8 @@
 class RubyVM::MJIT::BlockStub < Struct.new(
-  :iseq, # @param [RubyVM::MJIT::CPointer::Struct_rb_iseq_struct] Jump target ISEQ
-  :pc,   # @param [Integer] Jump target pc
-  :ctx,  # @param [RubyVM::MJIT::Context] Jump target context
-  :addr, # @param [Integer] Jump source address to be re-generated
+  :iseq,       # @param [RubyVM::MJIT::CPointer::Struct_rb_iseq_struct] Stub target ISEQ
+  :pc,         # @param [Integer] Stub target pc
+  :ctx,        # @param [RubyVM::MJIT::Context] Stub target context
+  :start_addr, # @param [Integer] Stub source start address to be re-generated
+  :end_addr,   # @param [Integer] Stub source end address to be re-generated
 )
 end

lib/ruby_vm/mjit/code_block.rb

@@ -37,21 +37,25 @@ module RubyVM::MJIT
       start_addr
     end
 
-    def with_addr(addr)
-      old_write_pos = @write_pos
+    def set_write_addr(addr)
       @write_pos = addr - @mem_block
-      @comments.delete(addr) # TODO: clean up old comments for all overwritten insns?
+      @comments.delete(addr) # TODO: clean up old comments for all the overwritten range?
+    end
+
+    def with_write_addr(addr)
+      old_write_pos = @write_pos
+      set_addr(addr)
       yield
     ensure
       @write_pos = old_write_pos
     end
 
-    private
-
     def write_addr
       @mem_block + @write_pos
     end
 
+    private
+
     def dump_disasm(from, to)
       C.dump_disasm(from, to).each do |address, mnemonic, op_str|
         @comments.fetch(address, []).each do |comment|

lib/ruby_vm/mjit/compiler.rb

@@ -63,22 +63,28 @@ module RubyVM::MJIT
       # Update cfp->pc for `jit.at_current_insn?`
       cfp.pc = stub.pc
 
-      # Compile the jump target
-      new_addr = Assembler.new.then do |asm|
+      # 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)
-        @cb.write(asm)
       end
 
-      # Re-generate the jump source
-      @cb.with_addr(stub.addr) do
-        asm = Assembler.new
-        asm.comment('regenerate block stub')
-        asm.jmp(new_addr)
-        @cb.write(asm)
+      # Rewrite the stub
+      if @cb.write_addr == stub.end_addr
+        # If the stub jump is the last code, overwrite the jump with the new code.
+        @cb.set_write_addr(stub.start_addr)
+        @cb.write(new_asm)
+      else
+        # If the stub jump is old code, change the jump target to the new code.
+        new_addr = @cb.write(new_asm)
+        @cb.with_write_addr(stub.start_addr) do
+          asm = Assembler.new
+          asm.comment('regenerate block stub')
+          asm.jmp(new_addr)
+          @cb.write(asm)
+        end
       end
-      new_addr
     end
 
     private

lib/ruby_vm/mjit/insn_compiler.rb

@@ -224,8 +224,9 @@ module RubyVM::MJIT
       end
 
       asm.comment('defer_compilation: block stub')
-      asm.stub(block_stub)
-      asm.jmp(stub_hit)
+      asm.stub(block_stub) do
+        asm.jmp(stub_hit)
+      end
     end
 
     # @param jit [RubyVM::MJIT::JITState]