Implement getblockparamproxy

6 files changed, 109 insertions, 3 deletions

diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index a4516e73e0..5c0ad97966 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -24,14 +24,14 @@ module RubyVM::MJIT
       asm.incr_counter(:mjit_insns_count)
       asm.comment("Insn: #{insn.name}")
 
-      # 63/101
+      # 64/101
       case insn.name
       when :nop then nop(jit, ctx, asm)
       when :getlocal then getlocal(jit, ctx, asm)
       when :setlocal then setlocal(jit, ctx, asm)
       # getblockparam
       # setblockparam
-      # getblockparamproxy
+      when :getblockparamproxy then getblockparamproxy(jit, ctx, asm)
       # getspecial
       # setspecial
       when :getinstancevariable then getinstancevariable(jit, ctx, asm)
@@ -165,7 +165,77 @@ module RubyVM::MJIT
 
     # getblockparam
     # setblockparam
-    # getblockparamproxy
+
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
+    def getblockparamproxy(jit, ctx, asm)
+      # To get block_handler
+      unless jit.at_current_insn?
+        defer_compilation(jit, ctx, asm)
+        return EndBlock
+      end
+
+      starting_context = ctx.dup # make a copy for use with jit_chain_guard
+
+      # A mirror of the interpreter code. Checking for the case
+      # where it's pushing rb_block_param_proxy.
+      side_exit = side_exit(jit, ctx)
+
+      # EP level
+      level = jit.operand(1)
+
+      # Peek at the block handler so we can check whether it's nil
+      comptime_handler = jit.peek_at_block_handler(level)
+
+      # When a block handler is present, it should always be a GC-guarded
+      # pointer (VM_BH_ISEQ_BLOCK_P)
+      if comptime_handler != 0 && comptime_handler & 0x3 != 0x1
+        asm.incr_counter(:getblockpp_not_gc_guarded)
+        return CantCompile
+      end
+
+      # Load environment pointer EP from CFP
+      ep_reg = :rax
+      jit_get_ep(asm, level, reg: ep_reg)
+
+      # Bail when VM_ENV_FLAGS(ep, VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM) is non zero
+      asm.test([ep_reg, C.VALUE.size * C.VM_ENV_DATA_INDEX_FLAGS], C.VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM)
+      asm.jnz(counted_exit(side_exit, :getblockpp_block_param_modified))
+
+      # Load the block handler for the current frame
+      # note, VM_ASSERT(VM_ENV_LOCAL_P(ep))
+      block_handler = :rax
+      asm.mov(block_handler, [ep_reg, C.VALUE.size * C.VM_ENV_DATA_INDEX_SPECVAL])
+
+      # Specialize compilation for the case where no block handler is present
+      if comptime_handler == 0
+        # Bail if there is a block handler
+        asm.cmp(block_handler, 0)
+
+        jit_chain_guard(:jnz, jit, starting_context, asm, counted_exit(side_exit, :getblockpp_block_handler_none))
+
+        putobject(jit, ctx, asm, val: Qnil)
+      else
+        # Block handler is a tagged pointer. Look at the tag. 0x03 is from VM_BH_ISEQ_BLOCK_P().
+        asm.and(block_handler, 0x3)
+
+        # Bail unless VM_BH_ISEQ_BLOCK_P(bh). This also checks for null.
+        asm.cmp(block_handler, 0x1)
+
+        jit_chain_guard(:jnz, jit, starting_context, asm, counted_exit(side_exit, :getblockpp_not_iseq_block))
+
+        # Push rb_block_param_proxy. It's a root, so no need to use jit_mov_gc_ptr.
+        top = ctx.stack_push
+        asm.mov(:rax, C.rb_block_param_proxy)
+        asm.mov(top, :rax)
+      end
+
+      jump_to_next_insn(jit, ctx, asm)
+
+      EndBlock
+    end
+
     # getspecial
     # setspecial
 
diff --git a/lib/ruby_vm/mjit/jit_state.rb b/lib/ruby_vm/mjit/jit_state.rb
index aec4e51c48..decafee579 100644
--- a/lib/ruby_vm/mjit/jit_state.rb
+++ b/lib/ruby_vm/mjit/jit_state.rb
@@ -37,5 +37,21 @@ module RubyVM::MJIT
     def peek_at_self
       C.to_ruby(cfp.self)
     end
+
+    def peek_at_block_handler(level)
+      ep = ep_at_level(cfp, level:)
+      ep[C.VM_ENV_DATA_INDEX_SPECVAL]
+    end
+
+    private
+
+    def ep_at_level(cfp, level:)
+      ep = cfp.ep
+      level.times do
+        # VM_ENV_PREV_EP
+        ep = C.VALUE.new(ep[C.VM_ENV_DATA_INDEX_SPECVAL] & ~0x03)
+      end
+      ep
+    end
   end
 end
diff --git a/lib/ruby_vm/mjit/stats.rb b/lib/ruby_vm/mjit/stats.rb
index 834b9dae61..b605f3c1ad 100644
--- a/lib/ruby_vm/mjit/stats.rb
+++ b/lib/ruby_vm/mjit/stats.rb
@@ -36,6 +36,7 @@ module RubyVM::MJIT
 
       print_counters(stats, prefix: 'send_', prompt: 'method call exit reasons')
       print_counters(stats, prefix: 'invokesuper_', prompt: 'invokesuper exit reasons')
+      print_counters(stats, prefix: 'getblockpp_', prompt: 'getblockparamproxy exit reasons')
       print_counters(stats, prefix: 'getivar_', prompt: 'getinstancevariable exit reasons')
       print_counters(stats, prefix: 'setivar_', prompt: 'setinstancevariable exit reasons')
       print_counters(stats, prefix: 'optaref_', prompt: 'opt_aref exit reasons')
diff --git a/mjit_c.h b/mjit_c.h
index 80c4f47681..4e01dff7b3 100644
--- a/mjit_c.h
+++ b/mjit_c.h
@@ -190,6 +190,11 @@ MJIT_RUNTIME_COUNTERS(
     expandarray_not_array,
     expandarray_rhs_too_small,
 
+    getblockpp_block_param_modified,
+    getblockpp_block_handler_none,
+    getblockpp_not_gc_guarded,
+    getblockpp_not_iseq_block,
+
     compiled_block_count
 )
 #undef MJIT_RUNTIME_COUNTERS
diff --git a/mjit_c.rb b/mjit_c.rb
index b3c86456ac..b95b9f0c32 100644
--- a/mjit_c.rb
+++ b/mjit_c.rb
@@ -702,6 +702,10 @@ module RubyVM::MJIT # :nodoc: all
     Primitive.cexpr! %q{ UINT2NUM(VM_FRAME_FLAG_CFRAME_KW) }
   end
 
+  def C.VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM
+    Primitive.cexpr! %q{ UINT2NUM(VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM) }
+  end
+
   def C.VM_FRAME_MAGIC_CFUNC
     Primitive.cexpr! %q{ UINT2NUM(VM_FRAME_MAGIC_CFUNC) }
   end
@@ -826,6 +830,10 @@ module RubyVM::MJIT # :nodoc: all
     Primitive.cexpr! %q{ ULONG2NUM(SHAPE_MASK) }
   end
 
+  def C.rb_block_param_proxy
+    Primitive.cexpr! %q{ PTR2NUM(rb_block_param_proxy) }
+  end
+
   def C.rb_cFalseClass
     Primitive.cexpr! %q{ PTR2NUM(rb_cFalseClass) }
   end
@@ -1408,6 +1416,10 @@ module RubyVM::MJIT # :nodoc: all
       expandarray_postarg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), expandarray_postarg)")],
       expandarray_not_array: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), expandarray_not_array)")],
       expandarray_rhs_too_small: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), expandarray_rhs_too_small)")],
+      getblockpp_block_param_modified: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), getblockpp_block_param_modified)")],
+      getblockpp_block_handler_none: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), getblockpp_block_handler_none)")],
+      getblockpp_not_gc_guarded: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), getblockpp_not_gc_guarded)")],
+      getblockpp_not_iseq_block: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), getblockpp_not_iseq_block)")],
       compiled_block_count: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), compiled_block_count)")],
     )
   end
diff --git a/tool/mjit/bindgen.rb b/tool/mjit/bindgen.rb
index 105f878cdc..fbfd98960a 100755
--- a/tool/mjit/bindgen.rb
+++ b/tool/mjit/bindgen.rb
@@ -408,6 +408,7 @@ generator = BindingGenerator.new(
       VM_FRAME_MAGIC_CFUNC
       VM_FRAME_FLAG_CFRAME
       VM_FRAME_FLAG_CFRAME_KW
+      VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM
       VM_METHOD_TYPE_ISEQ
       VM_METHOD_TYPE_NOTIMPLEMENTED
       VM_METHOD_TYPE_CFUNC
@@ -447,6 +448,7 @@ generator = BindingGenerator.new(
       rb_cNilClass
       rb_cSymbol
       rb_cTrueClass
+      rb_block_param_proxy
     ],
   },
   types: %w[