ZJIT: Spill to the stack using arguments instead of FrameState

The FrameState on the SendWithoutBlock represents the entry state,
but for instructions that push to the stack in the middle of the
instruction before actually doing the send like opt_aref_with,
the FrameState is incorrect. We need to write to the stack using the
arguments for the instruction.

2 files changed, 17 insertions, 7 deletions

diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index 4452f413f0..796851a9bf 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -487,6 +487,14 @@ class TestZJIT < Test::Unit::TestCase
     }, call_threshold: 5, num_profiles: 3
   end
 
+  def test_opt_aref_with
+    assert_compiles ':ok', %q{
+      def aref_with(hash) = hash["key"]
+
+      aref_with({ "key" => :ok })
+    }
+  end
+
   # tool/ruby_vm/views/*.erb relies on the zjit instructions a) being contiguous and
   # b) being reliably ordered after all the other instructions.
   def test_instruction_order
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 844ac5df42..d5202486f1 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -257,7 +257,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::Jump(branch) => return gen_jump(jit, asm, branch),
         Insn::IfTrue { val, target } => return gen_if_true(jit, asm, opnd!(val), target),
         Insn::IfFalse { val, target } => return gen_if_false(jit, asm, opnd!(val), target),
-        Insn::SendWithoutBlock { call_info, cd, state, .. } => gen_send_without_block(jit, asm, call_info, *cd, &function.frame_state(*state))?,
+        Insn::SendWithoutBlock { call_info, cd, state, self_val, args, .. } => gen_send_without_block(jit, asm, call_info, *cd, &function.frame_state(*state), self_val, args)?,
         Insn::SendWithoutBlockDirect { iseq, self_val, args, .. } => gen_send_without_block_direct(cb, jit, asm, *iseq, opnd!(self_val), args)?,
         Insn::Return { val } => return Some(gen_return(asm, opnd!(val))?),
         Insn::FixnumAdd { left, right, state } => gen_fixnum_add(asm, opnd!(left), opnd!(right), &function.frame_state(*state))?,
@@ -443,10 +443,12 @@ fn gen_send_without_block(
     call_info: &CallInfo,
     cd: *const rb_call_data,
     state: &FrameState,
+    self_val: &InsnId,
+    args: &Vec<InsnId>,
 ) -> Option<lir::Opnd> {
-    // Spill the virtual stack onto the stack. They need to be marked by GC and may be caller-saved registers.
+    // Spill the receiver and the arguments onto the stack. They need to be marked by GC and may be caller-saved registers.
     // TODO: Avoid spilling operands that have been spilled before.
-    for (idx, &insn_id) in state.stack().enumerate() {
+    for (idx, &insn_id) 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, jit.get_opnd(insn_id)?);
@@ -454,7 +456,7 @@ fn gen_send_without_block(
 
     // Save PC and SP
     gen_save_pc(asm, state);
-    gen_save_sp(asm, state);
+    gen_save_sp(asm, 1 + args.len()); // +1 for receiver
 
     asm_comment!(asm, "call #{} with dynamic dispatch", call_info.method_name);
     unsafe extern "C" {
@@ -678,13 +680,13 @@ fn gen_save_pc(asm: &mut Assembler, state: &FrameState) {
 }
 
 /// Save the current SP on the CFP
-fn gen_save_sp(asm: &mut Assembler, state: &FrameState) {
+fn gen_save_sp(asm: &mut Assembler, stack_size: usize) {
     // Update cfp->sp which will be read by the interpreter. We also have the SP register in JIT
     // code, and ZJIT's codegen currently assumes the SP register doesn't move, e.g. gen_param().
     // So we don't update the SP register here. We could update the SP register to avoid using
     // an extra register for asm.lea(), but you'll need to manage the SP offset like YJIT does.
-    asm_comment!(asm, "save SP to CFP: {}", state.stack_size());
-    let sp_addr = asm.lea(Opnd::mem(64, SP, state.stack_size() as i32 * SIZEOF_VALUE_I32));
+    asm_comment!(asm, "save SP to CFP: {}", stack_size);
+    let sp_addr = asm.lea(Opnd::mem(64, SP, stack_size as i32 * SIZEOF_VALUE_I32));
     let cfp_sp = Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SP);
     asm.mov(cfp_sp, sp_addr);
 }