1 files changed, 44 insertions, 43 deletions

diff --git a/lib/ruby_vm/mjit/compiler.rb b/lib/ruby_vm/mjit/compiler.rb
index 30f2557f5b..892b250895 100644
--- a/lib/ruby_vm/mjit/compiler.rb
+++ b/lib/ruby_vm/mjit/compiler.rb
@@ -131,7 +131,7 @@ module RubyVM::MJIT
     # @param pc [Integer]
     def invalidate_blocks(iseq, pc)
       list_blocks(iseq, pc).each do |block|
-        invalidate_block(iseq, block)
+        invalidate_block(block)
       end
 
       # If they were the ISEQ's first blocks, re-compile MJIT entry as well
@@ -141,6 +141,48 @@ module RubyVM::MJIT
       end
     end
 
+    def invalidate_block(block)
+      iseq = block.iseq
+      # Remove this block from the version array
+      remove_block(iseq, block)
+
+      # Invalidate the block with entry exit
+      unless block.invalidated
+        @cb.with_write_addr(block.start_addr) do
+          asm = Assembler.new
+          asm.comment('invalidate_block')
+          asm.jmp(block.entry_exit)
+          @cb.write(asm)
+        end
+        block.invalidated = true
+      end
+
+      # Re-stub incoming branches
+      block.incoming.each do |branch_stub|
+        target = [branch_stub.target0, branch_stub.target1].compact.find do |target|
+          target.pc == block.pc && target.ctx == block.ctx
+        end
+        next if target.nil?
+        # TODO: Could target.address be a stub address? Is invalidation not needed in that case?
+
+        # If the target being re-generated is currently a fallthrough block,
+        # the fallthrough code must be rewritten with a jump to the stub.
+        if target.address == branch_stub.end_addr
+          branch_stub.shape = Default
+        end
+
+        target.address = Assembler.new.then do |ocb_asm|
+          @exit_compiler.compile_branch_stub(block.ctx, ocb_asm, branch_stub, target == branch_stub.target0)
+          @ocb.write(ocb_asm)
+        end
+        @cb.with_write_addr(branch_stub.start_addr) do
+          branch_asm = Assembler.new
+          branch_stub.compile.call(branch_asm)
+          @cb.write(branch_asm)
+        end
+      end
+    end
+
     private
 
     # Callee-saved: rbx, rsp, rbp, r12, r13, r14, r15
@@ -170,7 +212,7 @@ module RubyVM::MJIT
     # @param asm [RubyVM::MJIT::Assembler]
     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
-      block = Block.new(pc:, ctx: ctx.dup)
+      block = Block.new(iseq: jit.iseq, pc:, ctx: ctx.dup)
       jit.block = block
       asm.block(block)
 
@@ -222,47 +264,6 @@ module RubyVM::MJIT
       end
     end
 
-    def invalidate_block(iseq, block)
-      # Remove this block from the version array
-      remove_block(iseq, block)
-
-      # Invalidate the block with entry exit
-      unless block.invalidated
-        @cb.with_write_addr(block.start_addr) do
-          asm = Assembler.new
-          asm.comment('invalidate_block')
-          asm.jmp(block.entry_exit)
-          @cb.write(asm)
-        end
-        block.invalidated = true
-      end
-
-      # Re-stub incoming branches
-      block.incoming.each do |branch_stub|
-        target = [branch_stub.target0, branch_stub.target1].compact.find do |target|
-          target.pc == block.pc && target.ctx == block.ctx
-        end
-        next if target.nil?
-        # TODO: Could target.address be a stub address? Is invalidation not needed in that case?
-
-        # If the target being re-generated is currently a fallthrough block,
-        # the fallthrough code must be rewritten with a jump to the stub.
-        if target.address == branch_stub.end_addr
-          branch_stub.shape = Default
-        end
-
-        target.address = Assembler.new.then do |ocb_asm|
-          @exit_compiler.compile_branch_stub(block.ctx, ocb_asm, branch_stub, target == branch_stub.target0)
-          @ocb.write(ocb_asm)
-        end
-        @cb.with_write_addr(branch_stub.start_addr) do
-          branch_asm = Assembler.new
-          branch_stub.compile.call(branch_asm)
-          @cb.write(branch_asm)
-        end
-      end
-    end
-
     def list_blocks(iseq, pc)
       mjit_blocks(iseq)[pc].values
     end