ZJIT: Unsupported call feature accounting, and new `send_fallback_fancy_call_feature`

In cases we fall back when the callee has an unsupported signature, it
was a little inaccurate to use `send_fallback_send_not_optimized_method_type`.
We do support the method type in other situations.

Add a new `send_fallback_fancy_call_feature` for these situations. Also,
`send_fallback_bmethod_non_iseq_proc` so we can stop using
`not_optimized_method_type` completely for bmethods.

Add accompanying `fancy_arg_pass_*` counters. These don't sum to the number
of unoptimized calls that run, but establishes the level of support the
optimizer provides for a given workload.

4 files changed, 103 insertions, 17 deletions

diff --git a/zjit.rb b/zjit.rb
index 39e353f327..72a6e58651 100644
--- a/zjit.rb
+++ b/zjit.rb
@@ -164,6 +164,11 @@ class << RubyVM::ZJIT
     print_counters_with_prefix(prefix: 'not_optimized_yarv_insn_', prompt: 'not optimized instructions', buf:, stats:, limit: 20)
     print_counters_with_prefix(prefix: 'send_fallback_', prompt: 'send fallback reasons', buf:, stats:, limit: 20)
 
+    # Show most popular unsupported call features. Because each call can
+    # use multiple fancy features, a decrease in this number does not
+    # necessarily mean an increase in number of optimized calls.
+    print_counters_with_prefix(prefix: 'fancy_arg_pass_', prompt: 'popular unsupported argument-parameter features', buf:, stats:, limit: 10)
+
     # Show exit counters, ordered by the typical amount of exits for the prefix at the time
     print_counters_with_prefix(prefix: 'unhandled_yarv_insn_', prompt: 'unhandled YARV insns', buf:, stats:, limit: 20)
     print_counters_with_prefix(prefix: 'compile_error_', prompt: 'compile error reasons', buf:, stats:, limit: 20)
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index d82d5837fe..013322537e 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -568,6 +568,11 @@ pub enum SendFallbackReason {
     SendNotOptimizedMethodType(MethodType),
     CCallWithFrameTooManyArgs,
     ObjToStringNotString,
+    /// The Proc object for a BMETHOD is not defined by an ISEQ. (See `enum rb_block_type`.)
+    BmethodNonIseqProc,
+    /// The call has at least one feature on the caller or callee side that the optimizer does not
+    /// support.
+    FancyFeatureUse,
     /// Initial fallback reason for every instruction, which should be mutated to
     /// a more actionable reason when an attempt to specialize the instruction fails.
     NotOptimizedInstruction(ruby_vminsn_type),
@@ -1384,14 +1389,22 @@ pub enum ValidationError {
     MiscValidationError(InsnId, String),
 }
 
-fn can_direct_send(iseq: *const rb_iseq_t) -> bool {
-    if unsafe { rb_get_iseq_flags_has_rest(iseq) } { false }
-    else if unsafe { rb_get_iseq_flags_has_opt(iseq) } { false }
-    else if unsafe { rb_get_iseq_flags_has_kw(iseq) } { false }
-    else if unsafe { rb_get_iseq_flags_has_kwrest(iseq) } { false }
-    else if unsafe { rb_get_iseq_flags_has_block(iseq) } { false }
-    else if unsafe { rb_get_iseq_flags_forwardable(iseq) } { false }
-    else { true }
+fn can_direct_send(function: &mut Function, block: BlockId, iseq: *const rb_iseq_t) -> bool {
+    let mut can_send = true;
+    let mut count_failure = |counter| {
+        can_send = false;
+        function.push_insn(block, Insn::IncrCounter(counter));
+    };
+
+    use Counter::*;
+    if unsafe { rb_get_iseq_flags_has_rest(iseq) }    { count_failure(fancy_arg_pass_param_rest) }
+    if unsafe { rb_get_iseq_flags_has_opt(iseq) }     { count_failure(fancy_arg_pass_param_opt) }
+    if unsafe { rb_get_iseq_flags_has_kw(iseq) }      { count_failure(fancy_arg_pass_param_kw) }
+    if unsafe { rb_get_iseq_flags_has_kwrest(iseq) }  { count_failure(fancy_arg_pass_param_kwrest) }
+    if unsafe { rb_get_iseq_flags_has_block(iseq) }   { count_failure(fancy_arg_pass_param_block) }
+    if unsafe { rb_get_iseq_flags_forwardable(iseq) } { count_failure(fancy_arg_pass_param_forwardable) }
+
+    can_send
 }
 
 /// A [`Function`], which is analogous to a Ruby ISeq, is a control-flow graph of [`Block`]s
@@ -2277,8 +2290,8 @@ impl Function {
                             // Only specialize positional-positional calls
                             // TODO(max): Handle other kinds of parameter passing
                             let iseq = unsafe { get_def_iseq_ptr((*cme).def) };
-                            if !can_direct_send(iseq) {
-                                self.set_dynamic_send_reason(insn_id, SendWithoutBlockNotOptimizedMethodType(MethodType::Iseq));
+                            if !can_direct_send(self, block, iseq) {
+                                self.set_dynamic_send_reason(insn_id, FancyFeatureUse);
                                 self.push_insn_id(block, insn_id); continue;
                             }
                             self.push_insn(block, Insn::PatchPoint { invariant: Invariant::MethodRedefined { klass, method: mid, cme }, state });
@@ -2297,21 +2310,18 @@ impl Function {
                             // Target ISEQ bmethods. Can't handle for example, `define_method(:foo, &:foo)`
                             // which makes a `block_type_symbol` bmethod.
                             if proc_block.type_ != block_type_iseq {
-                                self.set_dynamic_send_reason(insn_id, SendWithoutBlockNotOptimizedMethodType(MethodType::Bmethod));
+                                self.set_dynamic_send_reason(insn_id, BmethodNonIseqProc);
                                 self.push_insn_id(block, insn_id); continue;
                             }
                             let capture = unsafe { proc_block.as_.captured.as_ref() };
                             let iseq = unsafe { *capture.code.iseq.as_ref() };
 
-                            if !can_direct_send(iseq) {
-                                self.set_dynamic_send_reason(insn_id, SendWithoutBlockNotOptimizedMethodType(MethodType::Bmethod));
+                            if !can_direct_send(self, block, iseq) {
+                                self.set_dynamic_send_reason(insn_id, FancyFeatureUse);
                                 self.push_insn_id(block, insn_id); continue;
                             }
                             // Can't pass a block to a block for now
-                            if (unsafe { rb_vm_ci_flag(ci) } & VM_CALL_ARGS_BLOCKARG) != 0 {
-                                self.set_dynamic_send_reason(insn_id, SendWithoutBlockNotOptimizedMethodType(MethodType::Bmethod));
-                                self.push_insn_id(block, insn_id); continue;
-                            }
+                            assert!((unsafe { rb_vm_ci_flag(ci) } & VM_CALL_ARGS_BLOCKARG) == 0, "SendWithoutBlock but has a block arg");
 
                             // Patch points:
                             // Check for "defined with an un-shareable Proc in a different Ractor"
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
index d697065da9..f29af66bde 100644
--- a/zjit/src/hir/opt_tests.rs
+++ b/zjit/src/hir/opt_tests.rs
@@ -2416,6 +2416,7 @@ mod hir_opt_tests {
           Jump bb2(v4)
         bb2(v6:BasicObject):
           v10:Fixnum[1] = Const Value(1)
+          IncrCounter fancy_arg_pass_param_opt
           v12:BasicObject = SendWithoutBlock v6, :foo, v10
           CheckInterrupts
           Return v12
@@ -2499,6 +2500,7 @@ mod hir_opt_tests {
           Jump bb2(v4)
         bb2(v6:BasicObject):
           v10:Fixnum[1] = Const Value(1)
+          IncrCounter fancy_arg_pass_param_rest
           v12:BasicObject = SendWithoutBlock v6, :foo, v10
           CheckInterrupts
           Return v12
@@ -2833,6 +2835,9 @@ mod hir_opt_tests {
           v12:NilClass = Const Value(nil)
           PatchPoint MethodRedefined(Hash@0x1008, new@0x1010, cme:0x1018)
           v43:HashExact = ObjectAllocClass Hash:VALUE(0x1008)
+          IncrCounter fancy_arg_pass_param_opt
+          IncrCounter fancy_arg_pass_param_kw
+          IncrCounter fancy_arg_pass_param_block
           v18:BasicObject = SendWithoutBlock v43, :initialize
           CheckInterrupts
           CheckInterrupts
@@ -7022,6 +7027,58 @@ mod hir_opt_tests {
     }
 
     #[test]
+    fn counting_fancy_feature_use_for_fallback() {
+        eval("
+            define_method(:fancy) { |_a, *_b, kw: 100, **kw_rest, &block| }
+            def test = fancy(1)
+            test
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb2(v1)
+        bb1(v4:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v4)
+        bb2(v6:BasicObject):
+          v10:Fixnum[1] = Const Value(1)
+          IncrCounter fancy_arg_pass_param_rest
+          IncrCounter fancy_arg_pass_param_kw
+          IncrCounter fancy_arg_pass_param_kwrest
+          IncrCounter fancy_arg_pass_param_block
+          v12:BasicObject = SendWithoutBlock v6, :fancy, v10
+          CheckInterrupts
+          Return v12
+        ");
+    }
+
+    #[test]
+    fn call_method_forwardable_param() {
+        eval("
+           def forwardable(...) = itself(...)
+           def call_forwardable = forwardable
+           call_forwardable
+        ");
+        assert_snapshot!(hir_string("call_forwardable"), @r"
+        fn call_forwardable@<compiled>:3:
+        bb0():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb2(v1)
+        bb1(v4:BasicObject):
+          EntryPoint JIT(0)
+          Jump bb2(v4)
+        bb2(v6:BasicObject):
+          IncrCounter fancy_arg_pass_param_forwardable
+          v11:BasicObject = SendWithoutBlock v6, :forwardable
+          CheckInterrupts
+          Return v11
+        ");
+    }
+
+    #[test]
     fn test_elide_string_length() {
         eval(r#"
             def test(s)
diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs
index ad027ef593..17ae27ac01 100644
--- a/zjit/src/stats.rs
+++ b/zjit/src/stats.rs
@@ -173,6 +173,10 @@ make_counters! {
         send_fallback_send_no_profiles,
         send_fallback_send_not_optimized_method_type,
         send_fallback_ccall_with_frame_too_many_args,
+        // The call has at least one feature on the caller or callee side
+        // that the optimizer does not support.
+        send_fallback_fancy_call_feature,
+        send_fallback_bmethod_non_iseq_proc,
         send_fallback_obj_to_string_not_string,
         send_fallback_not_optimized_instruction,
     }
@@ -250,6 +254,14 @@ make_counters! {
     unspecialized_send_def_type_refined,
     unspecialized_send_def_type_null,
 
+    // Unsupported parameter features
+    fancy_arg_pass_param_rest,
+    fancy_arg_pass_param_opt,
+    fancy_arg_pass_param_kw,
+    fancy_arg_pass_param_kwrest,
+    fancy_arg_pass_param_block,
+    fancy_arg_pass_param_forwardable,
+
     // Writes to the VM frame
     vm_write_pc_count,
     vm_write_sp_count,
@@ -401,6 +413,8 @@ pub fn send_fallback_counter(reason: crate::hir::SendFallbackReason) -> Counter
         SendWithoutBlockDirectTooManyArgs         => send_fallback_send_without_block_direct_too_many_args,
         SendPolymorphic                           => send_fallback_send_polymorphic,
         SendNoProfiles                            => send_fallback_send_no_profiles,
+        FancyFeatureUse                           => send_fallback_fancy_call_feature,
+        BmethodNonIseqProc                        => send_fallback_bmethod_non_iseq_proc,
         SendNotOptimizedMethodType(_)             => send_fallback_send_not_optimized_method_type,
         CCallWithFrameTooManyArgs                 => send_fallback_ccall_with_frame_too_many_args,
         ObjToStringNotString                      => send_fallback_obj_to_string_not_string,