ZJIT: Check arg limit before pushing SendWithoutBLockDirect insn (#15854)

This reduces some processing and makes the HIR more accurate.

3 files changed, 15 insertions, 29 deletions

diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index d6a2ebec74..4a186d960c 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -401,9 +401,6 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::Send { cd, blockiseq, state, reason, .. } => gen_send(jit, asm, cd, blockiseq, &function.frame_state(state), reason),
         &Insn::SendForward { cd, blockiseq, state, reason, .. } => gen_send_forward(jit, asm, cd, blockiseq, &function.frame_state(state), reason),
         &Insn::SendWithoutBlock { cd, state, reason, .. } => gen_send_without_block(jit, asm, cd, &function.frame_state(state), reason),
-        // Give up SendWithoutBlockDirect for 6+ args since asm.ccall() doesn't support it.
-        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), SendFallbackReason::TooManyArgsForLir),
         Insn::SendWithoutBlockDirect { cme, iseq, recv, args, state, .. } => gen_send_without_block_direct(cb, jit, asm, *cme, *iseq, opnd!(recv), opnds!(args), &function.frame_state(*state)),
         &Insn::InvokeSuper { cd, blockiseq, state, reason, .. } => gen_invokesuper(jit, asm, cd, blockiseq, &function.frame_state(state), reason),
         &Insn::InvokeBlock { cd, state, reason, .. } => gen_invokeblock(jit, asm, cd, &function.frame_state(state), reason),
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 6cdb7993d0..a959ab07f4 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -1662,6 +1662,12 @@ fn can_direct_send(function: &mut Function, block: BlockId, iseq: *const rb_iseq
         return false;
     }
 
+    // asm.ccall() doesn't support 6+ args
+    if args.len() + 1 > C_ARG_OPNDS.len() { // +1 for self
+        function.set_dynamic_send_reason(send_insn, TooManyArgsForLir);
+        return false;
+    }
+
     // Because we exclude e.g. post parameters above, they are also excluded from the sum below.
     let lead_num = params.lead_num;
     let opt_num = params.opt_num;
@@ -2708,16 +2714,9 @@ impl Function {
                             let (send_state, processed_args) = if !kwarg.is_null() {
                                 match self.reorder_keyword_arguments(&args, kwarg, iseq) {
                                     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)
-                                        }
+                                        let new_state = self.frame_state(state).with_reordered_args(&reordered);
+                                        let snapshot = self.push_insn(block, Insn::Snapshot { state: new_state });
+                                        (snapshot, reordered)
                                     }
                                     Err(reason) => {
                                         self.set_dynamic_send_reason(insn_id, reason);
@@ -2768,16 +2767,9 @@ impl Function {
                             let (send_state, processed_args) = if !kwarg.is_null() {
                                 match self.reorder_keyword_arguments(&args, kwarg, iseq) {
                                     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)
-                                        }
+                                        let new_state = self.frame_state(state).with_reordered_args(&reordered);
+                                        let snapshot = self.push_insn(block, Insn::Snapshot { state: new_state });
+                                        (snapshot, reordered)
                                     }
                                     Err(reason) => {
                                         self.set_dynamic_send_reason(insn_id, reason);
diff --git a/zjit/src/hir/tests.rs b/zjit/src/hir/tests.rs
index 44dceba27f..2ae31e862d 100644
--- a/zjit/src/hir/tests.rs
+++ b/zjit/src/hir/tests.rs
@@ -152,14 +152,11 @@ mod snapshot_tests {
           v23:Fixnum[7] = Const Value(7)
           v25:Fixnum[8] = Const Value(8)
           v26:Any = Snapshot FrameState { pc: 0x1008, stack: [v6, v11, v13, v15, v17, v19, v21, v23, v25], locals: [] }
-          PatchPoint MethodRedefined(Object@0x1010, foo@0x1018, cme:0x1020)
-          PatchPoint NoSingletonClass(Object@0x1010)
-          v34:HeapObject[class_exact*:Object@VALUE(0x1010)] = GuardType v6, HeapObject[class_exact*:Object@VALUE(0x1010)]
-          v35:BasicObject = SendWithoutBlockDirect v34, :foo (0x1048), v11, v13, v19, v21, v17, v15, v23, v25
-          v28:Any = Snapshot FrameState { pc: 0x1050, stack: [v35], locals: [] }
+          v27:BasicObject = SendWithoutBlock v6, :foo, v11, v13, v15, v17, v19, v21, v23, v25 # SendFallbackReason: Too many arguments for LIR
+          v28:Any = Snapshot FrameState { pc: 0x1010, stack: [v27], locals: [] }
           PatchPoint NoTracePoint
           CheckInterrupts
-          Return v35
+          Return v27
         ");
     }
 }