3 files changed, 67 insertions, 47 deletions
diff --git a/lib/ruby_vm/mjit/code_block.rb b/lib/ruby_vm/mjit/code_block.rb
new file mode 100644
index 0000000000..e47f5fbc0a
--- /dev/null
+++ b/lib/ruby_vm/mjit/code_block.rb
@@ -0,0 +1,58 @@
+module RubyVM::MJIT
+  class CodeBlock
+    # @param mem_block [Integer] JIT buffer address
+    # @param mem_size  [Integer] JIT buffer size
+    def initialize(mem_block:, mem_size:)
+      @comments  = Hash.new { |h, k| h[k] = [] }
+      @mem_block = mem_block
+      @mem_size  = mem_size
+      @write_pos = 0
+    end
+
+    # @param asm [RubyVM::MJIT::X86Assembler]
+    def compile(asm)
+      return 0 if @write_pos + asm.size >= @mem_size
+
+      start_addr = write_addr
+
+      # Write machine code
+      C.mjit_mark_writable
+      @write_pos += asm.compile(start_addr)
+      C.mjit_mark_executable
+
+      end_addr = write_addr
+
+      # Convert comment indexes to addresses
+      asm.comments.each do |index, comments|
+        @comments[start_addr + index] += comments
+      end
+      asm.comments.clear
+
+      # Dump disasm if --mjit-dump-disasm
+      if C.mjit_opts.dump_disasm && start_addr < end_addr
+        dump_disasm(start_addr, end_addr)
+      end
+      start_addr
+    end
+
+    private
+
+    def write_addr
+      @mem_block + @write_pos
+    end
+
+    def dump_disasm(from, to)
+      C.dump_disasm(from, to).each do |address, mnemonic, op_str|
+        @comments.fetch(address, []).each do |comment|
+          puts bold("  # #{comment}")
+        end
+        puts "  0x#{format("%x", address)}: #{mnemonic} #{op_str}"
+      end
+      puts
+    end
+
+    def bold(text)
+      "\e[1m#{text}\e[0m"
+    end
+  end
+end
diff --git a/lib/ruby_vm/mjit/compiler.rb b/lib/ruby_vm/mjit/compiler.rb
index 99984980f1..c5a7380226 100644
--- a/lib/ruby_vm/mjit/compiler.rb
+++ b/lib/ruby_vm/mjit/compiler.rb
@@ -1,3 +1,4 @@
+require 'ruby_vm/mjit/code_block'
 require 'ruby_vm/mjit/context'
 require 'ruby_vm/mjit/exit_compiler'
 require 'ruby_vm/mjit/insn_compiler'
@@ -28,10 +29,9 @@ module RubyVM::MJIT
     end
 
     # @param mem_block [Integer] JIT buffer address
-    def initialize(mem_block)
-      @comments = Hash.new { |h, k| h[k] = [] }
-      @mem_block = mem_block
-      @write_pos = 0
+    # @param mem_size  [Integer] JIT buffer size
+    def initialize(mem_block, mem_size)
+      @cb = CodeBlock.new(mem_block:, mem_size:)
       @exit_compiler = ExitCompiler.new
       @insn_compiler = InsnCompiler.new
     end
@@ -45,41 +45,13 @@ module RubyVM::MJIT
       asm.comment("Block: #{iseq.body.location.label}@#{pathobj_path(iseq.body.location.pathobj)}:#{iseq.body.location.first_lineno}")
       compile_prologue(asm)
       compile_block(asm, iseq)
-      iseq.body.jit_func = compile(asm)
+      iseq.body.jit_func = @cb.compile(asm)
     rescue Exception => e
       $stderr.puts e.full_message # TODO: check verbose
     end
 
-    def write_addr
-      @mem_block + @write_pos
-    end
-
     private
 
-    # @param asm [RubyVM::MJIT::X86Assembler]
-    def compile(asm)
-      start_addr = write_addr
-
-      # Write machine code
-      C.mjit_mark_writable
-      @write_pos += asm.compile(start_addr)
-      C.mjit_mark_executable
-
-      end_addr = write_addr
-
-      # Convert comment indexes to addresses
-      asm.comments.each do |index, comments|
-        @comments[start_addr + index] += comments
-      end
-      asm.comments.clear
-
-      # Dump disasm if --mjit-dump-disasm
-      if C.mjit_opts.dump_disasm && start_addr < end_addr
-        dump_disasm(start_addr, end_addr)
-      end
-      start_addr
-    end
-
     #  ec: rdi
     # cfp: rsi
     #
@@ -226,20 +198,6 @@ module RubyVM::MJIT
       end
     end
 
-    def dump_disasm(from, to)
-      C.dump_disasm(from, to).each do |address, mnemonic, op_str|
-        @comments.fetch(address, []).each do |comment|
-          puts bold("  # #{comment}")
-        end
-        puts "  0x#{format("%x", address)}: #{mnemonic} #{op_str}"
-      end
-      puts
-    end
-
-    def bold(text)
-      "\e[1m#{text}\e[0m"
-    end
-
     # vm_core.h: pathobj_path
     def pathobj_path(pathobj)
       if pathobj.is_a?(String)
diff --git a/lib/ruby_vm/mjit/x86_assembler.rb b/lib/ruby_vm/mjit/x86_assembler.rb
index fddd39f933..11f99cc78d 100644
--- a/lib/ruby_vm/mjit/x86_assembler.rb
+++ b/lib/ruby_vm/mjit/x86_assembler.rb
@@ -32,6 +32,10 @@ module RubyVM::MJIT
       @bytes.clear
     end
 
+    def size
+      @bytes.size
+    end
+
     def add(dst, src)
       case [dst, src]
       # ADD r/m64, imm8 (Mod 11)