Implement proxy blockarg

3 files changed, 21 insertions, 8 deletions

diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index 4598bef697..4378a1e9a3 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -2903,15 +2903,20 @@ module RubyVM::MJIT
     def jit_caller_setup_arg_block(jit, ctx, asm, ci, blockiseq, is_super)
       side_exit = side_exit(jit, ctx)
       if C.vm_ci_flag(ci) & C.VM_CALL_ARGS_BLOCKARG != 0
-        # TODO: Skip the check using Context?
+        # TODO: Skip cmp + jne using Context?
         block_code = jit.peek_at_stack(0)
         block_opnd = ctx.stack_opnd(0) # to be popped after eliminating side exit possibility
         if block_code.nil?
           asm.cmp(block_opnd, Qnil)
           asm.jne(counted_exit(side_exit, :send_block_not_nil))
           return C.VM_BLOCK_HANDLER_NONE
+        elsif C.to_value(block_code) == C.rb_block_param_proxy
+          asm.mov(:rax, C.rb_block_param_proxy)
+          asm.cmp(block_opnd, :rax)
+          asm.jne(counted_exit(side_exit, :send_block_not_proxy))
+          return C.rb_block_param_proxy
         else
-          asm.incr_counter(:send_blockarg)
+          asm.incr_counter(:send_blockarg_not_nil_or_proxy)
           return CantCompile
         end
       elsif blockiseq != 0
@@ -3011,10 +3016,6 @@ module RubyVM::MJIT
         asm.incr_counter(:send_kw_splat)
         return CantCompile
       end
-      if flags & C.VM_CALL_ARGS_BLOCKARG != 0
-        asm.incr_counter(:send_blockarg)
-        return CantCompile
-      end
 
       # Ensure we haven't rebound this method onto an incompatible class.
       # In the interpreter we try to avoid making this check by performing some
@@ -3519,6 +3520,16 @@ module RubyVM::MJIT
       # ep[-1]: block handler or prev env ptr
       if block_handler == C.VM_BLOCK_HANDLER_NONE
         asm.mov([SP, C.VALUE.size * (ep_offset - 1)], C.VM_BLOCK_HANDLER_NONE)
+      elsif block_handler == C.rb_block_param_proxy
+        # vm_caller_setup_arg_block:
+        #   VALUE handler = VM_CF_BLOCK_HANDLER(reg_cfp);
+        #   reg_cfp->block_code = (const void *) handler;
+        jit_get_lep(jit, asm, reg: :rax)
+        asm.mov(:rax, [:rax, C.VALUE.size * C.VM_ENV_DATA_INDEX_SPECVAL]) # handler
+        asm.mov([CFP, C.rb_control_frame_t.offsetof(:block_code)], :rax)
+
+        asm.mov(:rax, C.rb_block_param_proxy)
+        asm.mov([SP, C.VALUE.size * (ep_offset - 1)], :rax)
       else # assume blockiseq
         asm.mov(:rax, block_handler)
         asm.mov([CFP, C.rb_control_frame_t.offsetof(:block_code)], :rax)
diff --git a/mjit_c.h b/mjit_c.h
index e14f73a223..40336ca0ca 100644
--- a/mjit_c.h
+++ b/mjit_c.h
@@ -133,11 +133,12 @@ MJIT_RUNTIME_COUNTERS(
     send_arity,
     send_c_tracing,
 
-    send_blockarg,
+    send_blockarg_not_nil_or_proxy,
     send_blockiseq,
     send_block_handler,
     send_block_setup,
     send_block_not_nil,
+    send_block_not_proxy,
 
     send_iseq_not_only_optparam,
     send_iseq_kw_splat,
diff --git a/mjit_c.rb b/mjit_c.rb
index a50c3ebe0c..6dca2e2441 100644
--- a/mjit_c.rb
+++ b/mjit_c.rb
@@ -1442,11 +1442,12 @@ module RubyVM::MJIT # :nodoc: all
       send_stackoverflow: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_stackoverflow)")],
       send_arity: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_arity)")],
       send_c_tracing: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_c_tracing)")],
-      send_blockarg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_blockarg)")],
+      send_blockarg_not_nil_or_proxy: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_blockarg_not_nil_or_proxy)")],
       send_blockiseq: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_blockiseq)")],
       send_block_handler: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_block_handler)")],
       send_block_setup: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_block_setup)")],
       send_block_not_nil: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_block_not_nil)")],
+      send_block_not_proxy: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_block_not_proxy)")],
       send_iseq_not_only_optparam: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_iseq_not_only_optparam)")],
       send_iseq_kw_splat: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_iseq_kw_splat)")],
       send_cfunc_variadic: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_cfunc_variadic)")],