ZJIT: Pass a &blk argument through to a direct send

`bar(&blk)` only became a direct send when ZJIT could prove `blk` was nil.
Anything else kept the site on `rb_vm_send`, which pays a method search, an
interpreter frame push and a `vm_exec` setjmp on every call. On lobsters that
is 2.06M of the 12.0M dynamic sends a run -- 17% of them -- concentrated in
Rails' `def _read_attribute(name, &block) = @attributes.fetch_value(name,
&block)` style forwarding methods.

`vm_caller_setup_arg_block` produces the callee's block handler from the block
argument, and two of its cases need no work at run time:

* A `Proc` is its own block handler, so a guard on the argument's class is
  enough to write it into the callee frame's specval.
* `rb_block_param_proxy` resolves to `VM_CF_BLOCK_HANDLER(cfp)`, this frame's
  own block handler, which is one load off the local EP. That is the case a
  `def foo(&blk) = bar(&blk)` site compiles to, because the block parameter is
  only ever read as `&blk`.

The proxy reaches the send through the join of `getblockparamproxy`'s
modified/unmodified branches, so its identity is not visible at the send.
Branch on it there instead: `getblockparamproxy` records the value it pushed
for this frame's local EP, and a site that profiled the proxy's (singleton)
class gets an arm that refines the argument to the proxy object, which
`type_specialize` then turns into the EP load. The other side of the branch is
a block parameter that `setblockparam` materialized, which is a legitimate
value to reach the site, so it keeps the dynamic send.

`SendDirect` grows a `block_arg` operand for the resulting handler, and the
inliner declines any call that has one: inlining does not push the callee
frame whose specval would carry it.

On lobsters (10 warmup + 5 bench iterations):

    dynamic_send_count      12,001,057 -> 10,527,202  (-12.3%)
    send_polymorphic         1,324,915 ->    466,975
    send_block_arg_not_nil     692,187 ->     57,695
    iseq_optimized_send_count 29,307,998 -> 30,822,668

[zjit/min port note] emit_polymorphic_send() ported without the per-arm profiled-shape
machinery (ARM_SHAPE_MIN_SHARE / as_polymorphic_arm / record_profiled_type) from
excluded commits, so each arm drops the receiver's profile entry the way master's
gen_send_chain does. Keeps master's try_inline_send_direct() entry point, its `cd`
argument to gen_send_iseq_direct(), and its 2-arg gen_block_handler_specval();
SendDirectData gains only `block_arg`, not the excluded `guard_state`.

Author: k0kubun

zjit/bindgen/src/main.rs

@@ -188,6 +188,7 @@ fn main() {
         .allowlist_var("rb_cSymbol")
         .allowlist_var("rb_cFloat")
         .allowlist_var("rb_cNumeric")
+        .allowlist_var("rb_cProc")
         .allowlist_var("rb_cRange")
         .allowlist_var("rb_cString")
         .allowlist_var("rb_cThread")

zjit/src/codegen.rs

@@ -695,11 +695,12 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::Send { cd, block: Some(BlockHandler::BlockArg), state, reason, .. } => gen_send(jit, asm, function, cd, std::ptr::null(), &function.frame_state(state), reason),
         &Insn::SendForward { cd, blockiseq, state, reason, .. } => gen_send_forward(jit, asm, function, cd, blockiseq, &function.frame_state(state), reason),
         Insn::SendDirect(insn) => {
-            let SendDirectData { cd, cme, iseq, recv, args, kw_bits, jit_entry_idx, block, state, .. } = &**insn;
+            let SendDirectData { cd, cme, iseq, recv, args, kw_bits, jit_entry_idx, block, block_arg, state, .. } = &**insn;
+            let block_arg = block_arg.map(|block_arg| opnd!(block_arg));
             gen_send_iseq_direct(
                 cb, jit, asm,
                 function, *cd, *cme, *iseq, opnd!(recv), opnds!(args),
-                *kw_bits, *jit_entry_idx, &function.frame_state(*state), *block,
+                *kw_bits, *jit_entry_idx, &function.frame_state(*state), *block, block_arg,
             )
         }
         Insn::PushInlineFrame { cme, iseq, recv, num_args, blockiseq, captured, state } => {
@@ -2159,6 +2160,7 @@ fn gen_send_iseq_direct(
     jit_entry_idx: u16,
     state: &FrameState,
     block: Option<BlockHandler>,
+    block_arg: Option<lir::Opnd>,
 ) -> lir::Opnd {
     gen_incr_counter(asm, Counter::iseq_optimized_send_count);
 
@@ -2181,11 +2183,16 @@ fn gen_send_iseq_direct(
     gen_spill_locals(jit, asm, state);
     asm.stack_map(stack_map, jit_frame, state.depth);
 
-    // This mirrors vm_caller_setup_arg_block() in for the `blockiseq != NULL` case.
-    // The HIR specialization guards ensure we will only reach here for literal blocks,
-    // not &block forwarding, &:foo, etc. Thise are rejected in `type_specialize` by
-    // `unspecializable_call_type`.
-    let block_handler = block.map(|bh| match bh { BlockHandler::BlockIseq(b) => gen_block_handler_specval(asm, b), BlockHandler::BlockArg => unreachable!("BlockArg in gen_send_iseq_direct") });
+    // This mirrors vm_caller_setup_arg_block(): `block` is its `blockiseq != NULL` case, and
+    // `block_arg` is a `&blk` argument that `type_specialize` reduced to the block handler that
+    // function would have produced -- the Proc itself, or this frame's own handler for the block
+    // param proxy. Every other kind of block argument, `&:foo` or anything with a `to_proc`,
+    // keeps the dynamic send.
+    let block_handler = match block {
+        Some(BlockHandler::BlockIseq(b)) => Some(gen_block_handler_specval(asm, b)),
+        Some(BlockHandler::BlockArg) => unreachable!("BlockArg in gen_send_iseq_direct"),
+        None => block_arg,
+    };
 
     let callee_is_bmethod = VM_METHOD_TYPE_BMETHOD == unsafe { get_cme_def_type(cme) };
 

zjit/src/codegen_tests.rs

@@ -2312,6 +2312,159 @@ fn test_send_nil_block_arg() {
     "), @"false");
 }
 
+#[test]
+fn test_attr_reader_two_shapes_per_class_in_polymorphic_arm() {
+    // A polymorphic call site branches on the receiver's class, so one arm can be entered by
+    // objects of that class with different shapes. Each shape the profile saw gets its own ivar
+    // load; the rest still have to read correctly through the C fallback.
+    assert_snapshot!(inspect("
+        class C
+          attr_reader :foo
+          def early = (@foo = 1; @bar = 2)
+          def late  = (@bar = 3; @foo = 4)
+          def third = (@baz = 5; @qux = 6; @foo = 7)
+        end
+        class D
+          attr_reader :foo
+          def initialize = @foo = :d
+        end
+        objs = []
+        200.times do |i|
+          c = C.new
+          case i % 3
+          when 0 then c.early
+          when 1 then c.late
+          else c.third
+          end
+          objs << c
+          objs << D.new
+        end
+        def read(o) = o.foo
+        out = objs.map { |o| read(o) }
+        [out.tally.sort_by(&:to_s), read(C.new)]
+    "), @"[[[1, 67], [4, 67], [7, 66], [:d, 200]], nil]");
+}
+
+#[test]
+fn test_send_mixed_nil_and_non_nil_block_arg() {
+    // A `&block` forwarding site that sees both nil and non-nil blocks is split on nil, so the
+    // no-block calls become direct sends. Both branches must still produce the right answer.
+    assert_snapshot!(inspect("
+        def callee(n, &block)
+          block ? block.call(n) : n
+        end
+        def forward(n, &block) = callee(n, &block)
+        results = []
+        100.times do |i|
+          results << forward(i)
+          results << forward(i) { |n| n * 2 }
+        end
+        [forward(7), forward(7) { |n| n + 1 }, results.last(2)]
+    "), @"[7, 8, [99, 198]]");
+}
+
+#[test]
+fn test_send_nil_block_arg_split_polymorphic_receiver() {
+    // The nil branch of the split dispatches on the receiver type, so a forwarding site shared by
+    // several receiver classes still has to pick the right method for each.
+    assert_snapshot!(inspect("
+        class A; def value(&block) = block ? block.call(1) : 1; end
+        class B; def value(&block) = block ? block.call(2) : 2; end
+        def forward(obj, &block) = obj.value(&block)
+        out = []
+        200.times do |i|
+          obj = i.even? ? A.new : B.new
+          out << forward(obj)
+          out << forward(obj) { |n| n * 10 } if i % 3 == 0
+        end
+        [forward(A.new), forward(B.new), forward(A.new) { |n| n * 10 }, out.sum]
+    "), @"[1, 2, 10, 1300]");
+}
+
+#[test]
+fn test_send_forwards_block_param_proxy() {
+    // `bar(&blk)` where `blk` comes from `getblockparamproxy` passes this frame's own block
+    // handler to the callee, so the callee's `yield` and `&b` parameter have to see the block
+    // the outermost caller gave. Every kind of handler goes through the same site: a literal
+    // block, no block at all, a Proc, and a symbol-to-proc.
+    assert_snapshot!(inspect("
+        def callee(n, &b)
+          [n, block_given? ? yield(n) : nil, b ? b.call(n) : nil]
+        end
+        def forward(n, &blk) = callee(n, &blk)
+        def pass_proc(n, p) = forward(n, &p)
+        out = []
+        200.times do |i|
+          out << forward(i) { |x| x + 1 }
+          out << forward(i)
+          out << pass_proc(i, ->(x) { x * 2 })
+          out << pass_proc(i, nil)
+        end
+        out.last(4)
+    "), @"[[199, 200, 200], [199, nil, nil], [199, 398, 398], [199, nil, nil]]");
+}
+
+#[test]
+fn test_send_forwards_block_param_proxy_after_setblockparam() {
+    // Assigning the block parameter makes `getblockparamproxy` hand back the materialized Proc
+    // instead of the proxy, so the branch on the proxy has to send those calls the other way.
+    assert_snapshot!(inspect("
+        def callee(n, &b) = [n, b ? b.call(n) : nil]
+        def forward(n, replace, &blk)
+          blk = ->(x) { x * 100 } if replace
+          callee(n, &blk)
+        end
+        out = []
+        200.times do |i|
+          out << forward(i, false) { |x| x + 1 }
+          out << forward(i, true) { |x| x + 1 }
+        end
+        out.last(2)
+    "), @"[[199, 200], [199, 19900]]");
+}
+
+#[test]
+fn test_send_proc_block_arg_passes_through() {
+    // A `&blk` argument holding a plain Proc is its own block handler in the interpreter, so the
+    // direct send installs it as the callee frame's specval.
+    assert_snapshot!(inspect("
+        def callee(n, &b) = [n, block_given?, b.call(n)]
+        def entry(n, p) = callee(n, &p)
+        doubler = ->(x) { x * 2 }
+        out = nil
+        200.times { |i| out = entry(i, doubler) }
+        [out, entry(5, proc { |x| x + 1 })]
+    "), @"[[199, true, 398], [5, true, 6]]");
+}
+
+#[test]
+fn test_send_proc_block_arg_guard_rejects_other_block_args() {
+    // The Proc guard has to send a `&:sym` or a Method through the interpreter, which converts it
+    // with `to_proc` rather than using it as the block handler directly.
+    assert_snapshot!(inspect("
+        def callee(n, &b) = b.call(n)
+        def entry(n, p) = callee(n, &p)
+        doubler = ->(x) { x * 2 }
+        200.times { |i| entry(i, doubler) }
+        [entry(5, doubler), entry(-6, :abs), entry(7, 2.method(:+))]
+    "), @"[10, 6, 9]");
+}
+
+#[test]
+fn test_send_block_param_proxy_from_block_body() {
+    // `foo(&blk)` inside a block reads the block parameter of the enclosing method, and
+    // `VM_CF_BLOCK_HANDLER` resolves through the local EP to that same frame.
+    assert_snapshot!(inspect("
+        def callee(n, &b) = [n, b ? b.call(n) : nil]
+        def forward(n, &blk)
+          [1].map { callee(n, &blk) }.first
+        end
+        out = nil
+        200.times { |i| out = forward(i) { |x| x + 3 } }
+        [out, forward(9)]
+    "), @"[[199, 202], [9, nil]]");
+}
+
 #[test]
 fn test_send_symbol_block_arg() {
     assert_snapshot!(inspect("