ZJIT: Pass a &blk argument through to a direct C call

A C method takes its block out of the frame's specval exactly like an ISEQ
method does, and ZJIT already writes a literal block's handler there for
`CCallWithFrame` and `CCallVariadic`. So the reduction the previous commits
added for `&blk` -- a Proc is its own block handler, and the block param proxy
resolves to this frame's -- works for a C callee too; only
`unspecializable_c_call_type` stood in the way.

The one difference is the stack: the ISEQ path takes the block argument out of
the caller's stack because the callee's parameters are laid out there, while a
C frame is simply pushed over the argument slots. So the C path keeps the
original frame state and only tells `gen_push_frame` that one more slot is
consumed. The inline bodies and the leaf `CCall` fast path still opt out, as
they do for a literal block: neither pushes a frame to carry the handler.

On lobsters this converts `Hash#fetch(k, &blk)`, `Array#bsearch_index(&blk)`
and friends:

    dynamic_send_count           10,527,955 -> 10,380,293
    one_or_more_complex_arg_pass  1,423,400 ->  1,018,033
    send_block_arg_not_nil           57,695 ->    319,890

The block-arg counter grows because a C site whose block argument cannot be
reduced now reports why it stayed dynamic instead of being lumped in with the
complex-argument sites; those calls were already dynamic before.

Author: k0kubun

zjit/src/codegen.rs

@@ -785,12 +785,14 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::PatchPoint { invariant, state } => no_output!(gen_patch_point(jit, asm, function, invariant, &function.frame_state(*state))),
         Insn::CCall { cfunc, recv, args, name, owner, return_type: _, elidable: _ } => gen_ccall(asm, *cfunc, *name, *owner, opnd!(recv), opnds!(args)),
         Insn::CCallWithFrame(insn) => {
-            let CCallWithFrameData { cfunc, recv, name, args, cme, state, block, .. } = &**insn;
-            gen_ccall_with_frame(jit, asm, function, *cfunc, *name, opnd!(recv), opnds!(args), *cme, *block, &function.frame_state(*state))
+            let CCallWithFrameData { cfunc, recv, name, args, cme, state, block, block_arg, .. } = &**insn;
+            let block_arg = block_arg.map(|block_arg| opnd!(block_arg));
+            gen_ccall_with_frame(jit, asm, function, *cfunc, *name, opnd!(recv), opnds!(args), *cme, *block, block_arg, &function.frame_state(*state))
         }
         Insn::CCallVariadic(insn) => {
-            let CCallVariadicData { cfunc, recv, name, args, cme, state, block, .. } = &**insn;
-            gen_ccall_variadic(jit, asm, function, *cfunc, *name, opnd!(recv), opnds!(args), *cme, *block, &function.frame_state(*state))
+            let CCallVariadicData { cfunc, recv, name, args, cme, state, block, block_arg, .. } = &**insn;
+            let block_arg = block_arg.map(|block_arg| opnd!(block_arg));
+            gen_ccall_variadic(jit, asm, function, *cfunc, *name, opnd!(recv), opnds!(args), *cme, *block, block_arg, &function.frame_state(*state))
         }
         Insn::GetIvar { self_val, id, ic, state } => gen_getivar(asm, opnd!(self_val), *id, *ic, &function.frame_state(*state)),
         Insn::SetGlobal { id, val, state } => no_output!(gen_setglobal(jit, asm, function, *id, opnd!(val), &function.frame_state(*state))),
@@ -1137,13 +1139,16 @@ fn gen_ccall_with_frame(
     args: Vec<Opnd>,
     cme: *const rb_callable_method_entry_t,
     block: Option<BlockHandler>,
+    block_arg: Option<lir::Opnd>,
     state: &FrameState,
 ) -> lir::Opnd {
     gen_incr_counter(asm, Counter::non_variadic_cfunc_optimized_send_count);
     gen_stack_overflow_check(jit, asm, function, state, state.stack_size());
 
-    let args_with_recv_len = args.len() + 1;
-    let caller_stack_size = state.stack().len() - args_with_recv_len;
+    // A `&blk` argument keeps its VM stack slot, above the arguments and below the frame this
+    // pushes, so the frame setup consumes one more slot than the C function has arguments.
+    let stack_slots_consumed = args.len() + 1 + usize::from(block_arg.is_some());
+    let caller_stack_size = state.stack().len() - stack_slots_consumed;
 
     // Can't use gen_prepare_non_leaf_call() because we need to adjust the SP
     // to account for the receiver and arguments (and block arguments if any)
@@ -1160,10 +1165,10 @@ fn gen_ccall_with_frame(
         let cfp_self_addr = asm.lea(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SELF));
         asm.or(cfp_self_addr, Opnd::Imm(1))
     } else {
-        VM_BLOCK_HANDLER_NONE.into()
+        block_arg.unwrap_or_else(|| VM_BLOCK_HANDLER_NONE.into())
     };
 
-    gen_push_frame(asm, args_with_recv_len, state, ControlFrame {
+    gen_push_frame(asm, stack_slots_consumed, state, ControlFrame {
         recv,
         iseq: None,
         cme,
@@ -1230,16 +1235,16 @@ fn gen_ccall_variadic(
     args: Vec<Opnd>,
     cme: *const rb_callable_method_entry_t,
     block: Option<BlockHandler>,
+    block_arg: Option<lir::Opnd>,
     state: &FrameState,
 ) -> lir::Opnd {
     gen_incr_counter(asm, Counter::variadic_cfunc_optimized_send_count);
     gen_stack_overflow_check(jit, asm, function, state, state.stack_size());
 
-    let args_with_recv_len = args.len() + 1;
-
-    // Compute the caller's stack size after consuming recv and args.
-    // state.stack() includes recv + args, so subtract both.
-    let caller_stack_size = state.stack_size() - args_with_recv_len;
+    // Compute the caller's stack size after consuming recv, args and, when the site passes one,
+    // the `&blk` argument's slot. state.stack() includes all of them.
+    let stack_slots_consumed = args.len() + 1 + usize::from(block_arg.is_some());
+    let caller_stack_size = state.stack_size() - stack_slots_consumed;
 
     // Can't use gen_prepare_non_leaf_call() because we need to adjust the SP
     // to account for the receiver and arguments (like gen_ccall_with_frame does)
@@ -1251,10 +1256,10 @@ fn gen_ccall_variadic(
     let block_handler_specval = if let Some(BlockHandler::BlockIseq(blockiseq)) = block {
         gen_block_handler_specval(asm, blockiseq)
     } else {
-        VM_BLOCK_HANDLER_NONE.into()
+        block_arg.unwrap_or_else(|| VM_BLOCK_HANDLER_NONE.into())
     };
 
-    gen_push_frame(asm, args_with_recv_len, state, ControlFrame {
+    gen_push_frame(asm, stack_slots_consumed, state, ControlFrame {
         recv,
         iseq: None,
         cme,

zjit/src/codegen_tests.rs

@@ -2465,6 +2465,28 @@ fn test_send_block_param_proxy_from_block_body() {
     "), @"[[199, 202], [9, nil]]");
 }
 
+#[test]
+fn test_send_forwards_block_arg_to_cfunc() {
+    // A C method takes its block from the frame's specval too, so `&blk` forwarding reaches
+    // `Hash#fetch` (variadic) and `Array#bsearch_index` (fixed arity) as a direct C call. The
+    // no-block calls must still see no block.
+    assert_snapshot!(inspect("
+        def fetch(h, k, &b) = h.fetch(k, &b)
+        def search(a, &b) = a.bsearch_index(&b)
+        def each_with_proc(a, p) = a.each(&p)
+        out = []
+        300.times do
+          out << fetch({ a: 1 }, :b) { |k| \"no #{k}\" }
+          out << fetch({ a: 1 }, :a) { |k| \"no #{k}\" }
+          out << search([1, 3, 5, 7]) { |x| x >= 5 }
+          out << (fetch({ a: 1 }, :b) rescue :keyerror)
+        end
+        seen = []
+        each_with_proc([1, 2], ->(x) { seen << x })
+        [out.last(4), seen]
+    "), @r#"[["no b", 1, 2, :keyerror], [1, 2]]"#);
+}
+
 #[test]
 fn test_send_symbol_block_arg() {
     assert_snapshot!(inspect("

zjit/src/hir.rs

@@ -1018,6 +1018,10 @@ pub struct CCallWithFrameData {
     pub return_type: Type,
     pub elidable: bool,
     pub block: Option<BlockHandler>,
+    /// See [`SendDirectData::block_arg`]. Unlike the ISEQ case this is *not* taken off the VM
+    /// stack: the frame the C method runs in is pushed over the argument slots, so the frame
+    /// setup counts the block argument's slot even though the C function never sees it.
+    pub block_arg: Option<InsnId>,
 }
 
 /// Payload of [`Insn::SendDirect`]. Boxed in the enum to keep `Insn` small.
@@ -1051,6 +1055,8 @@ pub struct CCallVariadicData {
     pub return_type: Type,
     pub elidable: bool,
     pub block: Option<BlockHandler>,
+    /// See [`CCallWithFrameData::block_arg`].
+    pub block_arg: Option<InsnId>,
 }
 
 /// An instruction in the SSA IR. The output of an instruction is referred to by the index of
@@ -1700,11 +1706,13 @@ macro_rules! for_each_operand_impl {
             Insn::CCallWithFrame(insn) => {
                 $visit_one!(insn.recv);
                 $visit_many!(insn.args);
+                $visit_many!(insn.block_arg);
                 $visit_one!(insn.state);
             }
             Insn::CCallVariadic(insn) => {
                 $visit_one!(insn.recv);
                 $visit_many!(insn.args);
+                $visit_many!(insn.block_arg);
                 $visit_one!(insn.state);
             }
             Insn::InvokeBlock { args, state, .. } => {
@@ -2483,7 +2491,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 Ok(())
             },
             Insn::CCallWithFrame(insn) => {
-                let CCallWithFrameData { cfunc, recv, args, name, cme, block, .. } = &**insn;
+                let CCallWithFrameData { cfunc, recv, args, name, cme, block, block_arg, .. } = &**insn;
                 write!(f, "CCallWithFrame {recv}, :{}@{:p}", qualified_method_name(unsafe { (**cme).owner }, *name), self.ptr_map.map_ptr(*cfunc))?;
                 write_separated!(f, ", ", ", ", args);
                 match block {
@@ -2493,12 +2501,18 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                         write!(f, ", block=&block")?,
                     None => {}
                 }
+                if let Some(block_arg) = block_arg {
+                    write!(f, ", block=&{block_arg}")?;
+                }
                 Ok(())
             },
             Insn::CCallVariadic(insn) => {
-                let CCallVariadicData { cfunc, recv, args, name, cme, .. } = &**insn;
+                let CCallVariadicData { cfunc, recv, args, name, cme, block_arg, .. } = &**insn;
                 write!(f, "CCallVariadic {recv}, :{}@{:p}", qualified_method_name(unsafe { (**cme).owner }, *name), self.ptr_map.map_ptr(*cfunc))?;
                 write_separated!(f, ", ", ", ", args);
+                if let Some(block_arg) = block_arg {
+                    write!(f, ", block=&{block_arg}")?;
+                }
                 Ok(())
             },
             Insn::IncrCounterPtr { .. } => write!(f, "IncrCounterPtr"),
@@ -5883,7 +5897,12 @@ impl Function {
                         }
                         let mut stripped_block_arg = false;
                         let mut send_block_arg = None;
-                        if send_block == Some(BlockHandler::BlockArg) && def_type == VM_METHOD_TYPE_ISEQ {
+                        // A C method's frame carries the block handler in its specval just like an
+                        // ISEQ frame's, so the same reduction applies; the difference is that the
+                        // block argument keeps its VM stack slot, which the C frame setup accounts
+                        // for. Nothing else reads `args` positionally for a C call.
+                        if send_block == Some(BlockHandler::BlockArg)
+                            && matches!(def_type, VM_METHOD_TYPE_ISEQ | VM_METHOD_TYPE_CFUNC) {
                             // The block arg is the last element in args
                             if let Some(&block_arg) = args.last() {
                                 let statically_nil = self.is_a(block_arg, types::NilClass);
@@ -6245,12 +6264,13 @@ impl Function {
                                 cme: *const rb_callable_method_entry_struct,
                                 method_id: ID,
                                 argc: u32,
+                                // The call site's flags with `VM_CALL_ARGS_BLOCKARG` cleared when
+                                // `block_arg` already holds the handler the interpreter would have
+                                // built from it.
+                                ci_flags: u32,
+                                block_arg: Option<InsnId>,
                             ) -> Result<(), ()> {
-                                let call_info = unsafe { (*cd).ci };
-
-                                let ci_flags = unsafe { vm_ci_flag(call_info) };
-                                // When seeing &block argument, fall back to dynamic dispatch for now
-                                // TODO: Support block forwarding
+                                // Argument shapes the C frame setup cannot reproduce.
                                 if unspecializable_c_call_type(ci_flags) {
                                     // Only count features NOT already counted in type_specialize.
                                     if !unspecializable_call_type(ci_flags) {
@@ -6265,6 +6285,10 @@ impl Function {
                                     Some(BlockHandler::BlockIseq(blockiseq)) => Some(blockiseq),
                                     None => None,
                                 };
+                                // A block reaches the callee either way, so neither the inline
+                                // bodies nor the leaf fast path (which push no frame to carry the
+                                // handler) can serve this call.
+                                let passes_block = blockiseq.is_some() || block_arg.is_some();
 
                                 let cfunc = unsafe { get_cme_def_body_cfunc(cme) };
                                 // Find the `argc` (arity) of the C method, which describes the parameters it expects
@@ -6279,10 +6303,8 @@ impl Function {
                                 }
                                 let props = props.unwrap_or_default();
                                 let return_type = props.return_type;
-                                let elidable = match blockiseq {
-                                    Some(_) => false, // Don't consider cfuncs with block arguments as elidable for now
-                                    None => props.elidable,
-                                };
+                                // Don't consider cfuncs with block arguments as elidable for now
+                                let elidable = !passes_block && props.elidable;
 
                                 match cfunc_argc {
                                     0.. => {
@@ -6309,7 +6331,7 @@ impl Function {
                                         }
 
                                         // Try inlining the cfunc into HIR. Only inline if we don't have a block argument
-                                        if blockiseq.is_none() {
+                                        if !passes_block {
                                             let tmp_block = fun.new_block(u32::MAX);
                                             if let Some(replacement) = (props.inline)(fun, tmp_block, recv, &args, state) {
                                                 // Copy contents of tmp_block to block
@@ -6352,6 +6374,7 @@ impl Function {
                                             return_type,
                                             elidable,
                                             block: blockiseq.map(BlockHandler::BlockIseq),
+                                            block_arg,
                                         })));
                                         fun.insn_types[ccall] = fun.infer_type(ccall);
                                         fun.make_equal_to(send_insn_id, ccall);
@@ -6376,7 +6399,7 @@ impl Function {
                                         }
 
                                         // Try inlining the cfunc into HIR. Only inline if we don't have a block argument
-                                        if blockiseq.is_none() {
+                                        if !passes_block {
                                             let tmp_block = fun.new_block(u32::MAX);
                                             if let Some(replacement) = (props.inline)(fun, tmp_block, recv, &args, state) {
                                                 // Copy contents of tmp_block to block
@@ -6419,6 +6442,7 @@ impl Function {
                                             return_type,
                                             elidable,
                                             block: blockiseq.map(BlockHandler::BlockIseq),
+                                            block_arg,
                                         })));
                                         fun.insn_types[ccall] = fun.infer_type(ccall);
                                         fun.make_equal_to(send_insn_id, ccall);
@@ -6434,7 +6458,7 @@ impl Function {
                             }
 
                             let ccall_argc = if send_mid_override.is_some() { args.len() as u32 } else { unsafe { vm_ci_argc(ci) } };
-                            if reduce_send_to_ccall(self, block, insn_id, recv, cd, send_block, args, state, klass, profiled_type, cme, mid, ccall_argc).is_ok() {
+                            if reduce_send_to_ccall(self, block, insn_id, recv, cd, send_block, args, state, klass, profiled_type, cme, mid, ccall_argc, flags_for_check, send_block_arg).is_ok() {
                                 continue;
                             }
 
@@ -6694,6 +6718,7 @@ impl Function {
                                             return_type,
                                             elidable,
                                             block: None,
+                                            block_arg: None,
                                         })))
                                     };
                                     self.make_equal_to(insn_id, ccall);
@@ -6743,6 +6768,7 @@ impl Function {
                                             return_type,
                                             elidable,
                                             block: None,
+                                            block_arg: None,
                                         })))
                                     };
                                     self.make_equal_to(insn_id, ccall);
@@ -9135,13 +9161,19 @@ impl Function {
                 for &arg in &insn.args {
                     self.assert_subtype(insn_id, arg, types::BasicObject)?;
                 }
+                if let Some(block_arg) = insn.block_arg {
+                    self.assert_subtype(insn_id, block_arg, types::BasicObject)?;
+                }
                 Ok(())
             }
             Insn::CCallVariadic(ref insn) => {
                 self.assert_subtype(insn_id, insn.recv, types::BasicObject)?;
                 for &arg in &insn.args {
                     self.assert_subtype(insn_id, arg, types::BasicObject)?;
                 }
+                if let Some(block_arg) = insn.block_arg {
+                    self.assert_subtype(insn_id, block_arg, types::BasicObject)?;
+                }
                 Ok(())
             }
             Insn::ArrayPackBuffer { ref elements, fmt, buffer, .. } => {
@@ -10132,9 +10164,10 @@ struct AddIseqResult {
     profiles: ProfileOracle,
 }
 
-/// Whether any receiver class this site profiled resolves the call to an ISEQ method, which is
-/// the only method type whose frame setup `type_specialize` can hand a `&blk` block handler to.
-fn profiled_recv_has_iseq_callee(
+/// Whether any receiver class this site profiled resolves the call to a method whose frame setup
+/// `type_specialize` can hand a `&blk` block handler to. Only ISEQ and C methods get such a
+/// frame; the rest keep the dynamic send whatever the block argument is.
+fn profiled_recv_takes_block_handler(
     fun: &Function,
     profiles: &ProfileOracle,
     recv: InsnId,
@@ -10151,7 +10184,7 @@ fn profiled_recv_has_iseq_callee(
             cme = unsafe { rb_aliased_callable_method_entry(cme) };
             def_type = unsafe { get_cme_def_type(cme) };
         }
-        def_type == VM_METHOD_TYPE_ISEQ
+        matches!(def_type, VM_METHOD_TYPE_ISEQ | VM_METHOD_TYPE_CFUNC)
     })
 }
 
@@ -11699,10 +11732,9 @@ fn add_iseq_to_hir(
                         && args.last().is_some_and(|arg| block_param_proxy_values.contains(arg))
                         && block_arg_summary.as_ref().is_some_and(|summary| summary.buckets().iter().any(|profiled_type|
                             !profiled_type.is_empty() && profiled_type.class() == proxy_class))
-                        // Only an ISEQ callee's frame setup takes the handler; a C method reads
-                        // its block from a frame ZJIT does not build for a `&blk` argument, so
-                        // its call stays dynamic and the branch would be dead weight.
-                        && profiled_recv_has_iseq_callee(fun, &profiles, recv, exit_id, cd);
+                        // Only an ISEQ or C callee's frame setup takes the handler; for anything
+                        // else the call stays dynamic and the branch would be dead weight.
+                        && profiled_recv_takes_block_handler(fun, &profiles, recv, exit_id, cd);
                     let proxy_join = if proxy_split {
                         let block_arg_insn = *args.last().unwrap();
                         let join_block = fun.new_block(insn_idx);