diff --git a/prism/templates/lib/prism/node.rb.erb b/prism/templates/lib/prism/node.rb.erb index fb13051aba0b3b..4d9a8ebf32b0f3 100644 --- a/prism/templates/lib/prism/node.rb.erb +++ b/prism/templates/lib/prism/node.rb.erb @@ -551,7 +551,16 @@ module Prism #: (Array[Symbol]? keys) -> Hash[Symbol, untyped] def deconstruct_keys(keys) # :nodoc: - { <%= (["node_id: node_id", "location: location"] + node.fields.map { |field| "#{field.name}: #{field.name}" }).join(", ") %> } + <%- deconstruct = [:node_id, :location, *node.fields.map { |field| field.name.to_sym }, *node.fields.select { |field| field.is_a?(Prism::Template::LocationField) || field.is_a?(Prism::Template::OptionalLocationField) }.map { |field| field.name.delete_suffix("_loc").to_sym }].uniq -%> + (keys || %i[<%= deconstruct.join(" ") %>]).each_with_object( + {} #: Hash[Symbol, untyped] + ) do |key, deconstructed| + case key + <%- deconstruct.each do |key| -%> + when <%= "%-24s then" % [key.inspect] %> deconstructed[<%= key.inspect %>] = self.<%= key %> + <%- end -%> + end + end end # See `Node#type`. diff --git a/test/prism/ruby/deconstruct_keys_test.rb b/test/prism/ruby/deconstruct_keys_test.rb new file mode 100644 index 00000000000000..f81b06ef2e18ec --- /dev/null +++ b/test/prism/ruby/deconstruct_keys_test.rb @@ -0,0 +1,22 @@ +# frozen_string_literal: true + +require_relative "../test_helper" + +module Prism + class DeconstructKeysTest < TestCase + def test_deconstruct_keys + node = Prism.parse_statement("1.to_s") + + deconstruct_all = node.deconstruct_keys(nil) + assert_equal deconstruct_all[:node_id], node.node_id + assert_equal deconstruct_all[:message], "to_s" + + deconstruct_receiver = node.deconstruct_keys([:receiver]) + assert_equal 1, deconstruct_receiver[:receiver].value + refute_includes deconstruct_receiver.keys, :message + + deconstruct_invalid = node.deconstruct_keys([:invalid]) + refute_includes deconstruct_invalid.keys, :invalid + end + end +end diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs index 1a4f84d60beb4d..970f6103061d6c 100644 --- a/zjit/src/codegen.rs +++ b/zjit/src/codegen.rs @@ -667,6 +667,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio Insn::InvokeBlockIfunc { cd, block_handler, args, state, .. } => gen_invokeblock_ifunc(jit, asm, function, *cd, opnd!(block_handler), opnds!(args), &function.frame_state(*state)), Insn::InvokeProc { recv, args, state, kw_splat } => gen_invokeproc(jit, asm, function, opnd!(recv), opnds!(args), *kw_splat, &function.frame_state(*state)), Insn::InvokeBuiltin { bf, leaf, args, state, .. } => gen_invokebuiltin(jit, asm, function, &function.frame_state(*state), unsafe { &**bf }, *leaf, opnds!(args)), + Insn::InvokeBlockIseqDirect { iseq, captured, args, state } => gen_invoke_block_iseq_direct(cb, jit, asm, function, *iseq, opnd!(captured), opnds!(args), &function.frame_state(*state)), &Insn::EntryPoint { jit_entry_idx } => no_output!(gen_entry_point(jit, asm, jit_entry_idx)), Insn::Return { val } => no_output!(gen_return(asm, opnd!(val))), Insn::FixnumAdd { left, right, state } => gen_fixnum_add(jit, asm, function, opnd!(left), opnd!(right), &function.frame_state(*state)), @@ -1911,6 +1912,86 @@ fn gen_invokeproc( ) } +/// Compile `yield`. Inlines the block ISEQ frame like `invokeblock` instead of calling vm_yield. +/// The block handler is read from the enclosing frame's LEP (`level` hops up), guarded to be the +/// comptime-known ISEQ block, and its frame is pushed here before jumping to the block's JIT entry. +/// On a guard miss, side-exit and recompile. The HIR gate ensures the block is simple + lead-only +/// + non-throwing. +fn gen_invoke_block_iseq_direct( + cb: &mut CodeBlock, + jit: &mut JITState, + asm: &mut Assembler, + function: &Function, + block_iseq: IseqPtr, + captured: Opnd, + args: Vec, + state: &FrameState, +) -> lir::Opnd { + gen_incr_counter(asm, Counter::block_iseq_direct_optimized_send_count); + + let local_size = unsafe { get_iseq_body_local_table_size(block_iseq) }.to_usize(); + let stack_growth = state.stack_size() + local_size + unsafe { get_iseq_body_stack_max(block_iseq) }.to_usize(); + gen_stack_overflow_check(jit, asm, function, state, stack_growth); + + // `captured` is the guarded `struct rb_captured_block *` (block handler with the ISEQ tag + // masked off). The HIR builder loaded it from the LEP and guarded the tag + iseq identity. + // TODO: During inlining, captured->self can be known. It should be put into HIR. + let captured_self = asm.load(Opnd::mem(64, captured, 0)); // captured->self + // TODO: During inlining, captured->ep can sometimes also be known. + let captured_ep = asm.load(Opnd::mem(64, captured, SIZEOF_VALUE_I32)); // captured->ep + // specval = VM_GUARDED_PREV_EP(captured->ep) = captured->ep | 0x01 + let specval = asm.or(captured_ep, Opnd::Imm(0x1)); + + let stack_size = state.stack().len() - args.len(); + let stack_map = build_stack_map(jit, function, &state.with_stack_size(stack_size)); + let jit_frame = gen_write_jit_frame(asm, state, stack_map.len()); + gen_save_sp(asm, stack_size); + + gen_spill_locals(jit, asm, state); + asm.stack_map(stack_map, jit_frame, state.depth); + + gen_push_frame(asm, args.len(), state, ControlFrame { + recv: captured_self, + iseq: Some(block_iseq), + cme: std::ptr::null(), + frame_type: VM_FRAME_MAGIC_BLOCK, + specval, + write_block_code: iseq_may_write_block_code(block_iseq), + }); + + asm_comment!(asm, "switch to new SP register"); + let sp_offset = (stack_size + local_size + VM_ENV_DATA_SIZE.to_usize()) * SIZEOF_VALUE; + let new_sp = asm.add(SP, sp_offset.into()); + asm.mov(SP, new_sp); + + asm_comment!(asm, "switch to new CFP"); + let new_cfp = asm.sub(CFP, RUBY_SIZEOF_CONTROL_FRAME.into()); + asm.mov(CFP, new_cfp); + asm.store(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), CFP); + + // JIT-to-JIT convention: self as c_args[0], then positional args. The block is + // gated to simple + lead-only + exact arity, so there are no optionals/kw/block. + let mut c_args = Vec::with_capacity(1 + args.len()); + c_args.push(captured_self); + c_args.extend(&args); + + let iseq_call = IseqCall::new(block_iseq, 0, args.len().try_into().expect("checked in HIR")); + let dummy_ptr = cb.get_write_ptr().raw_ptr(cb); + jit.iseq_calls.push(iseq_call.clone()); + let ret = asm.ccall_with_iseq_call(dummy_ptr, c_args, &iseq_call); + + // If the callee side-exits (returns Qundef), propagate to the caller. + asm_comment!(asm, "side-exit if callee side-exits"); + asm.cmp(ret, Qundef.into()); + asm.je(jit, ZJITState::get_exit_trampoline().into()); + + asm_comment!(asm, "restore SP register for the caller"); + let new_sp = asm.sub(SP, sp_offset.into()); + asm.mov(SP, new_sp); + + ret +} + /// Compile a dynamic dispatch for `super` fn gen_invokesuper( jit: &mut JITState, @@ -2871,6 +2952,7 @@ fn gen_guard_bit_equals(jit: &mut JITState, asm: &mut Assembler, function: &Func let expected_opnd: Opnd = match expected { crate::hir::Const::Value(v) => { Opnd::Value(v) } crate::hir::Const::CInt64(v) => { v.into() } + crate::hir::Const::CPtr(v) => { Opnd::const_ptr(v) } crate::hir::Const::CShape(v) => { Opnd::UImm(v.0 as u64) } _ => panic!("gen_guard_bit_equals: unexpected hir::Const {expected:?}"), }; @@ -2973,6 +3055,28 @@ pub(crate) fn iseq_may_write_block_code(iseq: IseqPtr) -> bool { false } +/// True if the block ISEQ contains a `throw` opcode (break, non-local return). ZJIT can't +/// compile `throw`, so inlining such a block's frame would side-exit + deopt on every call. +/// These blocks fall back to `vm_yield`, which handles the non-local exit in C. +/// (`next`/`redo` lower to `leave`/`jump`, not `throw`, so they stay inlinable.) +pub(crate) fn block_iseq_may_throw(iseq: IseqPtr) -> bool { + let encoded_size = unsafe { rb_iseq_encoded_size(iseq) }; + let mut insn_idx: u32 = 0; + + while insn_idx < encoded_size { + let pc = unsafe { rb_iseq_pc_at_idx(iseq, insn_idx) }; + let opcode = unsafe { rb_iseq_bare_opcode_at_pc(iseq, pc) } as u32; + + if opcode == YARVINSN_throw { + return true; + } + + insn_idx = insn_idx.saturating_add(unsafe { rb_insn_len(VALUE(opcode as usize)) }.try_into().unwrap()); + } + + false +} + /// Byte offset from NATIVE_BASE_PTR of the JITFrame storage slot for a frame at /// the given inlining depth. Depth 0 (the top-level frame) lives at /// `[NATIVE_BASE_PTR - 8]`; each deeper inlined frame gets the next slot below. diff --git a/zjit/src/codegen_tests.rs b/zjit/src/codegen_tests.rs index eed7e280e9a15a..8ecfdf98e39994 100644 --- a/zjit/src/codegen_tests.rs +++ b/zjit/src/codegen_tests.rs @@ -543,6 +543,166 @@ fn test_getblockparamproxy_polymorphic_none_and_iseq_and_proc() { assert_snapshot!(assert_compiles("val = proc { 2 }; test(&val)"), @"2"); } +#[test] +fn test_yield_inline_self_is_captured_self() { + // The inlined frame's self must be the block's captured self, not the yielding receiver. + set_call_threshold(2); + eval(" + class Yielder + def run = yield + end + class C + def initialize(v) = @v = v + def go(y) = y.run { @v * 2 } + end + Y = Yielder.new + C.new(21).go(Y) + C.new(21).go(Y) + "); + assert_snapshot!(assert_compiles("C.new(21).go(Y)"), @"42"); +} + +#[test] +fn test_yield_iseq_guard_miss_recompiles() { + set_call_threshold(2); + eval(" + def invoke = yield(41) + invoke { |x| x * 2 } + invoke { |x| x * 2 } + "); + assert_snapshot!(assert_compiles_allowing_exits("[invoke { |x| x + 1 }, invoke { |x| x * 2 }]"), @"[42, 82]"); +} + +#[test] +fn test_yield_inline_invocation_with_args() { + // Plain yield with two args to a matching-arity block inlines and returns correctly. + set_call_threshold(2); + eval(" + def foo = yield(3, 4) + def test = foo { |a, b| a + b } + test + test + "); + assert_snapshot!(assert_compiles("test"), @"7"); +} + +#[test] +fn test_yield_inline_invocation_live_stack_below_args() { + // A live value sits on the stack below the yield args; the no-receiver-slot SP math + // must preserve it so `x +` sees the right operand. + set_call_threshold(2); + eval(" + def foo(x) = x + yield(1, 2) + def test = foo(10) { |a, b| a + b } + test + test + "); + assert_snapshot!(assert_compiles("test"), @"13"); +} + +#[test] +fn test_yield_inlined_caller_block_dispatches_without_guards() { + // When the yielding method is inlined into a caller that passes a literal block, the block + // handler is written into the inlined frame's EP from a compile-time constant, so the yield + // dispatches with no tag/iseq guards. assert_inlines requires the method to actually inline + // and to run with no side exits, exercising the guard-free InvokeBlockIseqDirect machine code. + with_inlining(|| { + assert_snapshot!(assert_inlines(" + def two_yields = (yield 1) + (yield 2) + def test = two_yields { |x| x * 10 } + test + test + "), @"30"); + }); +} + +#[test] +fn test_yield_with_lambda_arg() { + // A lambda passed via &l is a proc handler (not imemo_iseq): yield falls back but runs. + set_call_threshold(2); + eval(" + def foo = yield(5) + def test = foo(&L) + L = ->(x) { x * 10 } + test + test + "); + assert_snapshot!(assert_compiles_allowing_exits("test"), @"50"); +} + +#[test] +fn test_yield_break() { + set_call_threshold(2); + eval(" + def foo = yield + def test = foo { break 5 } + test + test + "); + assert_snapshot!(assert_compiles_allowing_exits("test"), @"5"); +} + +#[test] +fn test_yield_non_local_return() { + set_call_threshold(2); + eval(" + def inner = yield + def test + inner { return 42 } + 99 + end + test + test + "); + assert_snapshot!(assert_compiles_allowing_exits("test"), @"42"); +} + +#[test] +fn test_yield_autosplat() { + // {|a, b|} auto-splats a single Array arg for yield (falls back). + set_call_threshold(2); + eval(" + def via_yield = yield([3, 4]) + def test_yield = via_yield { |a, b| a + b } + test_yield; test_yield + "); + assert_snapshot!(assert_compiles_allowing_exits("test_yield"), @"7"); +} + +#[test] +fn test_yield_next() { + // next(val) compiles to leave (not throw), so yield inlines invocation and returns val. + set_call_threshold(2); + eval(" + def via_yield = yield + def test_yield = via_yield { next 7 } + test_yield; test_yield + "); + assert_snapshot!(assert_compiles("test_yield"), @"7"); +} + +#[test] +fn test_yield_inline_ensure_runs() { + // The ensure body must run on the normal inlined invocation yield path. + set_call_threshold(2); + eval(" + def foo = yield + $log = [] + def driver + foo do + begin + 42 + ensure + $log << :ensured + end + end + end + driver + driver + "); + assert_snapshot!(assert_compiles_allowing_exits("$log.clear; [driver, $log]"), @"[42, [:ensured]]"); +} + #[test] fn test_getblockparam() { eval(" @@ -5473,10 +5633,13 @@ fn test_invokeblock() { def test yield end - test { 41 } + def entry + test { 42 } + end + entry "); assert_contains_opcode("test", YARVINSN_invokeblock); - assert_snapshot!(assert_compiles("test { 42 }"), @"42"); + assert_snapshot!(assert_compiles("entry"), @"42"); } #[test] @@ -5485,10 +5648,13 @@ fn test_invokeblock_with_args() { def test(x, y) yield x, y end - test(1, 2) { |a, b| a + b } + def entry + test(1, 2) { |a, b| a + b } + end + entry "); assert_contains_opcode("test", YARVINSN_invokeblock); - assert_snapshot!(assert_compiles("test(1, 2) { |a, b| a + b }"), @"3"); + assert_snapshot!(assert_compiles("entry"), @"3"); } #[test] @@ -5500,7 +5666,9 @@ fn test_invokeblock_no_block_given() { test { } "); assert_contains_opcode("test", YARVINSN_invokeblock); - assert_snapshot!(assert_compiles("test"), @":error"); + // Compiled expecting an ISEQ block; calling with none misses the handler guard and + // deopts, so the interpreter raises LocalJumpError (rescued to :error). + assert_snapshot!(assert_compiles_allowing_exits("test"), @":error"); } #[test] @@ -5511,14 +5679,15 @@ fn test_invokeblock_multiple_yields() { yield 2 yield 3 end - test { |x| x } + def entry + results = [] + test { |x| results << x } + results + end + entry "); assert_contains_opcode("test", YARVINSN_invokeblock); - assert_snapshot!(assert_compiles(" - results = [] - test { |x| results << x } - results - "), @"[1, 2, 3]"); + assert_snapshot!(assert_compiles("entry"), @"[1, 2, 3]"); } #[test] diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs index 1dd357e76bbc0f..cbc302320f6d61 100644 --- a/zjit/src/hir.rs +++ b/zjit/src/hir.rs @@ -559,6 +559,8 @@ pub enum SideExitReason { BlockParamProxyNotProc, BlockParamProxyFallbackMiss, BlockParamProxyProfileNotCovered, + InvokeBlockHandlerNotIseq, + InvokeBlockIseqChanged, BlockParamWbRequired, StackOverflow, FixnumModByZero, @@ -821,6 +823,7 @@ pub enum FieldName { VM_ENV_DATA_INDEX_SPECVAL, VM_ENV_DATA_INDEX_FLAGS, RBASIC_FLAGS, + code_iseq, shape_id, as_heap, fields_obj, @@ -1124,6 +1127,18 @@ pub enum Insn { state: InsnId, kw_splat: bool, }, + /// Fast-path `yield`: the enclosing frame's block is a known ISEQ block, so push its frame + /// inline and jump straight to its JIT entry. `iseq` is the comptime-known block iseq; the + /// tag + iseq-identity guards live in preceding HIR instructions, and `captured` is the + /// guarded, untagged `struct rb_captured_block *` codegen reads self/ep from. Args are the + /// positional arguments (lead-only, exact arity). + InvokeBlockIseqDirect { + iseq: IseqPtr, + /// Guarded `struct rb_captured_block *` (block handler with the ISEQ tag masked off). + captured: InsnId, + args: Vec, + state: InsnId, + }, /// Optimized ISEQ call SendDirect(Box), @@ -1490,6 +1505,11 @@ macro_rules! for_each_operand_impl { $visit_many!(args); $visit_one!(*state); } + Insn::InvokeBlockIseqDirect { captured, args, state, .. } => { + $visit_one!(*captured); + $visit_many!(args); + $visit_one!(*state); + } Insn::InvokeBlockIfunc { block_handler, args, state, .. } => { $visit_one!(*block_handler); $visit_many!(args); @@ -1809,6 +1829,7 @@ impl Insn { Insn::IncrCounterPtr { .. } => Effect::read_write(abstract_heaps::Empty, abstract_heaps::Other), Insn::CheckInterrupts { .. } => Effect::read_write(abstract_heaps::InterruptFlag, abstract_heaps::Control), Insn::InvokeProc { .. } => effects::Any, + Insn::InvokeBlockIseqDirect { .. } => effects::Any, Insn::RefineType { .. } => effects::Empty, Insn::HasType { expected, .. } => Effect::read_write( @@ -2125,6 +2146,11 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> { } Ok(()) } + Insn::InvokeBlockIseqDirect { iseq, captured, args, .. } => { + write!(f, "InvokeBlockIseqDirect ({:?}), {captured}", self.ptr_map.map_ptr(iseq))?; + write_separated!(f, ", ", ", ", args); + Ok(()) + } Insn::InvokeBuiltin { bf, args, leaf, .. } => { let bf_name = unsafe { CStr::from_ptr((**bf).name) }.to_str().unwrap(); write!(f, "InvokeBuiltin{} {}", @@ -2718,6 +2744,29 @@ fn iseq_get_return_value(iseq: IseqPtr, captured_opnd: Option, ci_flags: } } +/// True if the `invokeblock` call flags permit inlining the block dispatch. +fn block_call_inlinable(flags: u32) -> bool { + (flags & (VM_CALL_ARGS_SPLAT | VM_CALL_KW_SPLAT | VM_CALL_KWARG | VM_CALL_ARGS_BLOCKARG)) == 0 +} + +/// True if `yield` with `argc` positional args can dispatch by inlining the block ISEQ +/// frame. The block must take the simple callee-setup path (`rb_simple_iseq_p`) +/// with an exact arity match, avoid arg0 auto-splat, and contain no `throw` (break / +/// non-local return). Anything else falls back to the generic `invokeblock` dispatch. +fn block_call_inlinable_iseq(iseq: IseqPtr, argc: usize) -> bool { + if !unsafe { rb_simple_iseq_p(iseq) } { + return false; + } + let lead_num = unsafe { rb_get_iseq_body_param_lead_num(iseq) } as usize; + if argc != lead_num { + return false; + } + if argc == 1 && !unsafe { rb_get_iseq_flags_ambiguous_param0(iseq) } { + return false; + } + !crate::codegen::block_iseq_may_throw(iseq) +} + impl Function { fn new(iseq: *const rb_iseq_t) -> Function { Function { @@ -2798,6 +2847,36 @@ impl Function { self.load_field(block, str, FieldName::len, RUBY_OFFSET_RSTRING_LEN, types::CInt64) } + /// Emit the fast-path `yield` dispatch to a known ISEQ block. + /// When `guarded`, the block handler is read from the runtime LEP and guarded (tag + iseq + /// identity) because the profiled block can differ per caller. When the enclosing method is + /// inlined and the caller passed a literal block, `gen_push_inline_frame` wrote that exact + /// block into this frame's EP from a compile-time constant, so both guards are unnecessary. + fn push_invoke_block_iseq_direct(&mut self, block: BlockId, block_iseq: IseqPtr, level: u32, args: Vec, state: InsnId, guarded: bool) -> InsnId { + let ep = self.get_ep(block, level); + let block_handler = self.load_ep_env_field(block, ep, FieldName::VM_ENV_DATA_INDEX_SPECVAL, VM_ENV_DATA_INDEX_SPECVAL, types::CInt64); + + if guarded { + // Guard the handler is an ISEQ block: VM_BH_ISEQ_BLOCK_P is `& 0x3 == 0x1`. + let tag_mask = self.push_insn(block, Insn::Const { val: Const::CInt64(0x3) }); + let tag = self.push_insn(block, Insn::IntAnd { left: block_handler, right: tag_mask }); + self.push_insn(block, Insn::GuardBitEquals { val: tag, expected: Const::CInt64(0x1), reason: Box::new(SideExitReason::InvokeBlockHandlerNotIseq), state, recompile: Some(Recompile) }); + } + + // captured = block_handler & ~0x3 (struct rb_captured_block *) + let untag_mask = self.push_insn(block, Insn::Const { val: Const::CInt64(!0x3) }); + let captured = self.push_insn(block, Insn::IntAnd { left: block_handler, right: untag_mask }); + + if guarded { + // Guard captured->code.iseq is the comptime block iseq. Compare the raw imemo pointer: + // type inference (from_value) can't type an iseq imemo, so guard it as a CPtr identity. + let captured_iseq = self.load_field(block, captured, FieldName::code_iseq, 2 * SIZEOF_VALUE_I32, types::CPtr); + self.push_insn(block, Insn::GuardBitEquals { val: captured_iseq, expected: Const::CPtr(block_iseq as *const u8), reason: Box::new(SideExitReason::InvokeBlockIseqChanged), state, recompile: Some(Recompile) }); + } + + self.push_insn(block, Insn::InvokeBlockIseqDirect { iseq: block_iseq, captured, args, state }) + } + // Add an instruction to an SSA block fn push_insn_id(&mut self, block: BlockId, insn_id: InsnId) -> InsnId { self.blocks[block.0].insns.push(insn_id); @@ -3137,6 +3216,7 @@ impl Function { Insn::InvokeBlock { .. } => types::BasicObject, Insn::InvokeBlockIfunc { .. } => types::BasicObject, Insn::InvokeProc { .. } => types::BasicObject, + Insn::InvokeBlockIseqDirect { .. } => types::BasicObject, Insn::InvokeBuiltin { return_type, .. } => *return_type, Insn::Defined { pushval, .. } => Type::from_value(*pushval).union(types::NilClass), Insn::DefinedIvar { pushval, .. } => Type::from_value(*pushval).union(types::NilClass), @@ -4881,6 +4961,7 @@ impl Function { caller: post_send_caller, depth: caller_depth + 1, jit_entry_idx: passed_opt_num, + blockiseq, }; let add_result = match add_iseq_to_hir(self, iseq, mode) { Ok(r) => r, @@ -6542,6 +6623,7 @@ impl Function { } // Instructions with a Vec of Ruby objects Insn::InvokeBlock { ref args, .. } + | Insn::InvokeBlockIseqDirect { ref args, .. } | Insn::InvokeBlockIfunc { ref args, .. } | Insn::NewArray { elements: ref args, .. } | Insn::ArrayHash { elements: ref args, .. } @@ -7335,6 +7417,8 @@ enum AddIseqMode { depth: InlineDepth, /// The JIT entry index selected by the call site's argument count. jit_entry_idx: usize, + /// The literal block the caller passed to this frame, if any. + blockiseq: Option, }, } @@ -8914,15 +8998,51 @@ fn add_iseq_to_hir( } let args = state.stack_pop_n(crate::profile::num_arguments_on_stack(cd))?; - // Check if this is a monomorphic IFUNC block handler we can specialize + // The monomorphic block handler class the profile recorded, if any. Both the + // IFUNC and inline-ISEQ specializations below key off this single distribution. let block_handler_types = payload.profile.get_operand_types(exit_state.insn_idx); - let is_ifunc = (flags & (VM_CALL_ARGS_SPLAT | VM_CALL_KW_SPLAT)) == 0 - && block_handler_types.is_some_and(|types| types.len() == 1 && { - let summary = TypeDistributionSummary::new(&types[0]); - summary.is_monomorphic() && unsafe { rb_IMEMO_TYPE_P(summary.bucket(0).class(), imemo_ifunc) == 1 } - }); + let block_handler_class = block_handler_types.and_then(|types| { + if types.len() != 1 { return None; } + let summary = TypeDistributionSummary::new(&types[0]); + if !summary.is_monomorphic() { return None; } + Some(summary.bucket(0).class()) + }); - let result = if is_ifunc { + let is_ifunc = (flags & (VM_CALL_ARGS_SPLAT | VM_CALL_KW_SPLAT)) == 0 + && block_handler_class.is_some_and(|obj| unsafe { rb_IMEMO_TYPE_P(obj, imemo_ifunc) == 1 }); + + // If the block handler is a known simple ISEQ block with exact arity and no + // non-local exit, push its frame inline instead of calling rb_vm_invokeblock. + let inline_iseq = if block_call_inlinable(flags) { + block_handler_class.and_then(|obj| { + if unsafe { rb_IMEMO_TYPE_P(obj, imemo_iseq) == 1 } { + let iseq = obj.as_iseq(); + if block_call_inlinable_iseq(iseq, args.len()) { return Some(iseq); } + } + None + }) + } else { None }; + + let inlined_known_block = if let AddIseqMode::Inlined { blockiseq: Some(bi), .. } = mode { + if block_call_inlinable(flags) + // Only methods are inlined today, so exit_state.iseq is always a method iseq and this is + // always 0. That matters because the emit below is guard-free and bakes in both level 0 + // and *this* frame's block (bi) — sound only when the yield resolves to this exact frame. + // TODO: if block iseqs become inlinable, a yield here could resolve to an ancestor frame + // (level > 0). To stay guard-free we'd bake in get_lvar_level(...) as the level and fetch + // that ancestor's blockiseq from the inline caller chain instead of bi. + && get_lvar_level(exit_state.iseq) == 0 + && block_call_inlinable_iseq(bi, args.len()) { + Some(bi) + } else { None } + } else { None }; + + let result = if let Some(block_iseq) = inlined_known_block { + fun.push_invoke_block_iseq_direct(block, block_iseq, 0, args, exit_id, false) + } else if let Some(block_iseq) = inline_iseq { + let level = get_lvar_level(exit_state.iseq); + fun.push_invoke_block_iseq_direct(block, block_iseq, level, args, exit_id, true) + } else if is_ifunc { // Load the block handler from the current frame's LEP. In inlined // code, the function ISEQ is the caller while `exit_state.iseq` is the // callee containing this `invokeblock`. diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs index 2a59144b2d6c77..2d7e47b8d9c8bd 100644 --- a/zjit/src/hir/opt_tests.rs +++ b/zjit/src/hir/opt_tests.rs @@ -3816,6 +3816,118 @@ mod hir_opt_tests { "); } + #[test] + fn test_yield_no_args_inlines_invocation() { + // `foo` is inlined into `test`, which passes a literal block, so PushInlineFrame writes + // that exact block into the frame's EP from a compile-time constant. The yield's block + // handler is therefore statically known: dispatch has no tag/iseq GuardBitEquals, just an + // untag (IntAnd -4) feeding InvokeBlockIseqDirect. + let result = eval(" + def foo = yield + def test = foo { 42 } + test + test + "); + assert_eq!(VALUE::fixnum_from_usize(42), result); + assert_snapshot!(hir_string("test"), @" + fn test@:3: + bb1(): + EntryPoint interpreter + v1:BasicObject = LoadSelf + Jump bb3(v1) + bb2(): + EntryPoint JIT(0) + v4:BasicObject = LoadArg :self@0 + Jump bb3(v4) + bb3(v6:BasicObject): + PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010) + v18:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v6, ObjectSubclass[class_exact*:Object@VALUE(0x1000)] recompile + PushInlineFrame v18 (0x1038) + v25:CPtr = GetEP 0 + v26:CInt64 = LoadField v25, :VM_ENV_DATA_INDEX_SPECVAL@0x1040 + v27:CInt64[-4] = Const CInt64(-4) + v28:CInt64 = IntAnd v26, v27 + v29:BasicObject = InvokeBlockIseqDirect (0x1048), v28 + CheckInterrupts + PopInlineFrame + Return v29 + "); + } + + #[test] + fn test_yield_live_stack_below_args_inlines_invocation() { + // A live value sits on the stack below the yield args (base > 0): the no-receiver-slot + // SP math must preserve it. + let result = eval(" + def foo(x) = x + yield(1, 2) + def test = foo(10) { |a, b| a + b } + test + test + "); + assert_eq!(VALUE::fixnum_from_usize(13), result); + assert_snapshot!(hir_string("test"), @" + fn test@:3: + bb1(): + EntryPoint interpreter + v1:BasicObject = LoadSelf + Jump bb3(v1) + bb2(): + EntryPoint JIT(0) + v4:BasicObject = LoadArg :self@0 + Jump bb3(v4) + bb3(v6:BasicObject): + v11:Fixnum[10] = Const Value(10) + PatchPoint MethodRedefined(Object@0x1000, foo@0x1008, cme:0x1010) + v20:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v6, ObjectSubclass[class_exact*:Object@VALUE(0x1000)] recompile + PushInlineFrame v20 (0x1038), v11 + v29:Fixnum[1] = Const Value(1) + v31:Fixnum[2] = Const Value(2) + v33:CPtr = GetEP 0 + v34:CInt64 = LoadField v33, :VM_ENV_DATA_INDEX_SPECVAL@0x1040 + v35:CInt64[-4] = Const CInt64(-4) + v36:CInt64 = IntAnd v34, v35 + v37:BasicObject = InvokeBlockIseqDirect (0x1048), v36, v29, v31 + PatchPoint MethodRedefined(Integer@0x1050, +@0x1058, cme:0x1060) + v52:Fixnum = GuardType v37, Fixnum + v53:Fixnum = FixnumAdd v11, v52 + CheckInterrupts + PopInlineFrame + Return v53 + "); + } + + #[test] + fn test_yield_lambda_falls_back() { + // A lambda passed via &l becomes a proc block handler (not imemo_iseq), so it never inlines invocation. + // Compiles to Send. + let result = eval(" + def foo = yield(5) + def test(l) = foo(&l) + l = ->(x) { x * 10 } + test(l) + test(l) + "); + assert_eq!(VALUE::fixnum_from_usize(50), result); + assert_snapshot!(hir_string("test"), @" + fn test@:3: + bb1(): + EntryPoint interpreter + v1:BasicObject = LoadSelf + v2:CPtr = LoadSP + v3:BasicObject = LoadField v2, :l@0x1000 + Jump bb3(v1, v3) + bb2(): + EntryPoint JIT(0) + v6:BasicObject = LoadArg :self@0 + v7:BasicObject = LoadArg :l@1 + Jump bb3(v6, v7) + bb3(v9:BasicObject, v10:BasicObject): + v16:BasicObject = Send v9, &block, :foo, v10 # SendFallbackReason: Complex argument passing + CheckInterrupts + Return v16 + "); + } + #[test] fn reload_local_across_send() { eval(" @@ -9551,10 +9663,14 @@ mod hir_opt_tests { v18:ObjectSubclass[class_exact*:Object@VALUE(0x1000)] = GuardType v6, ObjectSubclass[class_exact*:Object@VALUE(0x1000)] recompile PushInlineFrame v18 (0x1038) v25:Fixnum[1] = Const Value(1) - v27:BasicObject = InvokeBlock v25 # SendFallbackReason: InvokeBlock: not yet specialized + v27:CPtr = GetEP 0 + v28:CInt64 = LoadField v27, :VM_ENV_DATA_INDEX_SPECVAL@0x1040 + v29:CInt64[-4] = Const CInt64(-4) + v30:CInt64 = IntAnd v28, v29 + v31:BasicObject = InvokeBlockIseqDirect (0x1048), v30, v25 CheckInterrupts PopInlineFrame - Return v27 + Return v31 "); } @@ -16619,11 +16735,20 @@ mod hir_opt_tests { v75:Array = RefineType v70, Array v76:CInt64 = UnboxFixnum v71 v77:BasicObject = ArrayAref v75, v76 - v79:BasicObject = InvokeBlock v77 # SendFallbackReason: InvokeBlock: not yet specialized - v83:Fixnum[1] = Const Value(1) - v84:Fixnum = FixnumAdd v71, v83 + v79:CPtr = GetEP 0 + v80:CInt64 = LoadField v79, :VM_ENV_DATA_INDEX_SPECVAL@0x1000 + v81:CInt64[3] = Const CInt64(3) + v82:CInt64 = IntAnd v80, v81 + v83:CInt64[1] = GuardBitEquals v82, CInt64(1) recompile + v84:CInt64[-4] = Const CInt64(-4) + v85:CInt64 = IntAnd v80, v84 + v86:CPtr = LoadField v85, :code_iseq@0x1001 + v87:CPtr[CPtr(0x1002)] = GuardBitEquals v86, CPtr(0x1002) recompile + v88:BasicObject = InvokeBlockIseqDirect (0x1008), v85, v77 + v92:Fixnum[1] = Const Value(1) + v93:Fixnum = FixnumAdd v71, v92 PatchPoint NoEPEscape(each) - Jump bb8(v70, v84) + Jump bb8(v70, v93) bb4(v23:BasicObject, v24:NilClass): v28:BasicObject = InvokeBuiltin , v23 CheckInterrupts @@ -19008,11 +19133,15 @@ mod hir_opt_tests { PatchPoint MethodRedefined(Object@0x1008, with_yield@0x1010, cme:0x1018) v25:ObjectSubclass[class_exact*:Object@VALUE(0x1008)] = GuardType v9, ObjectSubclass[class_exact*:Object@VALUE(0x1008)] recompile PushInlineFrame v25 (0x1040), v10 - v34:BasicObject = InvokeBlock v10 # SendFallbackReason: InvokeBlock: not yet specialized + v34:CPtr = GetEP 0 + v35:CInt64 = LoadField v34, :VM_ENV_DATA_INDEX_SPECVAL@0x1048 + v36:CInt64[-4] = Const CInt64(-4) + v37:CInt64 = IntAnd v35, v36 + v38:BasicObject = InvokeBlockIseqDirect (0x1050), v37, v10 CheckInterrupts PopInlineFrame PatchPoint NoEPEscape(test) - Return v34 + Return v38 "); } diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs index bd88b97eeed124..7468f96a5eab5f 100644 --- a/zjit/src/stats.rs +++ b/zjit/src/stats.rs @@ -231,6 +231,8 @@ make_counters! { exit_block_param_proxy_not_proc, exit_block_param_proxy_fallback_miss, exit_block_param_proxy_profile_not_covered, + exit_invoke_block_handler_not_iseq, + exit_invoke_block_iseq_changed, exit_block_param_wb_required, exit_too_many_keyword_parameters, exit_too_many_args_for_lir, @@ -305,6 +307,7 @@ make_counters! { inline_iseq_optimized_send_count, non_variadic_cfunc_optimized_send_count, variadic_cfunc_optimized_send_count, + block_iseq_direct_optimized_send_count, } // Ivar fallback counters that are summed as dynamic_setivar_count @@ -635,6 +638,8 @@ pub fn side_exit_counter(reason: crate::hir::SideExitReason) -> Counter { BlockParamProxyNotProc => exit_block_param_proxy_not_proc, BlockParamProxyFallbackMiss => exit_block_param_proxy_fallback_miss, BlockParamProxyProfileNotCovered => exit_block_param_proxy_profile_not_covered, + InvokeBlockHandlerNotIseq => exit_invoke_block_handler_not_iseq, + InvokeBlockIseqChanged => exit_invoke_block_iseq_changed, BlockParamWbRequired => exit_block_param_wb_required, TooManyKeywordParameters => exit_too_many_keyword_parameters, TooManyArgsForLir => exit_too_many_args_for_lir,