ZJIT: Optimize send with a nil block to SendWithoutBlockDirect

rwstauner · rwstauner · commit f4fe2ff0329e · 2026-07-03T15:57:04.000-07:00
A call like `foo(&amp;block)` sets VM_CALL_ARGS_BLOCKARG, which previously
forced a dynamic dispatch even when the block argument is nil and the
callee is an ISEQ method that could otherwise be reduced to a
SendDirect.

Passing a nil block argument is equivalent to calling without a block,
so when we can prove the block argument is nil we strip it and treat the
call as blockless. This happens when the block argument is either:

  * statically known to be nil, or
  * profiled as monomorphically nil, in which case we emit a
    GuardBitEquals against nil (side exit reason BlockArgNotNil) before
    stripping it so a non-nil block at runtime takes a side exit.

When the block argument can't be proven nil, we fall back to a dynamic
Send and record the SendBlockArgNotNil fallback reason.

The stripped-block frame state is used only for the callee frame
for the direct send. The pre-call guards (the nil-block GuardBitEquals,
the receiver GuardType, and the method-redefinition PatchPoint) keep
using the original frame state, which still has the block argument on
the stack so that it is there if we side exit.
diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
@@ -452,6 +452,56 @@ def test_float_arithmetic
     assert_compiles '-2', 'def test = (-2.9).to_i; test'
   end
 
+  def test_send_forwarded_block_arg_nil_then_non_nil
+    # Regression test: when a forwarded &block arg is profiled as nil, the nil
+    # block optimization must update the frame state to match the stripped args.
+    # Otherwise the saved SP is off by one, causing a stack consistency error
+    # when the guard side-exits for a non-nil block.
+    assert_runs ':ok', <<~RUBY, call_threshold: 2
+      def inner(callable = nil, &block)
+        callable || block
+      end
+
+      def outer(&block)
+        inner(&block)
+      end
+
+      100.times { outer }
+      result = outer { |x| x }
+      result.is_a?(Proc) ? :ok : :fail
+    RUBY
+  end
+
+  def test_send_forwarded_nil_block_arg_with_polymorphic_receiver
+    # Regression test: the nil block optimization strips the block arg from the
+    # frame state used to set up the callee frame, but the pre-call guards
+    # (receiver GuardType, method-redefinition PatchPoint) must keep using the
+    # original frame state that still has the block arg on the stack. Otherwise a
+    # guard side-exit re-executes the send with a stack that is missing the block
+    # arg slot, corrupting the pushed frame's EP (VM_ENV_FLAGS assertion failure).
+    # A polymorphic receiver forces the receiver GuardType to side-exit.
+    assert_runs ':ok', <<~RUBY, call_threshold: 2
+      class Base
+        def self.inner(model, name, &block)
+          block ? block.call : model
+        end
+        def self.outer(model, name, &block)
+          inner(model, name, &block)
+        end
+      end
+      class A < Base; end
+      class B < Base; end
+      class C < Base; end
+      class D < Base; end
+
+      1000.times do |i|
+        klass = [A, B, C, D][i % 4]
+        klass.outer(i, :n)
+      end
+      :ok
+    RUBY
+  end
+
   private
 
   # Assert that every method call in `test_script` can be compiled by ZJIT
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
@@ -531,6 +531,7 @@ pub enum SideExitReason {
     UnhandledYARVInsn(u32),
     UnhandledCallType(CallType),
     UnhandledBlockArg,
+    BlockArgNotNil,
     TooManyKeywordParameters,
     TooManyArgsForLir,
     FixnumAddOverflow,
@@ -694,6 +695,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,
@@ -768,6 +771,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"),
@@ -2488,7 +2492,7 @@ fn can_direct_send(function: &mut Function, block: BlockId, iseq: *const rb_iseq
     let params = unsafe { iseq.params() };
 
     let callee_has_block_param = 0 != params.flags.has_block();
-    let caller_passes_block_arg = (unsafe { rb_vm_ci_flag(ci) } & VM_CALL_ARGS_BLOCKARG) != 0;
+    let caller_passes_block_arg = has_block && (unsafe { rb_vm_ci_flag(ci) } & VM_CALL_ARGS_BLOCKARG) != 0;
 
     use Counter::*;
     if 0 != params.flags.has_rest()    { count_failure(complex_arg_pass_param_rest) }
@@ -3803,7 +3807,7 @@ impl Function {
                     Insn::Send { recv, block: None, args, state, cd, .. } if ruby_call_method_id(cd) == ID!(minusat) && args.is_empty() =>
                         self.try_rewrite_uminus(block, insn_id, recv, state),
                     ref send@Insn::Send { mut recv, cd, state, block: send_block, ref args, .. } => {
-                        let has_block = send_block.is_some();
+                        let mut has_block = send_block.is_some();
                         let (klass, profiled_type) = match self.resolve_receiver_type(recv, self.type_of(recv), state) {
                             ReceiverTypeResolution::StaticallyKnown { class } => (class, None),
                             ReceiverTypeResolution::Monomorphic { profiled_type }
@@ -3864,9 +3868,57 @@ 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.
+                        // `state` keeps referring to the pre-send frame state (block arg still on the
+                        // stack). Any guard that side-exits before the call re-executes the `send` in
+                        // the interpreter, so it must reconstruct the stack with the block arg present.
+                        // Only the direct-send frame setup uses `send_frame_state`, which has the nil
+                        // block arg stripped from the stack.
+                        let mut send_block = send_block;
+                        let mut send_frame_state = state;
+                        let mut args = args.to_vec();
+                        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 statically_nil = self.is_a(block_arg, types::NilClass);
+                                let profiled_nil = self.profiled_type_of_at(block_arg, state)
+                                    .map_or(false, |pt| pt.is_nil());
+                                if statically_nil || profiled_nil {
+                                    if !statically_nil {
+                                        // Guard needed when relying on profiled type. Uses the original
+                                        // `state` so a side exit re-executes the send with the block
+                                        // arg still on the VM stack.
+                                        self.push_insn(block, Insn::GuardBitEquals {
+                                            val: block_arg,
+                                            expected: Const::Value(Qnil),
+                                            reason: SideExitReason::BlockArgNotNil,
+                                            state,
+                                            recompile: None,
+                                        });
+                                    }
+                                    // Strip nil block arg and treat as no block
+                                    args = args[..args.len() - 1].to_vec();
+                                    send_block = None;
+                                    has_block = false;
+                                    stripped_nil_block = true;
+                                    // Frame state for the direct send only: the block arg is removed
+                                    // from the stack so the callee frame is laid out correctly.
+                                    let new_state = self.frame_state(state).with_replaced_args(&args, args.len() + 1);
+                                    send_frame_state = self.push_insn(block, Insn::Snapshot { state: Box::new(new_state) });
+                                } 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`) and C methods handle their own argument constraints (e.g., kw_splat for Proc call).
-                        if def_type != VM_METHOD_TYPE_OPTIMIZED && def_type != VM_METHOD_TYPE_CFUNC && 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 && def_type != VM_METHOD_TYPE_CFUNC && 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;
@@ -3882,7 +3934,7 @@ impl Function {
                             }
 
                             // Check if the args are compatible before emitting any assmptions
-                            let Ok((send_state, processed_args, kw_bits)) = self.prepare_direct_send_args(block, &args, ci, iseq, state)
+                            let Ok((send_state, processed_args, kw_bits)) = self.prepare_direct_send_args(block, &args, ci, iseq, send_frame_state)
                                 .inspect_err(|&reason| self.set_dynamic_send_reason(insn_id, reason)) else {
                                 self.push_insn_id(block, insn_id); continue;
                             };
@@ -3922,7 +3974,7 @@ impl Function {
                             }
 
                             // Check if the args are compatible before emitting any assmptions
-                            let Ok((send_state, processed_args, kw_bits)) = self.prepare_direct_send_args(block, &args, ci, iseq, state)
+                            let Ok((send_state, processed_args, kw_bits)) = self.prepare_direct_send_args(block, &args, ci, iseq, send_frame_state)
                                 .inspect_err(|&reason| self.set_dynamic_send_reason(insn_id, reason)) else {
                                 self.push_insn_id(block, insn_id); continue;
                             };
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
@@ -5207,7 +5207,7 @@ mod hir_opt_tests {
           v26:TrueClass = GuardBitEquals v25, Value(true) recompile
           Jump bb6(v24, v10)
         bb6(v16:BasicObject, v17:BasicObject):
-          v29:BasicObject = Send v14, &block, :then, v16 # SendFallbackReason: Complex argument passing
+          v29:BasicObject = Send v14, &block, :then, v16 # SendFallbackReason: Send: block argument is not nil
           CheckInterrupts
           Return v29
         ");
@@ -5263,7 +5263,7 @@ mod hir_opt_tests {
           v37:NilClass = Const Value(nil)
           Jump bb8(v37, v13)
         bb8(v27:BasicObject, v28:BasicObject):
-          v40:BasicObject = Send v25, &block, :then, v27 # SendFallbackReason: Complex argument passing
+          v40:BasicObject = Send v25, &block, :then, v27 # SendFallbackReason: Send: block argument is not nil
           CheckInterrupts
           Return v40
         bb4(v45:BasicObject, v46:Falsy, v47:BasicObject):
@@ -5428,7 +5428,7 @@ mod hir_opt_tests {
           v34:NilClass = Const Value(nil)
           Jump bb6(v34, v10)
         bb6(v16:BasicObject, v17:BasicObject):
-          v38:BasicObject = Send v14, &block, :then, v16 # SendFallbackReason: Complex argument passing
+          v38:BasicObject = Send v14, &block, :then, v16 # SendFallbackReason: Send: block argument is not nil
           CheckInterrupts
           Return v38
         bb10():
@@ -5496,7 +5496,7 @@ mod hir_opt_tests {
           v41:ObjectSubclass[BlockParamProxy] = Const Value(VALUE(0x1010))
           Jump bb6(v41, v10)
         bb6(v16:BasicObject, v17:BasicObject):
-          v45:BasicObject = Send v14, &block, :then, v16 # SendFallbackReason: Complex argument passing
+          v45:BasicObject = Send v14, &block, :then, v16 # SendFallbackReason: Send: block argument is not nil
           CheckInterrupts
           Return v45
         bb13():
@@ -9519,7 +9519,7 @@ mod hir_opt_tests {
           v26:ObjectSubclass[BlockParamProxy] = Const Value(VALUE(0x1008))
           Jump bb6(v26, v10)
         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
         ");
@@ -9559,9 +9559,12 @@ mod hir_opt_tests {
           v26:NilClass = Const Value(nil)
           Jump bb6(v26, v10)
         bb6(v16:BasicObject, v17:BasicObject):
-          v29:BasicObject = Send v14, &block, :map, v16 # SendFallbackReason: Complex argument passing
+          v35:NilClass = GuardBitEquals v16, Value(nil)
+          PatchPoint NoSingletonClass(Array@0x1008)
+          PatchPoint MethodRedefined(Array@0x1008, map@0x1010, cme:0x1018)
+          v40:BasicObject = SendDirect v14, 0x1040, :map (0x1050)
           CheckInterrupts
-          Return v29
+          Return v40
         ");
     }
 
@@ -9600,7 +9603,7 @@ mod hir_opt_tests {
           v21:ObjectSubclass[BlockParamProxy] = Const Value(VALUE(0x1008))
           Jump bb6(v21)
         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
         ");
@@ -9768,7 +9771,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
         ");
@@ -9799,9 +9802,11 @@ 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
+          PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
+          v27:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v8, ObjectSubclass[class_exact*:Object@VALUE(0x1000)] recompile
+          v28:Fixnum[42] = Const Value(42)
           CheckInterrupts
-          Return v19
+          Return v28
         ");
     }
 
@@ -9830,19 +9835,24 @@ mod hir_opt_tests {
           Jump bb3(v6, v7)
         bb3(v9:BasicObject, v10:BasicObject):
           v17:CPtr = GetEP 0
-          v18:CBool = IsBlockParamModified v17
-          IfTrue v18, bb4()
-          v23:CInt64 = LoadField v17, :_env_data_index_specval@0x1001
-          v24:CInt64[0] = GuardBitEquals v23, CInt64(0)
-          v25:NilClass = Const Value(nil)
-          Jump bb6(v25, v10)
+          v18:CUInt64 = LoadField v17, :VM_ENV_DATA_INDEX_FLAGS@0x1001
+          v19:CBool = IsBlockParamModified v18
+          CondBranch v19, bb4(), bb5()
         bb4():
           v21:BasicObject = LoadField v17, :block@0x1002
           Jump bb6(v21, v21)
+        bb5():
+          v23:CInt64 = LoadField v17, :VM_ENV_DATA_INDEX_SPECVAL@0x1003
+          v24:CInt64[0] = GuardBitEquals v23, CInt64(0) recompile
+          v25:NilClass = Const Value(nil)
+          Jump bb6(v25, v10)
         bb6(v15:BasicObject, v16:BasicObject):
-          v28:BasicObject = Send v9, &block, :foo, v15 # SendFallbackReason: Complex argument passing
+          v34:NilClass = GuardBitEquals v15, Value(nil)
+          PatchPoint MethodRedefined(Object@0x1008, foo@0x1010, cme:0x1018)
+          v37:ObjectSubclass[class_exact*:Object@VALUE(0x1008)] = GuardType v9, ObjectSubclass[class_exact*:Object@VALUE(0x1008)] recompile
+          v38:Fixnum[42] = Const Value(42)
           CheckInterrupts
-          Return v28
+          Return v38
         ");
     }
 
@@ -13203,7 +13213,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
         ");
@@ -13248,7 +13258,7 @@ mod hir_opt_tests {
           v26:ObjectSubclass[BlockParamProxy] = Const Value(VALUE(0x1008))
           Jump bb6(v26, v10)
         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
         ");
@@ -17095,7 +17105,7 @@ mod hir_opt_tests {
             test(true, block)
         ");
 
-        assert_snapshot!(hir_string("test"), @r"
+        assert_snapshot!(hir_string("test"), @"
         fn test@<compiled>:7:
         bb1():
           EntryPoint interpreter
@@ -17118,7 +17128,7 @@ mod hir_opt_tests {
         bb5():
           v22:Truthy = RefineType v12, Truthy
           v26:Fixnum[42] = Const Value(42)
-          v29:BasicObject = Send v11, &block, :passthrough_recompile_blockarg, v26, v13 # SendFallbackReason: Complex argument passing
+          v29:BasicObject = Send v11, &block, :passthrough_recompile_blockarg, v26, v13 # SendFallbackReason: Send: block argument is not nil
           CheckInterrupts
           Return v29
         bb4(v34:BasicObject, v35:Falsy, v36:BasicObject):
@@ -19326,7 +19336,7 @@ mod hir_opt_tests {
           v47:ObjectSubclass[BlockParamProxy] = Const Value(VALUE(0x1050))
           Jump bb8(v47, v55)
         bb8(v37:BasicObject, v38:BasicObject):
-          v50:BasicObject = Send v27, &block, :inner, v10, v37 # SendFallbackReason: Complex argument passing
+          v50:BasicObject = Send v27, &block, :inner, v10, v37 # SendFallbackReason: Send: block argument is not nil
           CheckInterrupts
           PopInlineFrame
           PatchPoint NoEPEscape(test)
diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs
