ZJIT: Exit with the block arg on the stack from SendDirect args

When a direct send drops a profiled-nil `&block` argument, it lays out
the callee frame from a snapshot with the block arg stripped from the VM
stack. emit_send_direct_args() also used that snapshot for instructions
that materialize the callee's arguments before the call, such as the rest
array of a `*args` callee. Those can side-exit, e.g. at the NoNewObjHook
patch point of the inline allocation once a NEWOBJ hook is enabled, and the
exit re-executes the send in the interpreter at the same PC. With the block
arg missing from the stack, the interpreter took the last positional
argument as the block:

    TypeError: no implicit conversion of Net::HTTP::Post into Proc

This happened in Shopify's Sonic test suite after a test enabled
ObjectSpace.trace_object_allocations: Net::HTTP#request passes `&block` to
Semian's `transport_request(*)`, and `request` had reached the version
limit, so its invalidated code kept exiting there.

Use the send's original frame state for those instructions.

Author: k0kubun

zjit/src/hir.rs

@@ -3984,11 +3984,16 @@ impl Function {
     }
 
     /// Materialize a validated SendDirect call in the selected runtime path.
-    fn emit_send_direct_args(&mut self, block: BlockId, call: SendDirectCall, original_args: &[InsnId], state: InsnId) -> SendDirectArgs {
+    ///
+    /// `state` is the frame state that the callee frame is laid out from, which may have a nil
+    /// block arg stripped from the stack. `exit_state` is the frame state of the send itself, used
+    /// by instructions that materialize arguments before the call (e.g. the rest array), since a
+    /// side exit from them re-executes the send in the interpreter with the original VM stack.
+    fn emit_send_direct_args(&mut self, block: BlockId, call: SendDirectCall, original_args: &[InsnId], state: InsnId, exit_state: InsnId) -> SendDirectArgs {
         let args: Vec<_> = call
             .args
             .into_iter()
-            .map(|arg| self.emit_send_direct_arg(block, arg, state))
+            .map(|arg| self.emit_send_direct_arg(block, arg, exit_state))
             .collect();
 
         // If args were reordered or synthesized, create a new snapshot with the updated stack.
@@ -4882,7 +4887,7 @@ impl Function {
                             }
 
                             let SendDirectArgs { state: send_state, args: send_args, kw_bits, jit_entry_idx } =
-                                self.emit_send_direct_args(block, call, &args, send_frame_state);
+                                self.emit_send_direct_args(block, call, &args, send_frame_state, state);
                             let replacement = self.try_inline_send_direct(block, Insn::SendDirect(Box::new(SendDirectData { recv, cd, cme, iseq, args: send_args, kw_bits, jit_entry_idx, state: send_state, block: send_block })));
                             self.make_equal_to(insn_id, replacement);
                         } else if !has_block && def_type == VM_METHOD_TYPE_BMETHOD {
@@ -4923,7 +4928,7 @@ impl Function {
                             }
 
                             let SendDirectArgs { state: send_state, args: send_args, kw_bits, jit_entry_idx } =
-                                self.emit_send_direct_args(block, call, &args, send_frame_state);
+                                self.emit_send_direct_args(block, call, &args, send_frame_state, state);
                             let replacement = self.try_inline_send_direct(block, Insn::SendDirect(Box::new(SendDirectData { recv, cd, cme, iseq, args: send_args, kw_bits, jit_entry_idx, state: send_state, block: None })));
                             self.make_equal_to(insn_id, replacement);
                         } else if !has_block && def_type == VM_METHOD_TYPE_IVAR && args.is_empty() {
@@ -5456,7 +5461,7 @@ impl Function {
                             emit_super_call_guards(self, block, super_cme, current_cme, mid, state, frame_state_iseq);
 
                             let SendDirectArgs { state: send_state, args: send_args, kw_bits, jit_entry_idx } =
-                                self.emit_send_direct_args(block, call, &args, state);
+                                self.emit_send_direct_args(block, call, &args, state, state);
                             // Use SendDirect with the super method's CME and ISEQ.
                             let replacement = self.try_inline_send_direct(block, Insn::SendDirect(Box::new(SendDirectData {
                                 recv,

zjit/src/hir/opt_tests.rs

@@ -26,6 +26,37 @@ mod hir_opt_tests {
         hir_string_proc(&format!("{}.method(:{})", "self", method))
     }
 
+    // A direct send strips a profiled-nil `&block` argument from the stack that the callee
+    // frame is laid out from, but instructions that materialize the callee's arguments, like
+    // the rest array, can side-exit and re-execute the send in the interpreter. Their frame
+    // state must keep the block argument on the stack, or the interpreter takes the last
+    // positional argument as the block.
+    #[test]
+    fn test_send_direct_rest_array_keeps_stripped_block_arg_in_exit_state() {
+        eval("
+            def rest_callee(*args) = args
+            def test(obj, &block) = rest_callee(obj, &block)
+            test(1)
+            test(1)
+        ");
+        let iseq = crate::cruby::with_rubyvm(|| get_method_iseq("self", "test"));
+        unsafe { crate::cruby::rb_zjit_profile_disable(iseq) };
+        let mut function = iseq_to_hir(iseq).unwrap();
+        function.optimize();
+        function.validate().unwrap();
+
+        let new_array_stack_sizes: Vec<usize> = (0..function.num_insns())
+            .map(InsnId::from)
+            .filter_map(|insn_id| match function.find(insn_id) {
+                Insn::NewArray { state, .. } => Some(function.frame_state(state).stack().len()),
+                _ => None,
+            })
+            .collect();
+        // [self, obj, block]
+        assert!(!new_array_stack_sizes.is_empty(), "{}", hir_string_function(&function));
+        assert!(new_array_stack_sizes.iter().all(|&size| size == 3), "{new_array_stack_sizes:?}");
+    }
+
     #[test]
     fn test_fold_iftrue_away() {
         eval("