YJIT: Drop extra arguments passed by yield

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 take no work.

bootstraptest/test_yjit.rb

@@ -1,3 +1,35 @@
+# 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
+}
+
 # regression test for send stack shifting
 assert_normal_exit %q{
   def foo(a, b)

yjit/src/codegen.rs

@@ -6097,6 +6097,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) };
@@ -6117,10 +6118,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;
@@ -6138,11 +6139,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 } {
@@ -6230,7 +6237,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)
@@ -6581,6 +6587,19 @@ fn gen_send_iseq(
         asm.store(rest_param, rest_param_array);
     }
 
+    // Pop surplus positional arguments when yielding
+    if arg_setup_block {
+        // Checked earlier. If there are keyword args, then
+        // the positional arguments are not at the stack top.
+        assert_eq!(0, kw_arg_num);
+
+        let extras = argc - required_num - opt_num;
+        if extras > 0 {
+            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.
@@ -7005,11 +7024,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)
 }

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_kwargs_req_and_opt_missing,