YJIT: Support nil and blockparamproxy as blockarg in send (#6492)

Co-authored-by: John Hawthorn <john@hawthorn.email>

Co-authored-by: John Hawthorn <john@hawthorn.email>

7 files changed, 221 insertions, 47 deletions

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 2dc460c628..7361d2a725 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -3337,3 +3337,46 @@ assert_equal '[[1, nil, 2]]', %q{
 
   5.times.map { opt_and_kwargs(1, c: 2) }.uniq
 }
+
+# bmethod with forwarded block
+assert_equal '2', %q{
+  define_method(:foo) do |&block|
+    block.call
+  end
+
+  def bar(&block)
+    foo(&block)
+  end
+
+  bar { 1 }
+  bar { 2 }
+}
+
+# bmethod with forwarded block and arguments
+assert_equal '5', %q{
+  define_method(:foo) do |n, &block|
+    n + block.call
+  end
+
+  def bar(n, &block)
+    foo(n, &block)
+  end
+
+  bar(0) { 1 }
+  bar(3) { 2 }
+}
+
+# bmethod with forwarded unwanted block
+assert_equal '1', %q{
+  one = 1
+  define_method(:foo) do
+    one
+  end
+
+  def bar(&block)
+    foo(&block)
+  end
+
+  bar { }
+  bar { }
+}
diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index 09b5989a06..d740870736 100644
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@@ -513,9 +513,8 @@ class TestYJIT < Test::Unit::TestCase
     RUBY
   end
 
-  def test_getblockparamproxy_with_no_block
-    # Currently side exits on the send
-    assert_compiles(<<~'RUBY', insns: [:getblockparamproxy], exits: { send: 2 })
+  def test_send_blockarg
+    assert_compiles(<<~'RUBY', insns: [:getblockparamproxy, :send], exits: {})
       def bar
       end
 
@@ -526,6 +525,9 @@ class TestYJIT < Test::Unit::TestCase
 
       foo
       foo
+
+      foo { }
+      foo { }
     RUBY
   end
 
diff --git a/yjit.c b/yjit.c
index 7e6fc9e3fb..c53444d5a3 100644
--- a/yjit.c
+++ b/yjit.c
@@ -592,6 +592,12 @@ rb_get_iseq_body_local_iseq(const rb_iseq_t *iseq)
     return iseq->body->local_iseq;
 }
 
+const rb_iseq_t *
+rb_get_iseq_body_parent_iseq(const rb_iseq_t *iseq)
+{
+    return iseq->body->parent_iseq;
+}
+
 unsigned int
 rb_get_iseq_body_local_table_size(const rb_iseq_t *iseq)
 {
diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index 60a9d1b87d..21aaec84cb 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -350,6 +350,7 @@ fn main() {
         .allowlist_function("rb_get_def_bmethod_proc")
         .allowlist_function("rb_iseq_encoded_size")
         .allowlist_function("rb_get_iseq_body_local_iseq")
+        .allowlist_function("rb_get_iseq_body_parent_iseq")
         .allowlist_function("rb_get_iseq_body_iseq_encoded")
         .allowlist_function("rb_get_iseq_body_stack_max")
         .allowlist_function("rb_get_iseq_flags_has_opt")
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index d6ea8996e1..51645f0744 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -1542,6 +1542,22 @@ fn gen_get_ep(asm: &mut Assembler, level: u32) -> Opnd {
     ep_opnd
 }
 
+// Gets the EP of the ISeq of the containing method, or "local level".
+// Equivalent of GET_LEP() macro.
+fn gen_get_lep(jit: &mut JITState, asm: &mut Assembler) -> Opnd {
+    // Equivalent of get_lvar_level() in compile.c
+    fn get_lvar_level(iseq: IseqPtr) -> u32 {
+        if iseq == unsafe { rb_get_iseq_body_local_iseq(iseq) } {
+            0
+        } else {
+            1 + get_lvar_level(unsafe { rb_get_iseq_body_parent_iseq(iseq) })
+        }
+    }
+
+    let level = get_lvar_level(jit.get_iseq());
+    gen_get_ep(asm, level)
+}
+
 fn gen_getlocal_generic(
     ctx: &mut Context,
     asm: &mut Assembler,
@@ -3995,9 +4011,14 @@ unsafe extern "C" fn build_kwhash(ci: *const rb_callinfo, sp: *const VALUE) -> V
     hash
 }
 
-enum BlockHandler {
+// SpecVal is a single value in an iseq invocation's environment on the stack,
+// at sp[-2]. Depending on the frame type, it can serve different purposes,
+// which are covered here by enum variants.
+enum SpecVal {
     None,
-    CurrentFrame,
+    BlockISeq(IseqPtr),
+    BlockParamProxy,
+    PrevEP(*const VALUE),
 }
 
 struct ControlFrame {
@@ -4006,8 +4027,7 @@ struct ControlFrame {
     iseq: Option<IseqPtr>,
     pc: Option<u64>,
     frame_type: u32,
-    block_handler: BlockHandler,
-    prev_ep: Option<*const VALUE>,
+    specval: SpecVal,
     cme: *const rb_callable_method_entry_t,
     local_size: i32
 }
@@ -4028,7 +4048,7 @@ struct ControlFrame {
 //   * Stack overflow is not checked (should be done by the caller)
 //   * Interrupts are not checked (should be done by the caller)
 fn gen_push_frame(
-    _jit: &mut JITState,
+    jit: &mut JITState,
     _ctx: &mut Context,
     asm: &mut Assembler,
     set_sp_cfp: bool, // if true CFP and SP will be switched to the callee
@@ -4060,19 +4080,33 @@ fn gen_push_frame(
     // Write special value at sp[-2]. It's either a block handler or a pointer to
     // the outer environment depending on the frame type.
     // sp[-2] = specval;
-    let specval: Opnd = match (frame.prev_ep, frame.block_handler) {
-        (None, BlockHandler::None) => {
+    let specval: Opnd = match frame.specval {
+        SpecVal::None => {
             VM_BLOCK_HANDLER_NONE.into()
         }
-        (None, BlockHandler::CurrentFrame) => {
+        SpecVal::BlockISeq(block_iseq) => {
+            // Change cfp->block_code in the current frame. See vm_caller_setup_arg_block().
+            // VM_CFP_TO_CAPTURED_BLOCK does &cfp->self, rb_captured_block->code.iseq aliases
+            // with cfp->block_code.
+            asm.store(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_BLOCK_CODE), VALUE::from(block_iseq).into());
+
             let cfp_self = asm.lea(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SELF));
             asm.or(cfp_self, Opnd::Imm(1))
         }
-        (Some(prev_ep), BlockHandler::None) => {
+        SpecVal::BlockParamProxy => {
+            let ep_opnd = gen_get_lep(jit, asm);
+            let block_handler = asm.load(
+                Opnd::mem(64, ep_opnd, (SIZEOF_VALUE as i32) * (VM_ENV_DATA_INDEX_SPECVAL as i32))
+            );
+
+            asm.store(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_BLOCK_CODE), block_handler);
+
+            block_handler
+        }
+        SpecVal::PrevEP(prev_ep) => {
             let tagged_prev_ep = (prev_ep as usize) | 1;
             VALUE(tagged_prev_ep).into()
         }
-        (_, _) => panic!("specval can only be one of prev_ep or block_handler")
     };
     asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -2), specval);
 
@@ -4218,6 +4252,43 @@ fn gen_send_cfunc(
         return CantCompile;
     }
 
+    let block_arg = flags & VM_CALL_ARGS_BLOCKARG != 0;
+    let block_arg_type = if block_arg {
+        Some(ctx.get_opnd_type(StackOpnd(0)))
+    } else {
+        None
+    };
+
+    match block_arg_type {
+        Some(Type::Nil | Type::BlockParamProxy) => {
+            // We'll handle this later
+        }
+        None => {
+            // Nothing to do
+        }
+        _ => {
+            gen_counter_incr!(asm, send_block_arg);
+            return CantCompile;
+        }
+    }
+
+    match block_arg_type {
+        Some(Type::Nil) => {
+            // We have a nil block arg, so let's pop it off the args
+            ctx.stack_pop(1);
+        }
+        Some(Type::BlockParamProxy) => {
+            // We don't need the actual stack value
+            ctx.stack_pop(1);
+        }
+        None => {
+            // Nothing to do
+        }
+        _ => {
+            assert!(false);
+        }
+    }
+
     // This is a .send call and we need to adjust the stack
     if flags & VM_CALL_OPT_SEND != 0 {
         handle_opt_send_shift_stack(asm, argc as i32, ctx);
@@ -4229,21 +4300,16 @@ fn gen_send_cfunc(
     // Store incremented PC into current control frame in case callee raises.
     jit_save_pc(jit, asm);
 
-    if let Some(block_iseq) = block {
-        // Change cfp->block_code in the current frame. See vm_caller_setup_arg_block().
-        // VM_CFP_TO_CAPTURED_BLOCK does &cfp->self, rb_captured_block->code.iseq aliases
-        // with cfp->block_code.
-        asm.mov(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_BLOCK_CODE), VALUE::from(block_iseq).into());
-    }
-
     // Increment the stack pointer by 3 (in the callee)
     // sp += 3
     let sp = asm.lea(ctx.sp_opnd((SIZEOF_VALUE as isize) * 3));
 
-    let frame_block_handler = if let Some(_block_iseq) = block {
-        BlockHandler::CurrentFrame
+    let specval = if block_arg_type == Some(Type::BlockParamProxy) {
+        SpecVal::BlockParamProxy
+    } else if let Some(block_iseq) = block {
+        SpecVal::BlockISeq(block_iseq)
     } else {
-        BlockHandler::None
+        SpecVal::None
     };
 
     let mut frame_type = VM_FRAME_MAGIC_CFUNC | VM_FRAME_FLAG_CFRAME | VM_ENV_FLAG_LOCAL;
@@ -4253,11 +4319,10 @@ fn gen_send_cfunc(
 
     gen_push_frame(jit, ctx, asm, false, ControlFrame {
         frame_type,
-        block_handler: frame_block_handler,
+        specval,
         cme,
         recv,
         sp,
-        prev_ep: None,
         pc: Some(0),
         iseq: None,
         local_size: 0,
@@ -4606,6 +4671,26 @@ fn gen_send_iseq(
         return CantCompile;
     }
 
+    let block_arg = flags & VM_CALL_ARGS_BLOCKARG != 0;
+    let block_arg_type = if block_arg {
+        Some(ctx.get_opnd_type(StackOpnd(0)))
+    } else {
+        None
+    };
+
+    match block_arg_type {
+        Some(Type::Nil | Type::BlockParamProxy) => {
+            // We'll handle this later
+        }
+        None => {
+            // Nothing to do
+        }
+        _ => {
+            gen_counter_incr!(asm, send_block_arg);
+            return CantCompile;
+        }
+    }
+
     // If we have unfilled optional arguments and keyword arguments then we
     // would need to move adjust the arguments location to account for that.
     // For now we aren't handling this case.
@@ -4700,6 +4785,23 @@ fn gen_send_iseq(
     // Check for interrupts
     gen_check_ints(asm, side_exit);
 
+    match block_arg_type {
+        Some(Type::Nil) => {
+            // We have a nil block arg, so let's pop it off the args
+            ctx.stack_pop(1);
+        }
+        Some(Type::BlockParamProxy) => {
+            // We don't need the actual stack value
+            ctx.stack_pop(1);
+        }
+        None => {
+            // Nothing to do
+        }
+        _ => {
+            assert!(false);
+        }
+    }
+
     let leaf_builtin_raw = unsafe { rb_leaf_builtin_function(iseq) };
     let leaf_builtin: Option<*const rb_builtin_function> = if leaf_builtin_raw.is_null() {
         None
@@ -4914,32 +5016,30 @@ fn gen_send_iseq(
     // Store the next PC in the current frame
     jit_save_pc(jit, asm);
 
-    if let Some(block_val) = block {
-        // Change cfp->block_code in the current frame. See vm_caller_setup_arg_block().
-        // VM_CFP_TO_CAPTURED_BLCOK does &cfp->self, rb_captured_block->code.iseq aliases
-        // with cfp->block_code.
-        asm.store(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_BLOCK_CODE), VALUE(block_val as usize).into());
-    }
-
     // Adjust the callee's stack pointer
     let offs =
         (SIZEOF_VALUE as isize) * (3 + (num_locals as isize) + if doing_kw_call { 1 } else { 0 });
     let callee_sp = asm.lea(ctx.sp_opnd(offs));
 
-    let frame_block_handler = if let Some(_) = block {
-        BlockHandler::CurrentFrame
+    let specval = if let Some(prev_ep) = prev_ep {
+        // We've already side-exited if the callee expects a block, so we
+        // ignore any supplied block here
+        SpecVal::PrevEP(prev_ep)
+    } else if block_arg_type == Some(Type::BlockParamProxy) {
+        SpecVal::BlockParamProxy
+    } else if let Some(block_val) = block {
+        SpecVal::BlockISeq(block_val)
     } else {
-        BlockHandler::None
+        SpecVal::None
     };
 
     // Setup the new frame
     gen_push_frame(jit, ctx, asm, true, ControlFrame {
         frame_type,
-        block_handler: frame_block_handler,
+        specval,
         cme,
         recv,
         sp: callee_sp,
-        prev_ep,
         iseq: Some(iseq),
         pc: None, // We are calling into jitted code, which will set the PC as necessary
         local_size: num_locals
@@ -5146,26 +5246,23 @@ fn gen_send_general(
         return CantCompile;
     }
 
-    if flags & VM_CALL_ARGS_BLOCKARG != 0 {
-        gen_counter_incr!(asm, send_block_arg);
-        return CantCompile;
-    }
-
     // Defer compilation so we can specialize on class of receiver
     if !jit_at_current_insn(jit) {
         defer_compilation(jit, ctx, asm, ocb);
         return EndBlock;
     }
 
-    let comptime_recv = jit_peek_at_stack(jit, ctx, argc as isize);
+    let recv_idx = argc + if flags & VM_CALL_ARGS_BLOCKARG != 0 { 1 } else { 0 };
+
+    let comptime_recv = jit_peek_at_stack(jit, ctx, recv_idx as isize);
     let comptime_recv_klass = comptime_recv.class_of();
 
     // Guard that the receiver has the same class as the one from compile time
     let side_exit = get_side_exit(jit, ocb, ctx);
 
     // Points to the receiver operand on the stack
-    let recv = ctx.stack_opnd(argc);
-    let recv_opnd = StackOpnd(argc.try_into().unwrap());
+    let recv = ctx.stack_opnd(recv_idx);
+    let recv_opnd = StackOpnd(recv_idx.try_into().unwrap());
     jit_guard_known_klass(
         jit,
         ctx,
@@ -5273,6 +5370,11 @@ fn gen_send_general(
 
                 let ivar_name = unsafe { get_cme_def_body_attr_id(cme) };
 
+                if flags & VM_CALL_ARGS_BLOCKARG != 0 {
+                    gen_counter_incr!(asm, send_block_arg);
+                    return CantCompile;
+                }
+
                 return gen_get_ivar(
                     jit,
                     ctx,
@@ -5298,6 +5400,9 @@ fn gen_send_general(
                     // See :attr-tracing:
                     gen_counter_incr!(asm, send_cfunc_tracing);
                     return CantCompile;
+                } else if flags & VM_CALL_ARGS_BLOCKARG != 0 {
+                    gen_counter_incr!(asm, send_block_arg);
+                    return CantCompile;
                 } else {
                     let ivar_name = unsafe { get_cme_def_body_attr_id(cme) };
                     return gen_set_ivar(jit, ctx, asm, comptime_recv, ivar_name, flags, argc);
@@ -5326,6 +5431,10 @@ fn gen_send_general(
             }
             // Send family of methods, e.g. call/apply
             VM_METHOD_TYPE_OPTIMIZED => {
+                if flags & VM_CALL_ARGS_BLOCKARG != 0 {
+                    gen_counter_incr!(asm, send_block_arg);
+                    return CantCompile;
+                }
 
                 let opt_type = unsafe { get_cme_def_body_optimized_type(cme) };
                 match opt_type {
@@ -6191,7 +6300,7 @@ fn gen_getblockparamproxy(
         // Push rb_block_param_proxy. It's a root, so no need to use jit_mov_gc_ptr.
         assert!(!unsafe { rb_block_param_proxy }.special_const_p());
 
-        let top = ctx.stack_push(Type::Unknown);
+        let top = ctx.stack_push(Type::BlockParamProxy);
         asm.mov(top, Opnd::const_ptr(unsafe { rb_block_param_proxy }.as_ptr()));
     }
 
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 19272350ed..9288a0188b 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -44,6 +44,9 @@ pub enum Type {
 
     TString, // An object with the T_STRING flag set, possibly an rb_cString
     CString, // An un-subclassed string of type rb_cString (can have instance vars in some cases)
+
+    BlockParamProxy, // A special sentinel value indicating the block parameter should be read from
+                     // the current surrounding cfp
 }
 
 // Default initialization
@@ -79,6 +82,12 @@ impl Type {
             if val.class_of() == unsafe { rb_cString } {
                 return Type::CString;
             }
+            // We likewise can't reference rb_block_param_proxy, but it's again an optimisation;
+            // we can just treat it as a normal Object.
+            #[cfg(not(test))]
+            if val == unsafe { rb_block_param_proxy } {
+                return Type::BlockParamProxy;
+            }
             match val.builtin_type() {
                 RUBY_T_ARRAY => Type::Array,
                 RUBY_T_HASH => Type::Hash,
@@ -142,7 +151,8 @@ impl Type {
             Type::Hash => Some(RUBY_T_HASH),
             Type::ImmSymbol | Type::HeapSymbol => Some(RUBY_T_SYMBOL),
             Type::TString | Type::CString => Some(RUBY_T_STRING),
-            Type::Unknown | Type::UnknownImm | Type::UnknownHeap => None
+            Type::Unknown | Type::UnknownImm | Type::UnknownHeap => None,
+            Type::BlockParamProxy => None,
         }
     }
 
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index a9242e05ca..da7b332e07 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -1419,6 +1419,9 @@ extern "C" {
     pub fn rb_get_iseq_body_local_iseq(iseq: *const rb_iseq_t) -> *const rb_iseq_t;
 }
 extern "C" {
+    pub fn rb_get_iseq_body_parent_iseq(iseq: *const rb_iseq_t) -> *const rb_iseq_t;
+}
+extern "C" {
     pub fn rb_get_iseq_body_local_table_size(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;
 }
 extern "C" {