Consider stack_size on side exits

lib/mjit/codegen.rb

@@ -4,12 +4,19 @@ module RubyVM::MJIT
   #  sp: rbx
   # scratch regs: rax
   class Codegen
-    def putnil(asm)
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::X86Assembler]
+    def putnil(ctx, asm)
       asm.mov([:rbx], Qnil)
+      ctx.stack_size += 1
       KeepCompiling
     end
 
-    def leave(asm)
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::X86Assembler]
+    def leave(ctx, asm)
+      assert_eq!(ctx.stack_size, 1)
+
       # pop the current frame (ec->cfp++)
       asm.add(:rsi, C.rb_control_frame_t.size) # rsi = cfp + 1
       asm.mov([:rdi, C.rb_execution_context_t.offsetof(:cfp)], :rsi) # ec->cfp = rsi
@@ -19,5 +26,13 @@ module RubyVM::MJIT
       asm.ret
       EndBlock
     end
+
+    private
+
+    def assert_eq!(left, right)
+      if left != right
+        raise "'#{left.inspect}' was not '#{right.inspect}'"
+      end
+    end
   end
 end

lib/mjit/context.rb

@@ -0,0 +1,8 @@
+class RubyVM::MJIT::Context < Struct.new(
+  :stack_size, # @param [Integer]
+)
+  def initialize(*)
+    super
+    self.stack_size ||= 0
+  end
+end

lib/ruby_vm/mjit/compiler.rb

@@ -1,4 +1,5 @@
 require 'mjit/codegen'
+require 'mjit/context'
 require 'mjit/instruction'
 require 'mjit/x86_assembler'
 
@@ -40,6 +41,7 @@ module RubyVM::MJIT
 
     private
 
+    # @param asm [RubyVM::MJIT::X86Assembler]
     def compile(asm)
       start_addr = write_addr
 
@@ -56,41 +58,50 @@ module RubyVM::MJIT
 
     #  ec: rdi
     # cfp: rsi
+    # @param asm [RubyVM::MJIT::X86Assembler]
     def compile_prologue(asm)
       asm.mov(:rbx, [:rsi, C.rb_control_frame_t.offsetof(:sp)]) # rbx = cfp->sp
     end
 
+    # @param asm [RubyVM::MJIT::X86Assembler]
     def compile_block(asm, iseq)
+      ctx = Context.new
       index = 0
       while index < iseq.body.iseq_size
         insn = decode_insn(iseq.body.iseq_encoded[index])
-        case compile_insn(asm, insn)
+        case compile_insn(ctx, asm, insn)
         when EndBlock
           break
         when CantCompile
-          compile_exit(asm, (iseq.body.iseq_encoded + index).to_i)
+          compile_exit(ctx, asm, (iseq.body.iseq_encoded + index).to_i)
           break
         end
         index += insn.len
       end
     end
 
-    def compile_insn(asm, insn)
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::X86Assembler]
+    def compile_insn(ctx, asm, insn)
       case insn.name
-      when :putnil then @codegen.putnil(asm)
-      #when :leave  then @codegen.leave(asm)
+      when :putnil then @codegen.putnil(ctx, asm)
+      when :leave  then @codegen.leave(ctx, asm)
       else CantCompile
       end
     end
 
-    def compile_exit(asm, exit_pc)
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::X86Assembler]
+    def compile_exit(ctx, asm, exit_pc)
       # update pc
       asm.mov(:rax, exit_pc) # rax = exit_pc
       asm.mov([:rsi, C.rb_control_frame_t.offsetof(:pc)], :rax) # cfp->pc = rax
 
-      # update sp (TODO: consider JIT state)
-      asm.add(:rbx, C.VALUE.size) # rbx += 1
-      asm.mov([:rsi, C.rb_control_frame_t.offsetof(:sp)], :rbx) # cfp->sp = rbx
+      # update sp
+      if ctx.stack_size > 0
+        asm.add(:rbx, C.VALUE.size * ctx.stack_size) # rbx += stack_size
+        asm.mov([:rsi, C.rb_control_frame_t.offsetof(:sp)], :rbx) # cfp->sp = rbx
+      end
 
       asm.mov(:rax, Qundef)
       asm.ret