ZJIT: Prevent specialization of splats instead of side-exiting (#15005)

2 files changed, 21 insertions, 4 deletions

diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index ea8dacd40b..d82d5837fe 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -2253,6 +2253,14 @@ impl Function {
                             (recv_type.class(), Some(recv_type))
                         };
                         let ci = unsafe { get_call_data_ci(cd) }; // info about the call site
+
+                        // If the call site info indicates that the `Function` has `VM_CALL_ARGS_SPLAT` set, then
+                        // do not optimize into a `SendWithoutBlockDirect`.
+                        let flags = unsafe { rb_vm_ci_flag(ci) };
+                        if unspecializable_call_type(flags) {
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+
                         let mid = unsafe { vm_ci_mid(ci) };
                         // Do method lookup
                         let mut cme = unsafe { rb_callable_method_entry(klass, mid) };
@@ -2676,7 +2684,7 @@ impl Function {
 
                     // When seeing &block argument, fall back to dynamic dispatch for now
                     // TODO: Support block forwarding
-                    if ci_flags & VM_CALL_ARGS_BLOCKARG != 0 {
+                    if unspecializable_call_type(ci_flags) {
                         return Err(());
                     }
 
@@ -4149,12 +4157,17 @@ fn num_locals(iseq: *const rb_iseq_t) -> usize {
 
 /// If we can't handle the type of send (yet), bail out.
 fn unhandled_call_type(flags: u32) -> Result<(), CallType> {
-    if (flags & VM_CALL_ARGS_SPLAT) != 0 { return Err(CallType::Splat); }
     if (flags & VM_CALL_KWARG) != 0 { return Err(CallType::Kwarg); }
     if (flags & VM_CALL_TAILCALL) != 0 { return Err(CallType::Tailcall); }
     Ok(())
 }
 
+/// If a given call uses splatting or block arguments, then we won't specialize.
+fn unspecializable_call_type(flags: u32) -> bool {
+    ((flags & VM_CALL_ARGS_SPLAT) != 0) ||
+    ((flags & VM_CALL_ARGS_BLOCKARG) != 0)
+}
+
 /// We have IseqPayload, which keeps track of HIR Types in the interpreter, but this is not useful
 /// or correct to query from inside the optimizer. Instead, ProfileOracle provides an API to look
 /// up profiled type information by HIR InsnId at a given ISEQ instruction.
diff --git a/zjit/src/hir/tests.rs b/zjit/src/hir/tests.rs
index 32c71dc5d3..c58e14ad4e 100644
--- a/zjit/src/hir/tests.rs
+++ b/zjit/src/hir/tests.rs
@@ -1527,7 +1527,9 @@ pub mod hir_build_tests {
           Jump bb2(v5, v6)
         bb2(v8:BasicObject, v9:BasicObject):
           v14:ArrayExact = ToArray v9
-          SideExit UnhandledCallType(Splat)
+          v16:BasicObject = SendWithoutBlock v8, :foo, v14
+          CheckInterrupts
+          Return v16
         ");
     }
 
@@ -1753,7 +1755,9 @@ pub mod hir_build_tests {
           v14:ArrayExact = ToNewArray v9
           v15:Fixnum[1] = Const Value(1)
           ArrayPush v14, v15
-          SideExit UnhandledCallType(Splat)
+          v19:BasicObject = SendWithoutBlock v8, :foo, v14
+          CheckInterrupts
+          Return v19
         ");
     }