From 0bf1749e9fcef24bf7bebbce2a62ee6c766d4c7c Mon Sep 17 00:00:00 2001 From: Alan Wu Date: Fri, 13 Oct 2023 10:41:53 -0400 Subject: [PATCH] YJIT: Fix argument clobbering in some block_arg+rest_param calls (#8647) Previously, for block argument callsites with some specific argument count and callee local variable count combinations, YJIT ended up writing over arguments that are supposed to be collected into a rest parameter array unmodified. Detect when clobbering would happen and avoid it. Also, place the block handler after the stack overflow check, since it writes to new stack space. Reported-by: Takashi Kokubun --- bootstraptest/test_yjit.rb | 12 +++++ yjit/src/codegen.rs | 100 +++++++++++++++++++++---------------- yjit/src/stats.rs | 2 + 3 files changed, 72 insertions(+), 42 deletions(-) diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb index 999c31fc37..be99d580c1 100644 --- a/bootstraptest/test_yjit.rb +++ b/bootstraptest/test_yjit.rb @@ -1,3 +1,15 @@ +# regression test for callee block handler overlapping with arguments +assert_equal '3', %q{ + def foo(_req, *args) = args.last + + def call_foo = foo(0, 1, 2, 3, &->{}) + + call_foo +} + +# call leaf builtin with a block argument +assert_equal '0', "0.abs(&nil)" + # regression test for invokeblock iseq guard assert_equal 'ok', %q{ return :ok unless defined?(GC.compact) diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs index 26049abfbd..2707932a3f 100644 --- a/yjit/src/codegen.rs +++ b/yjit/src/codegen.rs @@ -6043,48 +6043,7 @@ fn gen_send_iseq( } } - // We will not have None from here. You can use stack_pop / stack_pop. - - match block_arg_type { - Some(Type::Nil) => { - // We have a nil block arg, so let's pop it off the args - asm.stack_pop(1); - } - Some(Type::BlockParamProxy) => { - // We don't need the actual stack value - asm.stack_pop(1); - } - Some(Type::TProc) => { - // Place the proc as the block handler. We do this early because - // the block arg being at the top of the stack gets in the way of - // rest param handling later. Also, since there are C calls that - // come later, we can't hold this value in a register and place it - // near the end when we push a new control frame. - asm_comment!(asm, "guard block arg is a proc"); - // Simple predicate, no need for jit_prepare_routine_call(). - let is_proc = asm.ccall(rb_obj_is_proc as _, vec![asm.stack_opnd(0)]); - asm.cmp(is_proc, Qfalse.into()); - jit_chain_guard( - JCC_JE, - jit, - asm, - ocb, - SEND_MAX_DEPTH, - Counter::guard_send_block_arg_type, - ); - - let proc = asm.stack_pop(1); - let callee_ep = -argc + num_locals + VM_ENV_DATA_SIZE as i32 - 1; - let callee_specval = callee_ep + VM_ENV_DATA_INDEX_SPECVAL; - let callee_specval = asm.ctx.sp_opnd(callee_specval as isize * SIZEOF_VALUE as isize); - asm.store(callee_specval, proc); - } - None => { - // Nothing to do - } - _ => unreachable!(), - } - + // Shortcut for special `Primitive.attr! :leaf` builtins let builtin_attrs = unsafe { rb_yjit_iseq_builtin_attrs(iseq) }; let builtin_func_raw = unsafe { rb_yjit_builtin_function(iseq) }; let builtin_func = if builtin_func_raw.is_null() { None } else { Some(builtin_func_raw) }; @@ -6094,6 +6053,18 @@ fn gen_send_iseq( if builtin_argc + 1 < (C_ARG_OPNDS.len() as i32) { asm_comment!(asm, "inlined leaf builtin"); + // We pop the block arg without using it because: + // - the builtin is leaf, so it promises to not `yield`. + // - no leaf builtins have block param at the time of writing, and + // adding one requires interpreter changes to support. + if block_arg_type.is_some() { + if iseq_has_block_param { + gen_counter_incr(asm, Counter::send_iseq_leaf_builtin_block_arg_block_param); + return None; + } + asm.stack_pop(1); + } + // Skip this if it doesn't trigger GC if builtin_attrs & BUILTIN_ATTR_NO_GC == 0 { // The callee may allocate, e.g. Integer#abs on a Bignum. @@ -6135,6 +6106,51 @@ fn gen_send_iseq( asm.cmp(CFP, stack_limit); asm.jbe(Target::side_exit(Counter::guard_send_se_cf_overflow)); + match block_arg_type { + Some(Type::Nil) => { + // We have a nil block arg, so let's pop it off the args + asm.stack_pop(1); + } + Some(Type::BlockParamProxy) => { + // We don't need the actual stack value + asm.stack_pop(1); + } + Some(Type::TProc) => { + // Place the proc as the block handler. We do this early because + // the block arg being at the top of the stack gets in the way of + // rest param handling later. Also, since there are C calls that + // come later, we can't hold this value in a register and place it + // near the end when we push a new control frame. + asm_comment!(asm, "guard block arg is a proc"); + // Simple predicate, no need for jit_prepare_routine_call(). + let is_proc = asm.ccall(rb_obj_is_proc as _, vec![asm.stack_opnd(0)]); + asm.cmp(is_proc, Qfalse.into()); + jit_chain_guard( + JCC_JE, + jit, + asm, + ocb, + SEND_MAX_DEPTH, + Counter::guard_send_block_arg_type, + ); + + let callee_ep = -argc + num_locals + VM_ENV_DATA_SIZE as i32 - 1; + let callee_specval = callee_ep + VM_ENV_DATA_INDEX_SPECVAL; + if callee_specval < 0 { + // Can't write to sp[-n] since that's where the arguments are + gen_counter_incr(asm, Counter::send_iseq_clobbering_block_arg); + return None; + } + let proc = asm.stack_pop(1); // Pop first, as argc doesn't account for the block arg + let callee_specval = asm.ctx.sp_opnd(callee_specval as isize * SIZEOF_VALUE as isize); + asm.store(callee_specval, proc); + } + None => { + // Nothing to do + } + _ => unreachable!(), + } + // push_splat_args does stack manipulation so we can no longer side exit if let Some(array_length) = splat_array_length { if !iseq_has_rest { diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs index c9b21c2ae9..0bcae1cdc9 100644 --- a/yjit/src/stats.rs +++ b/yjit/src/stats.rs @@ -267,6 +267,8 @@ make_counters! { send_attrset_kwargs, send_iseq_tailcall, send_iseq_arity_error, + send_iseq_clobbering_block_arg, + send_iseq_leaf_builtin_block_arg_block_param, send_iseq_only_keywords, send_iseq_kwargs_req_and_opt_missing, send_iseq_kwargs_mismatch,