ZJIT: Fix profile_stack to skip block arg for ARGS_BLOCKARG sends (#16581)

For sends with ARGS_BLOCKARG (e.g. `foo(&block)`), the block arg sits
on the stack above the receiver and regular arguments. The profiling
(both interpreter and exit profiling) only records types for the
receiver and regular args (argc + 1 values), but profile_stack was
mapping those types onto stack positions starting from the top, which
incorrectly mapped them onto the block arg and the wrong operands.

Fix by detecting ARGS_BLOCKARG at the profile_stack call site and
passing a stack_offset of 1 to skip the block arg. This allows
resolve_receiver_type_from_profile to find the correct receiver type,
enabling method dispatch optimization for these sends after
recompilation via exit profiling.

2 files changed, 55 insertions, 5 deletions

diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 591ca0aad7..9227bbd730 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -6758,15 +6758,17 @@ impl ProfileOracle {
         Self { payload, types: Default::default() }
     }
 
-    /// Map the interpreter-recorded types of the stack onto the HIR operands on our compile-time virtual stack
-    fn profile_stack(&mut self, state: &FrameState) {
+    /// Map the interpreter-recorded types of the stack onto the HIR operands on our compile-time virtual stack.
+    /// `stack_offset` is the number of extra stack entries above the profiled operands (e.g. 1 for
+    /// sends with ARGS_BLOCKARG, where the block arg sits on top of the regular args).
+    fn profile_stack(&mut self, state: &FrameState, stack_offset: usize) {
         let iseq_insn_idx = state.insn_idx;
         let Some(operand_types) = self.payload.profile.get_operand_types(iseq_insn_idx) else { return };
         let entry = self.types.entry(iseq_insn_idx).or_default();
         // operand_types is always going to be <= stack size (otherwise it would have an underflow
         // at run-time) so use that to drive iteration.
         for (idx, insn_type_distribution) in operand_types.iter().rev().enumerate() {
-            let insn = state.stack_topn(idx).expect("Unexpected stack underflow in profiling");
+            let insn = state.stack_topn(idx + stack_offset).expect("Unexpected stack underflow in profiling");
             entry.push((insn, TypeDistributionSummary::new(insn_type_distribution)))
         }
     }
@@ -6969,7 +6971,17 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                 }
             }
             else {
-                profiles.profile_stack(&exit_state);
+                // For sends with ARGS_BLOCKARG, the block arg sits on the stack above
+                // the profiled operands (receiver + regular args). Skip it so that the
+                // profile types map onto the correct HIR operands.
+                let stack_offset = if opcode == YARVINSN_send || opcode == YARVINSN_opt_send_without_block {
+                    let cd: *const rb_call_data = get_arg(pc, 0).as_ptr();
+                    let flags = unsafe { vm_ci_flag(rb_get_call_data_ci(cd)) };
+                    usize::from(flags & VM_CALL_ARGS_BLOCKARG != 0)
+                } else {
+                    0
+                };
+                profiles.profile_stack(&exit_state, stack_offset);
             }
 
             // Flag a future getlocal/setlocal to add a patch point if this instruction is not leaf.
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
index 2d3e43d0f1..c7b2caf31a 100644
--- a/zjit/src/hir/opt_tests.rs
+++ b/zjit/src/hir/opt_tests.rs
@@ -11427,13 +11427,51 @@ mod hir_opt_tests {
           Jump bb3(v4)
         bb3(v6:BasicObject):
           v11:StaticSymbol[:the_block] = Const Value(VALUE(0x1000))
-          v13:BasicObject = Send v6, 0x1008, :callee, v11 # SendFallbackReason: Send: no profile data available
+          v13:BasicObject = Send v6, 0x1008, :callee, v11 # SendFallbackReason: Complex argument passing
           CheckInterrupts
           Return v13
         ");
     }
 
     #[test]
+    fn test_profile_stack_skips_block_arg() {
+        // Regression test: profile_stack must skip the &block arg on the stack when mapping
+        // profiled operand types. Without the fix, the receiver type would be mapped to the
+        // wrong stack slot, causing resolve_receiver_type to return NoProfile.
+        // With the fix, the receiver type is correctly resolved and the send gets past type
+        // resolution to hit the ARGS_BLOCKARG guard (ComplexArgPass) instead of NoProfile.
+        eval("
+            def test(&block) = [].map(&block)
+            test { |x| x }; test { |x| x }
+        ");
+        assert_snapshot!(hir_string("test"), @"
+        fn test@<compiled>:2:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:CPtr = LoadSP
+          v3:BasicObject = LoadField v2, :block@0x1000
+          Jump bb3(v1, v3)
+        bb2():
+          EntryPoint JIT(0)
+          v6:BasicObject = LoadArg :self@0
+          v7:BasicObject = LoadArg :block@1
+          Jump bb3(v6, v7)
+        bb3(v9:BasicObject, v10:BasicObject):
+          v14:ArrayExact = NewArray
+          v16:CPtr = GetEP 0
+          v17:CInt64 = LoadField v16, :_env_data_index_flags@0x1001
+          v18:CInt64 = GuardNoBitsSet v17, VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM=CUInt64(512)
+          v19:CInt64 = LoadField v16, :_env_data_index_specval@0x1002
+          v20:CInt64 = GuardAnyBitSet v19, CUInt64(1)
+          v21:ObjectSubclass[BlockParamProxy] = Const Value(VALUE(0x1008))
+          v23:BasicObject = Send v14, 0x1001, :map, v21 # SendFallbackReason: Complex argument passing
+          CheckInterrupts
+          Return v23
+        ");
+    }
+
+    #[test]
     fn test_optimize_stringexact_eq_stringexact() {
         eval(r#"
             def test(l, r) = l == r