YJIT: Support nil and blockparamproxy as blockarg in send (#6492)

Co-authored-by: John Hawthorn <john@hawthorn.email>

Co-authored-by: John Hawthorn <john@hawthorn.email>
This commit is contained in:
Matthew Draper 2022-10-27 05:57:59 +10:30 committed by GitHub
parent fa0adbad92
commit c746f380f2
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
Notes: git 2022-10-26 19:28:28 +00:00
Merged-By: maximecb <maximecb@ruby-lang.org>
7 changed files with 221 additions and 47 deletions

View File

@ -3337,3 +3337,46 @@ assert_equal '[[1, nil, 2]]', %q{
5.times.map { opt_and_kwargs(1, c: 2) }.uniq
}
# bmethod with forwarded block
assert_equal '2', %q{
define_method(:foo) do |&block|
block.call
end
def bar(&block)
foo(&block)
end
bar { 1 }
bar { 2 }
}
# bmethod with forwarded block and arguments
assert_equal '5', %q{
define_method(:foo) do |n, &block|
n + block.call
end
def bar(n, &block)
foo(n, &block)
end
bar(0) { 1 }
bar(3) { 2 }
}
# bmethod with forwarded unwanted block
assert_equal '1', %q{
one = 1
define_method(:foo) do
one
end
def bar(&block)
foo(&block)
end
bar { }
bar { }
}

View File

@ -513,9 +513,8 @@ class TestYJIT < Test::Unit::TestCase
RUBY
end
def test_getblockparamproxy_with_no_block
# Currently side exits on the send
assert_compiles(<<~'RUBY', insns: [:getblockparamproxy], exits: { send: 2 })
def test_send_blockarg
assert_compiles(<<~'RUBY', insns: [:getblockparamproxy, :send], exits: {})
def bar
end
@ -526,6 +525,9 @@ class TestYJIT < Test::Unit::TestCase
foo
foo
foo { }
foo { }
RUBY
end

6
yjit.c
View File

@ -592,6 +592,12 @@ rb_get_iseq_body_local_iseq(const rb_iseq_t *iseq)
return iseq->body->local_iseq;
}
const rb_iseq_t *
rb_get_iseq_body_parent_iseq(const rb_iseq_t *iseq)
{
return iseq->body->parent_iseq;
}
unsigned int
rb_get_iseq_body_local_table_size(const rb_iseq_t *iseq)
{

View File

@ -350,6 +350,7 @@ fn main() {
.allowlist_function("rb_get_def_bmethod_proc")
.allowlist_function("rb_iseq_encoded_size")
.allowlist_function("rb_get_iseq_body_local_iseq")
.allowlist_function("rb_get_iseq_body_parent_iseq")
.allowlist_function("rb_get_iseq_body_iseq_encoded")
.allowlist_function("rb_get_iseq_body_stack_max")
.allowlist_function("rb_get_iseq_flags_has_opt")

View File

@ -1542,6 +1542,22 @@ fn gen_get_ep(asm: &mut Assembler, level: u32) -> Opnd {
ep_opnd
}
// Gets the EP of the ISeq of the containing method, or "local level".
// Equivalent of GET_LEP() macro.
fn gen_get_lep(jit: &mut JITState, asm: &mut Assembler) -> Opnd {
// Equivalent of get_lvar_level() in compile.c
fn get_lvar_level(iseq: IseqPtr) -> u32 {
if iseq == unsafe { rb_get_iseq_body_local_iseq(iseq) } {
0
} else {
1 + get_lvar_level(unsafe { rb_get_iseq_body_parent_iseq(iseq) })
}
}
let level = get_lvar_level(jit.get_iseq());
gen_get_ep(asm, level)
}
fn gen_getlocal_generic(
ctx: &mut Context,
asm: &mut Assembler,
@ -3995,9 +4011,14 @@ unsafe extern "C" fn build_kwhash(ci: *const rb_callinfo, sp: *const VALUE) -> V
hash
}
enum BlockHandler {
// SpecVal is a single value in an iseq invocation's environment on the stack,
// at sp[-2]. Depending on the frame type, it can serve different purposes,
// which are covered here by enum variants.
enum SpecVal {
None,
CurrentFrame,
BlockISeq(IseqPtr),
BlockParamProxy,
PrevEP(*const VALUE),
}
struct ControlFrame {
@ -4006,8 +4027,7 @@ struct ControlFrame {
iseq: Option<IseqPtr>,
pc: Option<u64>,
frame_type: u32,
block_handler: BlockHandler,
prev_ep: Option<*const VALUE>,
specval: SpecVal,
cme: *const rb_callable_method_entry_t,
local_size: i32
}
@ -4028,7 +4048,7 @@ struct ControlFrame {
// * Stack overflow is not checked (should be done by the caller)
// * Interrupts are not checked (should be done by the caller)
fn gen_push_frame(
_jit: &mut JITState,
jit: &mut JITState,
_ctx: &mut Context,
asm: &mut Assembler,
set_sp_cfp: bool, // if true CFP and SP will be switched to the callee
@ -4060,19 +4080,33 @@ fn gen_push_frame(
// Write special value at sp[-2]. It's either a block handler or a pointer to
// the outer environment depending on the frame type.
// sp[-2] = specval;
let specval: Opnd = match (frame.prev_ep, frame.block_handler) {
(None, BlockHandler::None) => {
let specval: Opnd = match frame.specval {
SpecVal::None => {
VM_BLOCK_HANDLER_NONE.into()
}
(None, BlockHandler::CurrentFrame) => {
SpecVal::BlockISeq(block_iseq) => {
// Change cfp->block_code in the current frame. See vm_caller_setup_arg_block().
// VM_CFP_TO_CAPTURED_BLOCK does &cfp->self, rb_captured_block->code.iseq aliases
// with cfp->block_code.
asm.store(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_BLOCK_CODE), VALUE::from(block_iseq).into());
let cfp_self = asm.lea(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SELF));
asm.or(cfp_self, Opnd::Imm(1))
}
(Some(prev_ep), BlockHandler::None) => {
SpecVal::BlockParamProxy => {
let ep_opnd = gen_get_lep(jit, asm);
let block_handler = asm.load(
Opnd::mem(64, ep_opnd, (SIZEOF_VALUE as i32) * (VM_ENV_DATA_INDEX_SPECVAL as i32))
);
asm.store(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_BLOCK_CODE), block_handler);
block_handler
}
SpecVal::PrevEP(prev_ep) => {
let tagged_prev_ep = (prev_ep as usize) | 1;
VALUE(tagged_prev_ep).into()
}
(_, _) => panic!("specval can only be one of prev_ep or block_handler")
};
asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -2), specval);
@ -4218,6 +4252,43 @@ fn gen_send_cfunc(
return CantCompile;
}
let block_arg = flags & VM_CALL_ARGS_BLOCKARG != 0;
let block_arg_type = if block_arg {
Some(ctx.get_opnd_type(StackOpnd(0)))
} else {
None
};
match block_arg_type {
Some(Type::Nil | Type::BlockParamProxy) => {
// We'll handle this later
}
None => {
// Nothing to do
}
_ => {
gen_counter_incr!(asm, send_block_arg);
return CantCompile;
}
}
match block_arg_type {
Some(Type::Nil) => {
// We have a nil block arg, so let's pop it off the args
ctx.stack_pop(1);
}
Some(Type::BlockParamProxy) => {
// We don't need the actual stack value
ctx.stack_pop(1);
}
None => {
// Nothing to do
}
_ => {
assert!(false);
}
}
// This is a .send call and we need to adjust the stack
if flags & VM_CALL_OPT_SEND != 0 {
handle_opt_send_shift_stack(asm, argc as i32, ctx);
@ -4229,21 +4300,16 @@ fn gen_send_cfunc(
// Store incremented PC into current control frame in case callee raises.
jit_save_pc(jit, asm);
if let Some(block_iseq) = block {
// Change cfp->block_code in the current frame. See vm_caller_setup_arg_block().
// VM_CFP_TO_CAPTURED_BLOCK does &cfp->self, rb_captured_block->code.iseq aliases
// with cfp->block_code.
asm.mov(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_BLOCK_CODE), VALUE::from(block_iseq).into());
}
// Increment the stack pointer by 3 (in the callee)
// sp += 3
let sp = asm.lea(ctx.sp_opnd((SIZEOF_VALUE as isize) * 3));
let frame_block_handler = if let Some(_block_iseq) = block {
BlockHandler::CurrentFrame
let specval = if block_arg_type == Some(Type::BlockParamProxy) {
SpecVal::BlockParamProxy
} else if let Some(block_iseq) = block {
SpecVal::BlockISeq(block_iseq)
} else {
BlockHandler::None
SpecVal::None
};
let mut frame_type = VM_FRAME_MAGIC_CFUNC | VM_FRAME_FLAG_CFRAME | VM_ENV_FLAG_LOCAL;
@ -4253,11 +4319,10 @@ fn gen_send_cfunc(
gen_push_frame(jit, ctx, asm, false, ControlFrame {
frame_type,
block_handler: frame_block_handler,
specval,
cme,
recv,
sp,
prev_ep: None,
pc: Some(0),
iseq: None,
local_size: 0,
@ -4606,6 +4671,26 @@ fn gen_send_iseq(
return CantCompile;
}
let block_arg = flags & VM_CALL_ARGS_BLOCKARG != 0;
let block_arg_type = if block_arg {
Some(ctx.get_opnd_type(StackOpnd(0)))
} else {
None
};
match block_arg_type {
Some(Type::Nil | Type::BlockParamProxy) => {
// We'll handle this later
}
None => {
// Nothing to do
}
_ => {
gen_counter_incr!(asm, send_block_arg);
return CantCompile;
}
}
// If we have unfilled optional arguments and keyword arguments then we
// would need to move adjust the arguments location to account for that.
// For now we aren't handling this case.
@ -4700,6 +4785,23 @@ fn gen_send_iseq(
// Check for interrupts
gen_check_ints(asm, side_exit);
match block_arg_type {
Some(Type::Nil) => {
// We have a nil block arg, so let's pop it off the args
ctx.stack_pop(1);
}
Some(Type::BlockParamProxy) => {
// We don't need the actual stack value
ctx.stack_pop(1);
}
None => {
// Nothing to do
}
_ => {
assert!(false);
}
}
let leaf_builtin_raw = unsafe { rb_leaf_builtin_function(iseq) };
let leaf_builtin: Option<*const rb_builtin_function> = if leaf_builtin_raw.is_null() {
None
@ -4914,32 +5016,30 @@ fn gen_send_iseq(
// Store the next PC in the current frame
jit_save_pc(jit, asm);
if let Some(block_val) = block {
// Change cfp->block_code in the current frame. See vm_caller_setup_arg_block().
// VM_CFP_TO_CAPTURED_BLCOK does &cfp->self, rb_captured_block->code.iseq aliases
// with cfp->block_code.
asm.store(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_BLOCK_CODE), VALUE(block_val as usize).into());
}
// Adjust the callee's stack pointer
let offs =
(SIZEOF_VALUE as isize) * (3 + (num_locals as isize) + if doing_kw_call { 1 } else { 0 });
let callee_sp = asm.lea(ctx.sp_opnd(offs));
let frame_block_handler = if let Some(_) = block {
BlockHandler::CurrentFrame
let specval = if let Some(prev_ep) = prev_ep {
// We've already side-exited if the callee expects a block, so we
// ignore any supplied block here
SpecVal::PrevEP(prev_ep)
} else if block_arg_type == Some(Type::BlockParamProxy) {
SpecVal::BlockParamProxy
} else if let Some(block_val) = block {
SpecVal::BlockISeq(block_val)
} else {
BlockHandler::None
SpecVal::None
};
// Setup the new frame
gen_push_frame(jit, ctx, asm, true, ControlFrame {
frame_type,
block_handler: frame_block_handler,
specval,
cme,
recv,
sp: callee_sp,
prev_ep,
iseq: Some(iseq),
pc: None, // We are calling into jitted code, which will set the PC as necessary
local_size: num_locals
@ -5146,26 +5246,23 @@ fn gen_send_general(
return CantCompile;
}
if flags & VM_CALL_ARGS_BLOCKARG != 0 {
gen_counter_incr!(asm, send_block_arg);
return CantCompile;
}
// Defer compilation so we can specialize on class of receiver
if !jit_at_current_insn(jit) {
defer_compilation(jit, ctx, asm, ocb);
return EndBlock;
}
let comptime_recv = jit_peek_at_stack(jit, ctx, argc as isize);
let recv_idx = argc + if flags & VM_CALL_ARGS_BLOCKARG != 0 { 1 } else { 0 };
let comptime_recv = jit_peek_at_stack(jit, ctx, recv_idx as isize);
let comptime_recv_klass = comptime_recv.class_of();
// Guard that the receiver has the same class as the one from compile time
let side_exit = get_side_exit(jit, ocb, ctx);
// Points to the receiver operand on the stack
let recv = ctx.stack_opnd(argc);
let recv_opnd = StackOpnd(argc.try_into().unwrap());
let recv = ctx.stack_opnd(recv_idx);
let recv_opnd = StackOpnd(recv_idx.try_into().unwrap());
jit_guard_known_klass(
jit,
ctx,
@ -5273,6 +5370,11 @@ fn gen_send_general(
let ivar_name = unsafe { get_cme_def_body_attr_id(cme) };
if flags & VM_CALL_ARGS_BLOCKARG != 0 {
gen_counter_incr!(asm, send_block_arg);
return CantCompile;
}
return gen_get_ivar(
jit,
ctx,
@ -5298,6 +5400,9 @@ fn gen_send_general(
// See :attr-tracing:
gen_counter_incr!(asm, send_cfunc_tracing);
return CantCompile;
} else if flags & VM_CALL_ARGS_BLOCKARG != 0 {
gen_counter_incr!(asm, send_block_arg);
return CantCompile;
} else {
let ivar_name = unsafe { get_cme_def_body_attr_id(cme) };
return gen_set_ivar(jit, ctx, asm, comptime_recv, ivar_name, flags, argc);
@ -5326,6 +5431,10 @@ fn gen_send_general(
}
// Send family of methods, e.g. call/apply
VM_METHOD_TYPE_OPTIMIZED => {
if flags & VM_CALL_ARGS_BLOCKARG != 0 {
gen_counter_incr!(asm, send_block_arg);
return CantCompile;
}
let opt_type = unsafe { get_cme_def_body_optimized_type(cme) };
match opt_type {
@ -6191,7 +6300,7 @@ fn gen_getblockparamproxy(
// Push rb_block_param_proxy. It's a root, so no need to use jit_mov_gc_ptr.
assert!(!unsafe { rb_block_param_proxy }.special_const_p());
let top = ctx.stack_push(Type::Unknown);
let top = ctx.stack_push(Type::BlockParamProxy);
asm.mov(top, Opnd::const_ptr(unsafe { rb_block_param_proxy }.as_ptr()));
}

View File

@ -44,6 +44,9 @@ pub enum Type {
TString, // An object with the T_STRING flag set, possibly an rb_cString
CString, // An un-subclassed string of type rb_cString (can have instance vars in some cases)
BlockParamProxy, // A special sentinel value indicating the block parameter should be read from
// the current surrounding cfp
}
// Default initialization
@ -79,6 +82,12 @@ impl Type {
if val.class_of() == unsafe { rb_cString } {
return Type::CString;
}
// We likewise can't reference rb_block_param_proxy, but it's again an optimisation;
// we can just treat it as a normal Object.
#[cfg(not(test))]
if val == unsafe { rb_block_param_proxy } {
return Type::BlockParamProxy;
}
match val.builtin_type() {
RUBY_T_ARRAY => Type::Array,
RUBY_T_HASH => Type::Hash,
@ -142,7 +151,8 @@ 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::Unknown | Type::UnknownImm | Type::UnknownHeap => None
Type::Unknown | Type::UnknownImm | Type::UnknownHeap => None,
Type::BlockParamProxy => None,
}
}

View File

@ -1418,6 +1418,9 @@ extern "C" {
extern "C" {
pub fn rb_get_iseq_body_local_iseq(iseq: *const rb_iseq_t) -> *const rb_iseq_t;
}
extern "C" {
pub fn rb_get_iseq_body_parent_iseq(iseq: *const rb_iseq_t) -> *const rb_iseq_t;
}
extern "C" {
pub fn rb_get_iseq_body_local_table_size(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;
}