ZJIT: Snapshot FrameState with reordered args before direct send

You can see the reordered args in the new Snapshot right before the
DirectSend insn:

       v14:Any = Snapshot FrameState { pc: 0x00, stack: [v6, v11, v13], locals: [] }
       PatchPoint MethodRedefined(Object@0x00, a@0x00, cme:0x00)
       PatchPoint NoSingletonClass(Object@0x00)
       v22:HeapObject[class_exact*:Object@VALUE(0x00)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x00)]
    -  v23:BasicObject = SendWithoutBlockDirect v22, :a (0x00), v13, v11
    -  v16:Any = Snapshot FrameState { pc: 0x00, stack: [v23], locals: [] }
    +  v23:Any = Snapshot FrameState { pc: 0x00, stack: [v6, v13, v11], locals: [] }
    +  v24:BasicObject = SendWithoutBlockDirect v22, :a (0x00), v13, v11
    +  v16:Any = Snapshot FrameState { pc: 0x00, stack: [v24], locals: [] }

4 files changed, 95 insertions, 24 deletions

diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index 43db676d5d..cf3c46b3ed 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -383,6 +383,43 @@ class TestZJIT < Test::Unit::TestCase
     }, 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_kwargs_with_max_direct_send_arg_count
+    # Ensure that we only reorder the args when we _can_ use direct send (< 6 args).
+    assert_compiles '[[1, 2, 3, 4, 5, 6, 7, 8]]', %q{
+      def kwargs(five, six, a:, b:, c:, d:, e:, f:)
+        [a, b, c, d, five, six, e, f]
+      end
+
+      5.times.flat_map do
+        [
+          kwargs(5, 6, d: 4, c: 3, a: 1, b: 2, e: 7, f: 8),
+          kwargs(5, 6, d: 4, c: 3, b: 2, a: 1, e: 7, f: 8)
+        ]
+      end.uniq
+    }, call_threshold: 2
+  end
+
   def test_setlocal_on_eval
     assert_compiles '1', %q{
       @b = binding
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 41ddb3ae69..36aee8056c 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -6,6 +6,7 @@
 #![allow(clippy::if_same_then_else)]
 #![allow(clippy::match_like_matches_macro)]
 use crate::{
+    backend::lir::C_ARG_OPNDS,
     cast::IntoUsize, codegen::local_idx_to_ep_offset, cruby::*, payload::{get_or_create_iseq_payload, IseqPayload}, options::{debug, get_option, DumpHIR}, state::ZJITState, json::Json
 };
 use std::{
@@ -2699,19 +2700,30 @@ impl Function {
                             }
 
                             let kwarg = unsafe { rb_vm_ci_kwarg(ci) };
-                            let processed_args = if !kwarg.is_null() {
+                            let (send_state, processed_args) = if !kwarg.is_null() {
                                 match self.reorder_keyword_arguments(&args, kwarg, iseq) {
-                                    Ok(reordered) => reordered,
+                                    Ok(reordered) => {
+                                        // Only use reordered state if args fit in C registers.
+                                        // Fallback to interpreter needs original order for kwarg handling.
+                                        // NOTE: This needs to match with the condition in codegen.rs.
+                                        if reordered.len() + 1 <= C_ARG_OPNDS.len() {
+                                            let new_state = self.frame_state(state).with_reordered_args(&reordered);
+                                            let snapshot = self.push_insn(block, Insn::Snapshot { state: new_state });
+                                            (snapshot, reordered)
+                                        } else {
+                                            (state, reordered)
+                                        }
+                                    }
                                     Err(reason) => {
                                         self.set_dynamic_send_reason(insn_id, reason);
                                         self.push_insn_id(block, insn_id); continue;
                                     }
                                 }
                             } else {
-                                args.clone()
+                                (state, args.clone())
                             };
 
-                            let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args: processed_args, state });
+                            let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args: processed_args, state: send_state });
                             self.make_equal_to(insn_id, send_direct);
                         } else if def_type == VM_METHOD_TYPE_BMETHOD {
                             let procv = unsafe { rb_get_def_bmethod_proc((*cme).def) };
@@ -2748,19 +2760,30 @@ impl Function {
                             }
 
                             let kwarg = unsafe { rb_vm_ci_kwarg(ci) };
-                            let processed_args = if !kwarg.is_null() {
+                            let (send_state, processed_args) = if !kwarg.is_null() {
                                 match self.reorder_keyword_arguments(&args, kwarg, iseq) {
-                                    Ok(reordered) => reordered,
+                                    Ok(reordered) => {
+                                        // Only use reordered state if args fit in C registers.
+                                        // Fallback to interpreter needs original order for kwarg handling.
+                                        // NOTE: This needs to match with the condition in codegen.rs.
+                                        if reordered.len() + 1 <= C_ARG_OPNDS.len() {
+                                            let new_state = self.frame_state(state).with_reordered_args(&reordered);
+                                            let snapshot = self.push_insn(block, Insn::Snapshot { state: new_state });
+                                            (snapshot, reordered)
+                                        } else {
+                                            (state, reordered)
+                                        }
+                                    }
                                     Err(reason) => {
                                         self.set_dynamic_send_reason(insn_id, reason);
                                         self.push_insn_id(block, insn_id); continue;
                                     }
                                 }
                             } else {
-                                args.clone()
+                                (state, args.clone())
                             };
 
-                            let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args: processed_args, state });
+                            let send_direct = self.push_insn(block, Insn::SendWithoutBlockDirect { recv, cd, cme, iseq, args: processed_args, state: send_state });
                             self.make_equal_to(insn_id, send_direct);
                         } else if def_type == VM_METHOD_TYPE_IVAR && args.is_empty() {
                             // Check if we're accessing ivars of a Class or Module object as they require single-ractor mode.
@@ -5045,6 +5068,15 @@ impl FrameState {
         state.stack.clear();
         state
     }
+
+    /// Return itself with send args reordered. Used when kwargs are reordered for callee.
+    fn with_reordered_args(&self, reordered_args: &[InsnId]) -> Self {
+        let mut state = self.clone();
+        let args_start = state.stack.len() - reordered_args.len();
+        state.stack.truncate(args_start);
+        state.stack.extend_from_slice(reordered_args);
+        state
+    }
 }
 
 /// Print adaptor for [`FrameState`]. See [`PtrPrintMap`].
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
index 3b6f694810..93e59b35fb 100644
--- a/zjit/src/hir/opt_tests.rs
+++ b/zjit/src/hir/opt_tests.rs
@@ -2816,9 +2816,9 @@ mod hir_opt_tests {
           PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
           PatchPoint NoSingletonClass(Object@0x1000)
           v22:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
-          v23:BasicObject = SendWithoutBlockDirect v22, :foo (0x1038), v11, v13
+          v24:BasicObject = SendWithoutBlockDirect v22, :foo (0x1038), v11, v13
           CheckInterrupts
-          Return v23
+          Return v24
         ");
     }
 
@@ -2846,9 +2846,9 @@ mod hir_opt_tests {
           PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
           PatchPoint NoSingletonClass(Object@0x1000)
           v24:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
-          v25:BasicObject = SendWithoutBlockDirect v24, :foo (0x1038), v13, v15, v11
+          v26:BasicObject = SendWithoutBlockDirect v24, :foo (0x1038), v13, v15, v11
           CheckInterrupts
-          Return v25
+          Return v26
         ");
     }
 
@@ -2876,9 +2876,9 @@ mod hir_opt_tests {
           PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
           PatchPoint NoSingletonClass(Object@0x1000)
           v24:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
-          v25:BasicObject = SendWithoutBlockDirect v24, :foo (0x1038), v11, v15, v13
+          v26:BasicObject = SendWithoutBlockDirect v24, :foo (0x1038), v11, v15, v13
           CheckInterrupts
-          Return v25
+          Return v26
         ");
     }
 
@@ -2933,16 +2933,16 @@ mod hir_opt_tests {
           PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
           PatchPoint NoSingletonClass(Object@0x1000)
           v37:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
-          v38:BasicObject = SendWithoutBlockDirect v37, :foo (0x1038), v11, v13, v15
+          v39:BasicObject = SendWithoutBlockDirect v37, :foo (0x1038), v11, v13, v15
           v20:Fixnum[1] = Const Value(1)
           v22:Fixnum[2] = Const Value(2)
           v24:Fixnum[4] = Const Value(4)
           v26:Fixnum[3] = Const Value(3)
           PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010)
           PatchPoint NoSingletonClass(Object@0x1000)
-          v41:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
-          v42:BasicObject = SendWithoutBlockDirect v41, :foo (0x1038), v20, v22, v26, v24
-          v30:ArrayExact = NewArray v38, v42
+          v42:HeapObject[class_exact*:Object@VALUE(0x1000)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1000)]
+          v44:BasicObject = SendWithoutBlockDirect v42, :foo (0x1038), v20, v22, v26, v24
+          v30:ArrayExact = NewArray v39, v44
           CheckInterrupts
           Return v30
         ");
diff --git a/zjit/src/hir/tests.rs b/zjit/src/hir/tests.rs
index 4528901e9f..44dceba27f 100644
--- a/zjit/src/hir/tests.rs
+++ b/zjit/src/hir/tests.rs
@@ -79,11 +79,12 @@ mod snapshot_tests {
           PatchPoint MethodRedefined(Object@0x1010, foo@0x1018, cme:0x1020)
           PatchPoint NoSingletonClass(Object@0x1010)
           v24:HeapObject[class_exact*:Object@VALUE(0x1010)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1010)]
-          v25:BasicObject = SendWithoutBlockDirect v24, :foo (0x1048), v13, v15, v11
-          v18:Any = Snapshot FrameState { pc: 0x1050, stack: [v25], locals: [] }
+          v25:Any = Snapshot FrameState { pc: 0x1008, stack: [v6, v13, v15, v11], locals: [] }
+          v26:BasicObject = SendWithoutBlockDirect v24, :foo (0x1048), v13, v15, v11
+          v18:Any = Snapshot FrameState { pc: 0x1050, stack: [v26], locals: [] }
           PatchPoint NoTracePoint
           CheckInterrupts
-          Return v25
+          Return v26
         ");
     }
 
@@ -113,11 +114,12 @@ mod snapshot_tests {
           PatchPoint MethodRedefined(Object@0x1010, foo@0x1018, cme:0x1020)
           PatchPoint NoSingletonClass(Object@0x1010)
           v22:HeapObject[class_exact*:Object@VALUE(0x1010)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1010)]
-          v23:BasicObject = SendWithoutBlockDirect v22, :foo (0x1048), v11, v13
-          v16:Any = Snapshot FrameState { pc: 0x1050, stack: [v23], locals: [] }
+          v23:Any = Snapshot FrameState { pc: 0x1008, stack: [v6, v11, v13], locals: [] }
+          v24:BasicObject = SendWithoutBlockDirect v22, :foo (0x1048), v11, v13
+          v16:Any = Snapshot FrameState { pc: 0x1050, stack: [v24], locals: [] }
           PatchPoint NoTracePoint
           CheckInterrupts
-          Return v23
+          Return v24
         ");
     }