3 files changed, 51 insertions, 22 deletions

diff --git a/lib/mjit/x86_64/assembler.rb b/lib/mjit/x86_64/assembler.rb
index 5458746c9b..be6480bba7 100644
--- a/lib/mjit/x86_64/assembler.rb
+++ b/lib/mjit/x86_64/assembler.rb
@@ -5,22 +5,34 @@ class RubyVM::MJIT::Assembler
     @bytes = []
   end
 
-  def compile(compiler)
-    with_dump_disasm(compiler) do
-      RubyVM::MJIT::C.mjit_mark_writable
-      write_bytes(compiler.write_addr, @bytes)
-      RubyVM::MJIT::C.mjit_mark_executable
-
-      compiler.write_pos += @bytes.size
-      @bytes.clear
-    end
+  def compile(compiler) = with_dump_disasm(compiler) do
+    RubyVM::MJIT::C.mjit_mark_writable
+    write_bytes(compiler.write_addr, @bytes)
+    RubyVM::MJIT::C.mjit_mark_executable
+
+    compiler.write_pos += @bytes.size
+    @bytes.clear
   end
 
-  def mov(_reg, val)
-    @bytes.push(0xb8, val, 0x00, 0x00, 0x00)
+  def add(_reg, imm)
+    #           REX.W [83]  RSI   ib
+    @bytes.push(0x48, 0x83, 0xc6, imm)
+  end
+
+  def mov(reg, val)
+    case reg
+    when :rax
+      #           REX.W [C7]  RAX   imm32
+      @bytes.push(0x48, 0xc7, 0xc0, val, 0x00, 0x00, 0x00)
+    else
+      #           REX.W [89]  [rdi+val],rsi
+      @bytes.push(0x48, 0x89, 0x77, reg.last)
+    end
   end
 
   def ret
+    # Near return
+    #           [C3]
     @bytes.push(0xc3)
   end
 
diff --git a/lib/ruby_vm/mjit/c_pointer.rb b/lib/ruby_vm/mjit/c_pointer.rb
index 743003f230..f0f34e949e 100644
--- a/lib/ruby_vm/mjit/c_pointer.rb
+++ b/lib/ruby_vm/mjit/c_pointer.rb
@@ -47,12 +47,15 @@ module RubyVM::MJIT # :nodoc: all
         type[@addr + offset / 8] = value
       end
 
-      # @param sizeof [Integer]
+      # @param size [Integer]
       # @param members [Hash{ Symbol => [Integer, RubyVM::MJIT::CType::*] }]
-      def self.define(sizeof, members)
+      def self.define(size, members)
         Class.new(self) do
           # Return the size of this type
-          define_singleton_method(:sizeof) { sizeof }
+          define_singleton_method(:size) { size }
+
+          # Return the offset to a field
+          define_singleton_method(:offsetof) { |field| members.fetch(field).last / 8 }
 
           # Get the offset of a member named +name+
           define_singleton_method(:offsetof) { |name|
@@ -62,9 +65,9 @@ module RubyVM::MJIT # :nodoc: all
 
           define_method(:initialize) do |addr = nil|
             if addr.nil? # TODO: get rid of this feature later
-              addr = Fiddle.malloc(sizeof)
+              addr = Fiddle.malloc(size)
             end
-            super(addr, sizeof, members)
+            super(addr, size, members)
           end
 
           members.each do |member, (type, offset, to_ruby)|
diff --git a/lib/ruby_vm/mjit/compiler.rb b/lib/ruby_vm/mjit/compiler.rb
index c42f75bf4d..4b980c1400 100644
--- a/lib/ruby_vm/mjit/compiler.rb
+++ b/lib/ruby_vm/mjit/compiler.rb
@@ -19,15 +19,29 @@ class RubyVM::MJIT::Compiler
   # @param iseq [RubyVM::MJIT::CPointer::Struct]
   def compile(iseq)
     return if iseq.body.location.label == '<main>'
-
-    iseq.body.jit_func = write_addr
-    asm = Assembler.new
-    asm.mov(:eax, Qundef)
-    asm.ret
-    asm.compile(self)
+    iseq.body.jit_func = compile_iseq(iseq)
   end
 
   def write_addr
     @mem_block + @write_pos
   end
+
+  private
+
+  # ec -> RDI, cfp -> RSI
+  def compile_iseq(iseq)
+    addr = write_addr
+    asm = Assembler.new
+
+    # 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)
+
+    # return a value
+    asm.mov(:rax, 7)
+    asm.ret
+
+    asm.compile(self)
+    addr
+  end
 end