diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb index 5b53e8089bfbf4..0031dde61ef9d7 100644 --- a/bootstraptest/test_yjit.rb +++ b/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) diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs index 378b397fe493f9..d05dd78f4faf5d 100644 --- a/yjit/src/codegen.rs +++ 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) } diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs index 264843117a683b..0ff8195ae24cb6 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_kwargs_req_and_opt_missing,