From 95a21ece64040456ed8e2eba16d57d10295e6654 Mon Sep 17 00:00:00 2001 From: Alan Wu Date: Wed, 17 Jan 2024 19:25:56 -0500 Subject: [PATCH] YJIT: Drop extra arguments passed by yield --- yjit/src/codegen.rs | 39 +++++++++++++++++++++++++++++++-------- yjit/src/stats.rs | 1 + 2 files changed, 32 insertions(+), 8 deletions(-) diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs index 378b397fe493f9..b9a30c8c8f9e56 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,16 @@ fn gen_send_iseq( asm.store(rest_param, rest_param_array); } + // Drop extra arguments when yielding + if arg_setup_block { + assert_eq!(0, kw_arg_num); // Checked earlier + 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 +7021,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 place to put 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,