diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb index 0207f4aaf16577..5347028550dfca 100644 --- a/test/ruby/test_yjit.rb +++ b/test/ruby/test_yjit.rb @@ -1320,6 +1320,16 @@ def test_opt_aref_with RUBY end + def test_proc_block_arg + assert_compiles(<<~RUBY, result: [:proc, :no_block]) + def yield_if_given = block_given? ? yield : :no_block + + def call(block_arg = nil) = yield_if_given(&block_arg) + + [call(-> { :proc }), call] + RUBY + end + private def code_gc_helpers diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs index baa71fc6cf3fc0..efb7e9023985db 100644 --- a/yjit/bindgen/src/main.rs +++ b/yjit/bindgen/src/main.rs @@ -447,8 +447,8 @@ fn main() { .allowlist_function("rb_method_basic_definition_p") .allowlist_function("rb_yjit_array_len") .allowlist_function("rb_obj_class") - .allowlist_function("rb_obj_is_proc") .allowlist_function("rb_vm_base_ptr") + .allowlist_function("rb_obj_is_proc") // We define VALUE manually, don't import it .blocklist_type("VALUE") diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs index 8e83f128af733f..3461b1dd5004ff 100644 --- a/yjit/src/codegen.rs +++ b/yjit/src/codegen.rs @@ -5068,6 +5068,9 @@ enum SpecVal { BlockParamProxy, PrevEP(*const VALUE), PrevEPOpnd(Opnd), + // To avoid holding values across C calls for complex calls, + // we might need to set the SpecVal earlier in the call sequence + AlreadySet, } // Each variant represents a branch in vm_caller_setup_arg_block. @@ -5160,9 +5163,14 @@ fn gen_push_frame( } SpecVal::PrevEPOpnd(ep_opnd) => { asm.or(ep_opnd, 1.into()) - }, + } + SpecVal::AlreadySet => 0.into(), // unused }; - asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -2), specval); + if let SpecVal::AlreadySet = frame.specval { + asm.comment("specval should been set"); + } else { + asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -2), specval); + } // Write env flags at sp[-1] // sp[-1] = frame_type; @@ -5794,13 +5802,7 @@ fn gen_send_iseq( } let mut opts_missing: i32 = opt_num - opts_filled; - let block_arg = flags & VM_CALL_ARGS_BLOCKARG != 0; - let block_arg_type = if block_arg { - Some(asm.ctx.get_opnd_type(StackOpnd(0))) - } else { - None - }; exit_if_stack_too_large(iseq)?; exit_if_tail_call(asm, ci)?; @@ -5816,7 +5818,7 @@ fn gen_send_iseq( exit_if_wrong_number_arguments(asm, 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)?; - exit_if_unsupported_block_arg_type(asm, block_arg_type)?; + let block_arg_type = exit_if_unsupported_block_arg_type(jit, asm, block_arg)?; // Block parameter handling. This mirrors setup_parameters_complex(). if iseq_has_block_param { @@ -6003,12 +6005,35 @@ fn gen_send_iseq( // 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("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_CHAIN_DEPTH, + Counter::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 } - _ => { - assert!(false); - } + _ => unreachable!(), } let builtin_attrs = unsafe { rb_yjit_iseq_builtin_attrs(iseq) }; @@ -6434,6 +6459,8 @@ fn gen_send_iseq( SpecVal::PrevEPOpnd(ep_opnd) } else if block_arg_type == Some(Type::BlockParamProxy) { SpecVal::BlockParamProxy + } else if let Some(Type::TProc) = block_arg_type { + SpecVal::AlreadySet } else if let Some(block_handler) = block { SpecVal::BlockHandler(block_handler) } else { @@ -6642,20 +6669,35 @@ fn exit_if_has_rest_and_optional_and_block(asm: &mut Assembler, iseq_has_rest: b } #[must_use] -fn exit_if_unsupported_block_arg_type(asm: &mut Assembler, block_arg_type: Option) -> Option<()> { +fn exit_if_unsupported_block_arg_type( + jit: &mut JITState, + asm: &mut Assembler, + supplying_block_arg: bool +) -> Option> { + let block_arg_type = if supplying_block_arg { + asm.ctx.get_opnd_type(StackOpnd(0)) + } else { + // Passing no block argument + return Some(None); + }; + match block_arg_type { - Some(Type::Nil | Type::BlockParamProxy) => { + Type::Nil | Type::BlockParamProxy => { // We'll handle this later + Some(Some(block_arg_type)) } - None => { - // Nothing to do + _ if { + let sample_block_arg = jit.peek_at_stack(&asm.ctx, 0); + unsafe { rb_obj_is_proc(sample_block_arg) }.test() + } => { + // Speculate that we'll have a proc as the block arg + Some(Some(Type::TProc)) } _ => { gen_counter_incr(asm, Counter::send_block_arg); - return None + None } } - Some(()) } #[must_use] diff --git a/yjit/src/core.rs b/yjit/src/core.rs index 89061ac4e5ecec..40646ebd342fdf 100644 --- a/yjit/src/core.rs +++ b/yjit/src/core.rs @@ -59,6 +59,8 @@ pub enum Type { TArray, // An object with the T_ARRAY flag set, possibly an rb_cArray CArray, // An un-subclassed string of type rb_cArray (can have instance vars in some cases) + TProc, // A proc object. Could be an instance of a subclass of ::rb_cProc + BlockParamProxy, // A special sentinel value indicating the block parameter should be read from // the current surrounding cfp } @@ -110,6 +112,8 @@ impl Type { RUBY_T_ARRAY => Type::TArray, RUBY_T_HASH => Type::Hash, RUBY_T_STRING => Type::TString, + #[cfg(not(test))] + RUBY_T_DATA if unsafe { rb_obj_is_proc(val).test() } => Type::TProc, _ => Type::UnknownHeap, } } @@ -155,6 +159,7 @@ impl Type { Type::TString => true, Type::CString => true, Type::BlockParamProxy => true, + Type::TProc => true, _ => false, } } @@ -189,6 +194,7 @@ impl Type { Type::Hash => Some(RUBY_T_HASH), Type::ImmSymbol | Type::HeapSymbol => Some(RUBY_T_SYMBOL), Type::TString | Type::CString => Some(RUBY_T_STRING), + Type::TProc => Some(RUBY_T_DATA), Type::Unknown | Type::UnknownImm | Type::UnknownHeap => None, Type::BlockParamProxy => None, } diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs index ef4d25d051a35c..1afe417fc2ff99 100644 --- a/yjit/src/stats.rs +++ b/yjit/src/stats.rs @@ -210,6 +210,7 @@ make_counters! { send_args_splat_super, send_iseq_zsuper, send_block_arg, + send_block_arg_type, send_ivar_set_method, send_zsuper_method, send_undef_method,