ZJIT: Optimize send with a nil block to SendWithoutBlockDirect

rwstauner · rwstauner · commit 2e421eceaf56 · 2026-03-31T16:19:04.000-07:00
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
@@ -506,6 +506,7 @@ pub enum SideExitReason {
     UnhandledYARVInsn(u32),
     UnhandledCallType(CallType),
     UnhandledBlockArg,
+    BlockArgNotNil,
     TooManyKeywordParameters,
     FixnumAddOverflow,
     FixnumSubOverflow,
@@ -664,6 +665,8 @@ pub enum SendFallbackReason {
     SendCfuncArrayVariadic,
     SendNotOptimizedMethodType(MethodType),
     SendNotOptimizedNeedPermission,
+    /// The block argument is not nil, so we can't optimize to SendWithoutBlockDirect
+    SendBlockArgNotNil,
     CCallWithFrameTooManyArgs,
     ObjToStringNotString,
     TooManyArgsForLir,
@@ -738,6 +741,7 @@ impl Display for SendFallbackReason {
             SendCfuncVariadic => write!(f, "Send: C function is variadic"),
             SendCfuncArrayVariadic => write!(f, "Send: C function expects array variadic"),
             SendNotOptimizedMethodType(method_type) => write!(f, "Send: unsupported method type {:?}", method_type),
+            SendBlockArgNotNil => write!(f, "Send: block argument is not nil"),
             CCallWithFrameTooManyArgs => write!(f, "CCallWithFrame: too many arguments"),
             ObjToStringNotString => write!(f, "ObjToString: result is not a string"),
             TooManyArgsForLir => write!(f, "Too many arguments for LIR"),
@@ -3704,9 +3708,38 @@ impl Function {
                             def_type = unsafe { get_cme_def_type(cme) };
                         }
 
+                        // Check if we can optimize `foo(&block)` where block is nil to a send without block
+                        let mut send_block = send_block;
+                        let mut args = args;
+                        let mut stripped_nil_block = false;
+                        if send_block == Some(BlockHandler::BlockArg) && def_type == VM_METHOD_TYPE_ISEQ {
+                            // The block arg is the last element in args
+                            if let Some(&block_arg) = args.last() {
+                                let block_arg_type = self.type_of(block_arg);
+                                if block_arg_type.is_subtype(types::NilClass) {
+                                    // Block arg is known to be nil - strip it and treat as no block
+                                    self.push_insn(block, Insn::GuardBitEquals {
+                                        val: block_arg,
+                                        expected: Const::Value(Qnil),
+                                        reason: SideExitReason::BlockArgNotNil,
+                                        state,
+                                    });
+                                    args = args[..args.len() - 1].to_vec();
+                                    send_block = None;
+                                    stripped_nil_block = true;
+                                } else {
+                                    // Can't prove block arg is nil
+                                    self.set_dynamic_send_reason(insn_id, SendBlockArgNotNil);
+                                    self.push_insn_id(block, insn_id); continue;
+                                }
+                            }
+                        }
+
                         // If the call site info indicates that the `Function` has overly complex arguments, then do not optimize into a `SendDirect`.
                         // Optimized methods(`VM_METHOD_TYPE_OPTIMIZED`) handle their own argument constraints (e.g., kw_splat for Proc call).
-                        if def_type != VM_METHOD_TYPE_OPTIMIZED && unspecializable_call_type(flags) {
+                        // Mask out ARGS_BLOCKARG only if we've already handled the nil block arg case above.
+                        let flags_for_check = if stripped_nil_block { flags & !VM_CALL_ARGS_BLOCKARG } else { flags };
+                        if def_type != VM_METHOD_TYPE_OPTIMIZED && unspecializable_call_type(flags_for_check) {
                             self.count_complex_call_features(block, flags);
                             self.set_dynamic_send_reason(insn_id, ComplexArgPass);
                             self.push_insn_id(block, insn_id); continue;
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
@@ -8019,7 +8019,7 @@ mod hir_opt_tests {
           v22:BasicObject = LoadField v18, :block@0x1010
           Jump bb6(v22, v22)
         bb6(v16:BasicObject, v17:BasicObject):
-          v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Complex argument passing
+          v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Send: block argument is not nil
           CheckInterrupts
           Return v29
         ");
@@ -8057,7 +8057,7 @@ mod hir_opt_tests {
           v22:BasicObject = LoadField v18, :block@0x1002
           Jump bb6(v22, v22)
         bb6(v16:BasicObject, v17:BasicObject):
-          v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Complex argument passing
+          v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Send: block argument is not nil
           CheckInterrupts
           Return v29
         ");
@@ -8096,7 +8096,7 @@ mod hir_opt_tests {
           v17:BasicObject = LoadField v13, :block@0x1010
           Jump bb6(v17)
         bb6(v12:BasicObject):
-          v24:BasicObject = Send v10, &block, :map, v12 # SendFallbackReason: Complex argument passing
+          v24:BasicObject = Send v10, &block, :map, v12 # SendFallbackReason: Send: block argument is not nil
           CheckInterrupts
           Return v24
         ");
@@ -8193,6 +8193,7 @@ mod hir_opt_tests {
         ");
     }
 
+
     #[test]
     fn test_optimize_send_with_non_nil_block_arg() {
         eval(r#"
@@ -8218,7 +8219,7 @@ mod hir_opt_tests {
           Jump bb3(v5, v6)
         bb3(v8:BasicObject, v9:NilClass):
           v13:StaticSymbol[:to_s] = Const Value(VALUE(0x1000))
-          v19:BasicObject = Send v8, &block, :foo, v13 # SendFallbackReason: Complex argument passing
+          v19:BasicObject = Send v8, &block, :foo, v13 # SendFallbackReason: Send: block argument is not nil
           CheckInterrupts
           Return v19
         ");
@@ -8249,9 +8250,12 @@ mod hir_opt_tests {
           Jump bb3(v5, v6)
         bb3(v8:BasicObject, v9:NilClass):
           v13:NilClass = Const Value(nil)
-          v19:BasicObject = Send v8, &block, :foo, v13 # SendFallbackReason: Complex argument passing
+          v25:NilClass = GuardBitEquals v13, Value(nil)
+          PatchPoint MethodRedefined(NilClass@0x1000, foo@0x1008, cme:0x1010)
+          v27:NilClass = GuardType v8, NilClass
+          v29:Fixnum[42] = Const Value(42)
           CheckInterrupts
-          Return v19
+          Return v29
         ");
     }
 
@@ -11602,7 +11606,7 @@ mod hir_opt_tests {
           Jump bb3(v4)
         bb3(v6:BasicObject):
           v11:StaticSymbol[:the_block] = Const Value(VALUE(0x1000))
-          v13:BasicObject = Send v6, &block, :callee, v11 # SendFallbackReason: Complex argument passing
+          v13:BasicObject = Send v6, &block, :callee, v11 # SendFallbackReason: Send: block argument is not nil
           CheckInterrupts
           Return v13
         ");
@@ -11645,7 +11649,7 @@ mod hir_opt_tests {
           v22:BasicObject = LoadField v18, :block@0x1010
           Jump bb6(v22, v22)
         bb6(v16:BasicObject, v17:BasicObject):
-          v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Complex argument passing
+          v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Send: block argument is not nil
           CheckInterrupts
           Return v29
         ");
diff --git a/zjit/src/profile.rs b/zjit/src/profile.rs
@@ -97,6 +97,8 @@ fn profile_insn(bare_opcode: ruby_vminsn_type, ec: EcPtr) {
             let cd: *const rb_call_data = profiler.insn_opnd(0).as_ptr();
             let argc = unsafe { vm_ci_argc((*cd).ci) };
             // Profile all the arguments and self (+1).
+            // Note: For sends with ARGS_BLOCKARG, the block arg is NOT profiled here.
+            // The block arg type is determined from HIR static type inference.
             profile_operands(profiler, profile, (argc + 1) as usize);
         }
         YARVINSN_splatkw => profile_operands(profiler, profile, 2),
diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs
@@ -193,6 +193,7 @@ make_counters! {
         exit_unhandled_splat,
         exit_unhandled_kwarg,
         exit_unhandled_block_arg,
+        exit_block_arg_not_nil,
         exit_unknown_special_variable,
         exit_unhandled_hir_insn,
         exit_unhandled_yarv_insn,
@@ -263,6 +264,7 @@ make_counters! {
         send_fallback_send_no_profiles,
         send_fallback_send_not_optimized_method_type,
         send_fallback_send_not_optimized_need_permission,
+        send_fallback_send_block_arg_not_nil,
         send_fallback_ccall_with_frame_too_many_args,
         send_fallback_argc_param_mismatch,
         // The call has at least one feature on the caller or callee side
@@ -587,6 +589,7 @@ pub fn side_exit_counter(reason: crate::hir::SideExitReason) -> Counter {
         UnhandledHIRUnknown(_)        => exit_unhandled_hir_insn,
         UnhandledYARVInsn(_)          => exit_unhandled_yarv_insn,
         UnhandledBlockArg             => exit_unhandled_block_arg,
+        BlockArgNotNil                => exit_block_arg_not_nil,
         FixnumAddOverflow             => exit_fixnum_add_overflow,
         FixnumSubOverflow             => exit_fixnum_sub_overflow,
         FixnumMultOverflow            => exit_fixnum_mult_overflow,
@@ -676,6 +679,7 @@ pub fn send_fallback_counter(reason: crate::hir::SendFallbackReason) -> Counter
         BmethodNonIseqProc                        => send_fallback_bmethod_non_iseq_proc,
         SendNotOptimizedMethodType(_)             => send_fallback_send_not_optimized_method_type,
         SendNotOptimizedNeedPermission            => send_fallback_send_not_optimized_need_permission,
+        SendBlockArgNotNil                        => send_fallback_send_block_arg_not_nil,
         CCallWithFrameTooManyArgs                 => send_fallback_ccall_with_frame_too_many_args,
         ObjToStringNotString                      => send_fallback_obj_to_string_not_string,
         SuperCallWithBlock                        => send_fallback_super_call_with_block,

Original file line number	Diff line number	Diff line change
`@@ -97,6 +97,8 @@ fn profile_insn(bare_opcode: ruby_vminsn_type, ec: EcPtr) {`
`97`	`97`	`let cd: *const rb_call_data = profiler.insn_opnd(0).as_ptr();`
`98`	`98`	`let argc = unsafe { vm_ci_argc((*cd).ci) };`
`99`	`99`	`// Profile all the arguments and self (+1).`
	`100`	`+ // Note: For sends with ARGS_BLOCKARG, the block arg is NOT profiled here.`
	`101`	`+ // The block arg type is determined from HIR static type inference.`
`100`	`102`	`profile_operands(profiler, profile, (argc + 1) as usize);`
`101`	`103`	`}`
`102`	`104`	`YARVINSN_splatkw => profile_operands(profiler, profile, 2),`