3 files changed, 37 insertions, 37 deletions
diff --git a/lib/mjit/insn_compiler.rb b/lib/mjit/insn_compiler.rb
index c32dd808db..dbe3690dff 100644
--- a/lib/mjit/insn_compiler.rb
+++ b/lib/mjit/insn_compiler.rb
@@ -1,11 +1,11 @@
 module RubyVM::MJIT
   class InsnCompiler
-    def compile_putnil(_asm)
+    def on_putnil(_asm)
       # TODO
       KeepCompiling
     end
 
-    def compile_leave(asm)
+    def on_leave(asm)
       # pop the current frame (ec->cfp++)
       asm.add(:rsi, C.rb_control_frame_t.size)
       asm.mov([:rdi, C.rb_execution_context_t.offsetof(:cfp)], :rsi)
diff --git a/lib/mjit/x86_assembler.rb b/lib/mjit/x86_assembler.rb
index b89438f552..f04a9ea6b8 100644
--- a/lib/mjit/x86_assembler.rb
+++ b/lib/mjit/x86_assembler.rb
@@ -6,12 +6,15 @@ module RubyVM::MJIT
       @bytes = []
     end
 
-    def compile(compiler) = with_dump_disasm(compiler) do
-      C.mjit_mark_writable
-      write_bytes(compiler.write_addr, @bytes)
-      C.mjit_mark_executable
-
-      compiler.write_pos += @bytes.size
+    def compile(addr)
+      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.
+      @bytes.each_with_index do |byte, index|
+        writer[index] = byte
+      end
+      @bytes.size
+    ensure
       @bytes.clear
     end
 
@@ -36,27 +39,5 @@ module RubyVM::MJIT
       #           [C3]
       @bytes.push(0xc3)
     end
-
-    private
-
-    def with_dump_disasm(compiler)
-      from = compiler.write_addr
-      yield
-      to = compiler.write_addr
-      if C.mjit_opts.dump_disasm && from < to
-        C.dump_disasm(from, to).each do |address, mnemonic, op_str|
-          puts "  0x#{"%p" % address}: #{mnemonic} #{op_str}"
-        end
-      end
-    end
-
-    def write_bytes(addr, bytes)
-      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.
-      bytes.each_with_index do |byte, index|
-        writer[index] = byte
-      end
-    end
   end
 end
diff --git a/lib/ruby_vm/mjit/compiler.rb b/lib/ruby_vm/mjit/compiler.rb
index 1dfa3a3592..c9dadf6790 100644
--- a/lib/ruby_vm/mjit/compiler.rb
+++ b/lib/ruby_vm/mjit/compiler.rb
@@ -22,9 +22,9 @@ module RubyVM::MJIT
     end
 
     # @param iseq [RubyVM::MJIT::CPointer::Struct]
-    def compile(iseq)
+    def call(iseq)
       return if iseq.body.location.label == '<main>'
-      iseq.body.jit_func = compile_iseq(iseq)
+      iseq.body.jit_func = compile_block(iseq)
     rescue Exception => e
       # TODO: check --mjit-verbose
       $stderr.puts e.full_message
@@ -37,7 +37,7 @@ module RubyVM::MJIT
     private
 
     # ec -> RDI, cfp -> RSI
-    def compile_iseq(iseq)
+    def compile_block(iseq)
       addr = write_addr
       asm = X86Assembler.new
 
@@ -51,20 +51,39 @@ module RubyVM::MJIT
         index += insn.len
       end
 
-      asm.compile(self)
-      addr
+      compile(asm)
     end
 
     def compile_insn(asm, insn)
       case insn.name
-      when :putnil then @insn_compiler.compile_putnil(asm)
-      when :leave  then @insn_compiler.compile_leave(asm)
+      when :putnil then @insn_compiler.on_putnil(asm)
+      when :leave  then @insn_compiler.on_leave(asm)
       else raise NotImplementedError, "insn '#{insn.name}' is not supported yet"
       end
     end
 
+    def compile(asm)
+      start_addr = write_addr
+
+      C.mjit_mark_writable
+      @write_pos += asm.compile(start_addr)
+      C.mjit_mark_executable
+
+      end_addr = write_addr
+      if C.mjit_opts.dump_disasm && start_addr < end_addr
+        dump_disasm(start_addr, end_addr)
+      end
+      start_addr
+    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}"
+      end
+    end
   end
 end