YJIT: Drop extra arguments passed by yield (#9596)

Support dropping extra arguments passed by `yield` in blocks. For
example `10.times { work }` drops the count argument. This is common
enough that it's about 3% of fallback reasons in `lobsters`.

Only support simple cases where the surplus arguments are at the top of
the stack, that way they just need to be popped, which takes no work.

3 files changed, 101 insertions, 37 deletions

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 3bf715f888..926eaaf8f6 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -1,3 +1,42 @@
+# test discarding extra yield arguments
+assert_equal "2210150001501015", %q{
+  def splat_kw(ary) = yield *ary, a: 1
+
+  def splat(ary) = yield *ary
+
+  def kw = yield 1, 2, a: 0
+
+  def simple = yield 0, 1
+
+  def calls
+    [
+      splat([1, 1, 2]) { |x, y| x + y },
+      splat([1, 1, 2]) { |y, opt = raise| opt + y},
+      splat_kw([0, 1]) { |a:| a },
+      kw { |a:| a },
+      kw { |a| a },
+      simple { 5.itself },
+      simple { |a| a },
+      simple { |opt = raise| opt },
+      simple { |*rest| rest },
+      simple { |opt_kw: 5| opt_kw },
+      # autosplat ineractions
+      [0, 1, 2].yield_self { |a, b| [a, b] },
+      [0, 1, 2].yield_self { |a, opt = raise| [a, opt] },
+      [1].yield_self { |a, opt = 4| a + opt },
+    ]
+  end
+
+  calls.join
+}
+
+# test autosplat with empty splat
+assert_equal "ok", %q{
+  def m(pos, splat) = yield pos, *splat
+
+  m([:ok], []) {|v0,| v0 }
+}
+
 # regression test for send stack shifting
 assert_normal_exit %q{
   def foo(a, b)
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index b50cd10163..fcb4c47608 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -6117,6 +6117,7 @@ fn gen_send_iseq(
     let supplying_kws = unsafe { vm_ci_flag(ci) & VM_CALL_KWARG } != 0;
     let iseq_has_rest = unsafe { get_iseq_flags_has_rest(iseq) };
     let iseq_has_block_param = unsafe { get_iseq_flags_has_block(iseq) };
+    let arg_setup_block = captured_opnd.is_some(); // arg_setup_type: arg_setup_block (invokeblock)
 
     // For computing offsets to callee locals
     let num_params = unsafe { get_iseq_body_param_size(iseq) };
@@ -6137,10 +6138,10 @@ fn gen_send_iseq(
     // Arity handling and optional parameter setup
     let mut opts_filled = argc - required_num - kw_arg_num;
     let opt_num = unsafe { get_iseq_body_param_opt_num(iseq) };
-    // We have a rest parameter so there could be more args
-    // than are required + optional. Those will go in rest.
+    // With a rest parameter or a yield to a block,
+    // callers can pass more than required + optional.
     // So we cap ops_filled at opt_num.
-    if iseq_has_rest {
+    if iseq_has_rest || arg_setup_block {
         opts_filled = min(opts_filled, opt_num);
     }
     let mut opts_missing: i32 = opt_num - opts_filled;
@@ -6159,11 +6160,17 @@ fn gen_send_iseq(
     exit_if_supplying_kws_and_accept_no_kwargs(asm, supplying_kws, iseq)?;
     exit_if_splat_and_zsuper(asm, flags)?;
     exit_if_doing_kw_and_splat(asm, doing_kw_call, flags)?;
-    exit_if_wrong_number_arguments(asm, opts_filled, flags, opt_num, iseq_has_rest)?;
+    exit_if_wrong_number_arguments(asm, arg_setup_block, opts_filled, flags, opt_num, iseq_has_rest)?;
     exit_if_doing_kw_and_opts_missing(asm, doing_kw_call, opts_missing)?;
     exit_if_has_rest_and_optional_and_block(asm, iseq_has_rest, opt_num, iseq, block_arg)?;
     let block_arg_type = exit_if_unsupported_block_arg_type(jit, asm, block_arg)?;
 
+    // Bail if we can't drop extra arguments for a yield by just popping them
+    if supplying_kws && arg_setup_block && argc > (kw_arg_num + required_num + opt_num) {
+        gen_counter_incr(asm, Counter::send_iseq_complex_discard_extras);
+        return None;
+    }
+
     // Block parameter handling. This mirrors setup_parameters_complex().
     if iseq_has_block_param {
         if unsafe { get_iseq_body_local_iseq(iseq) == iseq } {
@@ -6249,35 +6256,6 @@ fn gen_send_iseq(
         }
     }
 
-    // Check if we need the arg0 splat handling of vm_callee_setup_block_arg()
-    // Also known as "autosplat" inside setup_parameters_complex()
-    let arg_setup_block = captured_opnd.is_some(); // arg_setup_type: arg_setup_block (invokeblock)
-    let block_arg0_splat = arg_setup_block && argc == 1 && unsafe {
-         (get_iseq_flags_has_lead(iseq) || opt_num > 1)
-             && !get_iseq_flags_ambiguous_param0(iseq)
-    };
-    if block_arg0_splat {
-        // If block_arg0_splat, we still need side exits after splat, but
-        // doing push_splat_args here disallows it. So bail out.
-        if flags & VM_CALL_ARGS_SPLAT != 0 && !iseq_has_rest {
-            gen_counter_incr(asm, Counter::invokeblock_iseq_arg0_args_splat);
-            return None;
-        }
-        // The block_arg0_splat implementation is for the rb_simple_iseq_p case,
-        // but doing_kw_call means it's not a simple ISEQ.
-        if doing_kw_call {
-            gen_counter_incr(asm, Counter::invokeblock_iseq_arg0_has_kw);
-            return None;
-        }
-        // The block_arg0_splat implementation cannot deal with optional parameters.
-        // This is a setup_parameters_complex() situation and interacts with the
-        // starting position of the callee.
-        if opt_num > 1 {
-            gen_counter_incr(asm, Counter::invokeblock_iseq_arg0_optional);
-            return None;
-        }
-    }
-
     let splat_array_length = if flags & VM_CALL_ARGS_SPLAT != 0 {
         let array = jit.peek_at_stack(&asm.ctx, if block_arg { 1 } else { 0 }) ;
         let array_length = if array == Qnil {
@@ -6318,6 +6296,33 @@ fn gen_send_iseq(
         None
     };
 
+    // Check if we need the arg0 splat handling of vm_callee_setup_block_arg()
+    // Also known as "autosplat" inside setup_parameters_complex().
+    // Autosplat checks argc == 1 after splat and kwsplat processing, so make
+    // sure to amend this if we start support kw_splat.
+    let block_arg0_splat = arg_setup_block
+        && (argc == 1 || (argc == 2 && splat_array_length == Some(0)))
+        && !supplying_kws && !doing_kw_call
+        && unsafe {
+            (get_iseq_flags_has_lead(iseq) || opt_num > 1)
+                && !get_iseq_flags_ambiguous_param0(iseq)
+        };
+    if block_arg0_splat {
+        // If block_arg0_splat, we still need side exits after splat, but
+        // the splat modifies the stack which breaks side exits. So bail out.
+        if flags & VM_CALL_ARGS_SPLAT != 0 {
+            gen_counter_incr(asm, Counter::invokeblock_iseq_arg0_args_splat);
+            return None;
+        }
+        // The block_arg0_splat implementation cannot deal with optional parameters.
+        // This is a setup_parameters_complex() situation and interacts with the
+        // starting position of the callee.
+        if opt_num > 1 {
+            gen_counter_incr(asm, Counter::invokeblock_iseq_arg0_optional);
+            return None;
+        }
+    }
+
     // Adjust `opts_filled` and `opts_missing` taking
     // into account the size of the splat expansion.
     if let Some(len) = splat_array_length {
@@ -6605,6 +6610,19 @@ fn gen_send_iseq(
         asm.store(rest_param, rest_param_array);
     }
 
+    // Pop surplus positional arguments when yielding
+    if arg_setup_block {
+        let extras = argc - required_num - opt_num;
+        if extras > 0 {
+            // Checked earlier. If there are keyword args, then
+            // the positional arguments are not at the stack top.
+            assert_eq!(0, kw_arg_num);
+
+            asm.stack_pop(extras as usize);
+            argc = required_num + opt_num;
+        }
+    }
+
     if doing_kw_call {
         // Here we're calling a method with keyword arguments and specifying
         // keyword arguments at this call site.
@@ -7034,11 +7052,18 @@ fn exit_if_doing_kw_and_splat(asm: &mut Assembler, doing_kw_call: bool, flags: u
 }
 
 #[must_use]
-fn exit_if_wrong_number_arguments(asm: &mut Assembler, opts_filled: i32, flags: u32, opt_num: i32, iseq_has_rest: bool) -> Option<()> {
+fn exit_if_wrong_number_arguments(
+    asm: &mut Assembler,
+    args_setup_block: bool,
+    opts_filled: i32,
+    flags: u32,
+    opt_num: i32,
+    iseq_has_rest: bool,
+) -> Option<()> {
     // Too few arguments and no splat to make up for it
     let too_few = opts_filled < 0 && flags & VM_CALL_ARGS_SPLAT == 0;
-    // Too many arguments and no place to put them (i.e. rest arg)
-    let too_many = opts_filled > opt_num && !iseq_has_rest;
+    // Too many arguments and no sink that take them
+    let too_many = opts_filled > opt_num && !(iseq_has_rest || args_setup_block);
 
     exit_if(asm, too_few || too_many, Counter::send_iseq_arity_error)
 }
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 1dd5f57b2b..7df01448a4 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -331,6 +331,7 @@ make_counters! {
     send_iseq_arity_error,
     send_iseq_block_arg_type,
     send_iseq_clobbering_block_arg,
+    send_iseq_complex_discard_extras,
     send_iseq_leaf_builtin_block_arg_block_param,
     send_iseq_only_keywords,
     send_iseq_kw_splat,
@@ -391,7 +392,6 @@ make_counters! {
     invokeblock_megamorphic,
     invokeblock_none,
     invokeblock_iseq_arg0_optional,
-    invokeblock_iseq_arg0_has_kw,
     invokeblock_iseq_arg0_args_splat,
     invokeblock_iseq_arg0_not_array,
     invokeblock_iseq_arg0_wrong_len,