ZJIT: Spill stack from direct send with args in callee order

Author: rwstauner

test/ruby/test_zjit.rb

@@ -383,6 +383,27 @@ def f(a)
     }, call_threshold: 2
   end
 
+  def test_kwargs_with_exit_and_local_invalidation
+    assert_compiles ':ok', %q{
+      def a(b:, c:)
+        if c == :b
+          return -> {}
+        end
+        Class # invalidate locals
+
+        raise "c is :b!" if c == :b
+      end
+
+      def test
+        # note opposite order of kwargs
+        a(c: :c, b: :b)
+      end
+
+      4.times { test }
+      :ok
+    }, call_threshold: 2
+  end
+
   def test_setlocal_on_eval
     assert_compiles '1', %q{
       @b = binding

zjit/src/codegen.rs

@@ -1346,7 +1346,7 @@ fn gen_send_without_block_direct(
     gen_save_sp(asm, state.stack().len() - args.len() - 1); // -1 for receiver
 
     gen_spill_locals(jit, asm, state);
-    gen_spill_stack(jit, asm, state);
+    gen_spill_stack_for_send(jit, asm, state, &args);
 
     let (frame_type, specval) = if VM_METHOD_TYPE_BMETHOD == unsafe { get_cme_def_type(cme) } {
         // Extract EP from the Proc instance
@@ -2211,6 +2211,24 @@ fn gen_spill_stack(jit: &JITState, asm: &mut Assembler, state: &FrameState) {
     }
 }
 
+/// Spill stack for a direct send, with args in callee order.
+/// This is needed because reorder_keyword_arguments may reorder the args in the HIR
+/// but the FrameState's stack remains in caller order.
+fn gen_spill_stack_for_send(jit: &JITState, asm: &mut Assembler, state: &FrameState, args: &[Opnd]) {
+    gen_incr_counter(asm, Counter::vm_write_stack_count);
+    let stack: Vec<_> = state.stack().collect();
+    let recv_idx = stack.len() - args.len() - 1;
+    asm_comment!(asm, "spill stack up to and including receiver");
+    for idx in 0..=recv_idx {
+        asm.mov(Opnd::mem(64, SP, idx as i32 * SIZEOF_VALUE_I32), jit.get_opnd(*stack[idx]));
+    }
+    asm_comment!(asm, "spill args in callee order");
+    for (i, arg) in args.iter().enumerate() {
+        let offset = (recv_idx + 1 + i) as i32 * SIZEOF_VALUE_I32;
+        asm.mov(Opnd::mem(64, SP, offset), *arg);
+    }
+}
+
 /// Prepare for calling a C function that may call an arbitrary method.
 /// Use gen_prepare_leaf_call_with_gc() if the method is leaf but allocates objects.
 fn gen_prepare_non_leaf_call(jit: &JITState, asm: &mut Assembler, state: &FrameState) {