ZJIT: Count caller side features for `complex_arg_pass`

After 34b0ac68b31, we use a fallback instead of side exit for splats.
Count splats under `send_fallback_one_or_more_complex_arg_pass`.

2 files changed, 55 insertions, 1 deletions

diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 09e293a0f6..b26d2ffa04 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -2290,6 +2290,8 @@ impl Function {
                         // do not optimize into a `SendWithoutBlockDirect`.
                         let flags = unsafe { rb_vm_ci_flag(ci) };
                         if unspecializable_call_type(flags) {
+                            self.count_complex_call_features(block, flags);
+                            self.set_dynamic_send_reason(insn_id, ComplexArgPass);
                             self.push_insn_id(block, insn_id); continue;
                         }
 
@@ -2770,6 +2772,8 @@ impl Function {
                     // When seeing &block argument, fall back to dynamic dispatch for now
                     // TODO: Support block forwarding
                     if unspecializable_call_type(ci_flags) {
+                        fun.count_complex_call_features(block, ci_flags);
+                        fun.set_dynamic_send_reason(send_insn_id, ComplexArgPass);
                         return Err(());
                     }
 
@@ -2877,6 +2881,7 @@ impl Function {
                     // Filter for simple call sites (i.e. no splats etc.)
                     if ci_flags & VM_CALL_ARGS_SIMPLE == 0 {
                         fun.count_complex_call_features(block, ci_flags);
+                        fun.set_dynamic_send_reason(send_insn_id, ComplexArgPass);
                         return Err(());
                     }
 
@@ -2948,7 +2953,10 @@ impl Function {
                     // func(int argc, VALUE *argv, VALUE recv)
                     let ci_flags = unsafe { vm_ci_flag(call_info) };
                     if ci_flags & VM_CALL_ARGS_SIMPLE == 0 {
-                        fun.count_complex_call_features(block, ci_flags);
+                        // TODO(alan): Add fun.count_complex_call_features() here without double
+                        // counting splat
+                        fun.set_dynamic_send_reason(send_insn_id, ComplexArgPass);
+                        return Err(());
                     } else {
                         fun.gen_patch_points_for_optimized_ccall(block, recv_class, method_id, cme, state);
 
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
index 97b93e3d7b..6fc890f385 100644
--- a/zjit/src/hir/opt_tests.rs
+++ b/zjit/src/hir/opt_tests.rs
@@ -4815,6 +4815,7 @@ mod hir_opt_tests {
           v14:ArrayExact = NewArray
           GuardBlockParamProxy l0
           v17:HeapObject[BlockParamProxy] = Const Value(VALUE(0x1000))
+          IncrCounter complex_arg_pass_caller_blockarg
           v19:BasicObject = Send v14, 0x1008, :map, v17
           CheckInterrupts
           Return v19
@@ -6842,6 +6843,51 @@ mod hir_opt_tests {
     }
 
     #[test]
+    fn test_splat() {
+        eval("
+            def foo = itself
+
+            def test
+              # Use a local to inhibit compile.c peephole optimization to ensure callsites have VM_CALL_ARGS_SPLAT
+              empty = []
+              foo(*empty)
+              ''.display(*empty)
+              itself(*empty)
+            end
+            test
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:6:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          v2:NilClass = Const Value(nil)
+          Jump bb2(v1, v2)
+        bb1(v5:BasicObject):
+          EntryPoint JIT(0)
+          v6:NilClass = Const Value(nil)
+          Jump bb2(v5, v6)
+        bb2(v8:BasicObject, v9:NilClass):
+          v14:ArrayExact = NewArray
+          v18:ArrayExact = ToArray v14
+          IncrCounter complex_arg_pass_caller_splat
+          v20:BasicObject = SendWithoutBlock v8, :foo, v18
+          v23:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+          v25:StringExact = StringCopy v23
+          PatchPoint NoEPEscape(test)
+          v29:ArrayExact = ToArray v14
+          IncrCounter complex_arg_pass_caller_splat
+          v31:BasicObject = SendWithoutBlock v25, :display, v29
+          PatchPoint NoEPEscape(test)
+          v37:ArrayExact = ToArray v14
+          IncrCounter complex_arg_pass_caller_splat
+          v39:BasicObject = SendWithoutBlock v8, :itself, v37
+          CheckInterrupts
+          Return v39
+        ");
+    }
+
+    #[test]
     fn test_inline_symbol_to_sym() {
         eval(r#"
             def test(o) = o.to_sym