ZJIT: Spill whole FrameState in `Insn::SendWithoutBlock`

Previously, we only spilled the arguments necessary for the particular
send. In case the callee raises and a rescue resumes the ISEQ, that
did not present a complete stack state. E.g. in `[1, (raise rescue 2)]`
the raise send only spills `self`, when `1` also needs to be spilled.

Spill the whole stack. Adjust parsing for `opt_aref_with` since the
key argument for the send now needs to be described by the frame state
of the send.

This changes the contract for `Insn::SendWithoutBlock` to use arguments
from the interpreter stack as described by its frame state.

3 files changed, 61 insertions, 33 deletions

diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index 8c4ddd598e..8f83e858a6 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -1902,6 +1902,17 @@ class TestZJIT < Test::Unit::TestCase
     }, call_threshold: 2
   end
 
+  def test_raise_in_second_argument
+    assert_compiles '{ok: true}', %q{
+      def write(hash, key)
+        hash[key] = raise rescue true
+        hash
+      end
+
+      write({}, :ok)
+    }
+  end
+
   private
 
   # Assert that every method call in `test_script` can be compiled by ZJIT
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 16d44c74b0..7ff5564d22 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -353,10 +353,10 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::Jump(branch) => no_output!(gen_jump(jit, asm, branch)),
         Insn::IfTrue { val, target } => no_output!(gen_if_true(jit, asm, opnd!(val), target)),
         Insn::IfFalse { val, target } => no_output!(gen_if_false(jit, asm, opnd!(val), target)),
-        Insn::SendWithoutBlock { cd, state, self_val, args, .. } => gen_send_without_block(jit, asm, *cd, &function.frame_state(*state), opnd!(self_val), opnds!(args)),
+        Insn::SendWithoutBlock { cd, state, .. } => gen_send_without_block(jit, asm, *cd, &function.frame_state(*state)),
         // Give up SendWithoutBlockDirect for 6+ args since asm.ccall() doesn't support it.
-        Insn::SendWithoutBlockDirect { cd, state, self_val, args, .. } if args.len() + 1 > C_ARG_OPNDS.len() => // +1 for self
-            gen_send_without_block(jit, asm, *cd, &function.frame_state(*state), opnd!(self_val), opnds!(args)),
+        Insn::SendWithoutBlockDirect { cd, state, args, .. } if args.len() + 1 > C_ARG_OPNDS.len() => // +1 for self
+            gen_send_without_block(jit, asm, *cd, &function.frame_state(*state)),
         Insn::SendWithoutBlockDirect { cme, iseq, self_val, args, state, .. } => gen_send_without_block_direct(cb, jit, asm, *cme, *iseq, opnd!(self_val), opnds!(args), &function.frame_state(*state)),
         // Ensure we have enough room fit ec, self, and arguments
         // TODO remove this check when we have stack args (we can use Time.new to test it)
@@ -859,26 +859,19 @@ fn gen_send_without_block(
     asm: &mut Assembler,
     cd: *const rb_call_data,
     state: &FrameState,
-    self_val: Opnd,
-    args: Vec<Opnd>,
 ) -> lir::Opnd {
-    gen_spill_locals(jit, asm, state);
-    // Spill the receiver and the arguments onto the stack.
-    // They need to be on the interpreter stack to let the interpreter access them.
-    // TODO: Avoid spilling operands that have been spilled before.
-    // TODO: Despite https://github.com/ruby/ruby/pull/13468, Kokubun thinks this should
-    // spill the whole stack in case it raises an exception. The HIR might need to change
-    // for opt_aref_with, which pushes to the stack in the middle of the instruction.
-    asm_comment!(asm, "spill receiver and arguments");
-    for (idx, &val) in [self_val].iter().chain(args.iter()).enumerate() {
-        // Currently, we don't move the SP register. So it's equal to the base pointer.
-        let stack_opnd = Opnd::mem(64, SP, idx as i32 * SIZEOF_VALUE_I32);
-        asm.mov(stack_opnd, val);
-    }
+    // Note that it's incorrect to use this frame state to side exit because
+    // the state might not be on the boundary of an interpreter instruction.
+    // For example, `opt_aref_with` pushes to the stack and then sends.
+    asm_comment!(asm, "spill frame state");
 
     // Save PC and SP
     gen_save_pc(asm, state);
-    gen_save_sp(asm, 1 + args.len()); // +1 for receiver
+    gen_save_sp(asm, state.stack().len());
+
+    // Spill locals and stack
+    gen_spill_locals(jit, asm, state);
+    gen_spill_stack(jit, asm, state);
 
     asm_comment!(asm, "call #{} with dynamic dispatch", ruby_call_method_name(cd));
     unsafe extern "C" {
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 1fdca5a688..1cd31497d8 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -818,7 +818,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 }
                 Ok(())
             },
-            Insn::Snapshot { state } => write!(f, "Snapshot {}", state),
+            Insn::Snapshot { state } => write!(f, "Snapshot {}", state.print(self.ptr_map)),
             Insn::Defined { op_type, v, .. } => {
                 // op_type (enum defined_type) printing logic from iseq.c.
                 // Not sure why rb_iseq_defined_string() isn't exhaustive.
@@ -2513,11 +2513,10 @@ pub struct FrameState {
     locals: Vec<InsnId>,
 }
 
-impl FrameState {
-    /// Get the opcode for the current instruction
-    pub fn get_opcode(&self) -> i32 {
-        unsafe { rb_iseq_opcode_at_pc(self.iseq, self.pc) }
-    }
+/// Print adaptor for [`FrameState`]. See [`PtrPrintMap`].
+pub struct FrameStatePrinter<'a> {
+    inner: &'a FrameState,
+    ptr_map: &'a PtrPrintMap,
 }
 
 /// Compute the index of a local variable from its slot index
@@ -2624,14 +2623,24 @@ impl FrameState {
         args.extend(self.locals.iter().chain(self.stack.iter()).map(|op| *op));
         args
     }
+
+    /// Get the opcode for the current instruction
+    pub fn get_opcode(&self) -> i32 {
+        unsafe { rb_iseq_opcode_at_pc(self.iseq, self.pc) }
+    }
+
+    pub fn print<'a>(&'a self, ptr_map: &'a PtrPrintMap) -> FrameStatePrinter<'a> {
+        FrameStatePrinter { inner: self, ptr_map }
+    }
 }
 
-impl std::fmt::Display for FrameState {
+impl Display for FrameStatePrinter<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
-        write!(f, "FrameState {{ pc: {:?}, stack: ", self.pc)?;
-        write_vec(f, &self.stack)?;
+        let inner = self.inner;
+        write!(f, "FrameState {{ pc: {:?}, stack: ", self.ptr_map.map_ptr(inner.pc))?;
+        write_vec(f, &inner.stack)?;
         write!(f, ", locals: ")?;
-        write_vec(f, &self.locals)?;
+        write_vec(f, &inner.locals)?;
         write!(f, " }}")
     }
 }
@@ -3195,9 +3204,11 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     let aref_arg = fun.push_insn(block, Insn::Const { val: Const::Value(get_arg(pc, 0)) });
                     let args = vec![aref_arg];
 
+                    let mut send_state = state.clone();
+                    send_state.stack_push(aref_arg);
+                    let send_state = fun.push_insn(block, Insn::Snapshot { state: send_state });
                     let recv = state.stack_pop()?;
-                    let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
-                    let send = fun.push_insn(block, Insn::SendWithoutBlock { self_val: recv, cd, args, state: exit_id });
+                    let send = fun.push_insn(block, Insn::SendWithoutBlock { self_val: recv, cd, args, state: send_state });
                     state.stack_push(send);
                 }
                 YARVINSN_opt_neq => {
@@ -3921,6 +3932,12 @@ mod tests {
     }
 
     #[track_caller]
+    pub fn assert_function_hir_with_frame_state(function: Function, expected_hir: Expect) {
+        let actual_hir = format!("{}", FunctionPrinter::with_snapshot(&function));
+        expected_hir.assert_eq(&actual_hir);
+    }
+
+    #[track_caller]
     fn assert_compile_fails(method: &str, reason: ParseError) {
         let iseq = crate::cruby::with_rubyvm(|| get_method_iseq("self", method));
         unsafe { crate::cruby::rb_zjit_profile_disable(iseq) };
@@ -5154,11 +5171,18 @@ mod tests {
         eval("
             def test(a) = a['string lit triggers aref_with']
         ");
-        assert_method_hir("test",  expect![[r#"
+
+        let iseq = crate::cruby::with_rubyvm(|| get_method_iseq("self", "test"));
+        assert!(iseq_contains_opcode(iseq, YARVINSN_opt_aref_with));
+        let function = iseq_to_hir(iseq).unwrap();
+        assert_function_hir_with_frame_state(function, expect![[r#"
             fn test@<compiled>:2:
             bb0(v0:BasicObject, v1:BasicObject):
-              v3:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+              v2:Any = Snapshot FrameState { pc: 0x1000, stack: [], locals: [v1] }
+              v3:StringExact[VALUE(0x1008)] = Const Value(VALUE(0x1008))
+              v4:Any = Snapshot FrameState { pc: 0x1010, stack: [v1, v3], locals: [v1] }
               v5:BasicObject = SendWithoutBlock v1, :[], v3
+              v6:Any = Snapshot FrameState { pc: 0x1018, stack: [v5], locals: [v1] }
               CheckInterrupts
               Return v5
         "#]]);