ZJIT: Dispatch a yield to a block that can break directly

`direct_invoke_block_adapt()` refused any block containing a `throw`
other than a plain non-local `return`, so every `yield` to a block with
a `break` in it called `rb_vm_invokeblock()` for the life of the
process. The refusal was there because `break` came out as an orphan
("break from proc-closure") or a segfault, which the previous commit
traces to `vm_throw_start()` reading the raw `cfp->_iseq` of a frame
ZJIT pushed. With that read going through `CFP_ISEQ()`, the throw
unwinds out of a JIT-pushed block frame the same way it unwinds out of
one the interpreter pushed: `rb_zjit_throw()` longjmps past every JIT
native frame to the enclosing `vm_exec()`, which resumes at the catch
entry the throw resolved to.

`Kernel#loop` and `Integer#downto` are the two sites this matters for on
lobsters, and both are monomorphic in their block, so they were already
one gate away from the direct dispatch.

lobsters, 15 iterations:

  rb_vm_invokeblock          1,350,799 -> 1,223,051  (-9.5%)
  invokeblock_may_throw        127,811 ->        0

Two master snapshots renumbered: optimize_send_with_block and
optimize_send_to_aliased_cfunc now inline Array#map, because inlining it is what
unlocks the direct yield inside it.

Author: k0kubun

zjit/src/codegen_tests.rs

@@ -8846,6 +8846,104 @@ fn test_invokeblock_truncated_block_with_return() {
     assert_snapshot!(assert_compiles_allowing_exits("entry"), @":returned");
 }
 
+/// A block that `break`s is dispatched directly, not through `rb_vm_invokeblock()`. The
+/// throw unwinds out of the JIT-pushed block frame, which reports its ISEQ through the
+/// JITFrame rather than `cfp->_iseq`.
+#[test]
+fn test_invokeblock_direct_dispatch_with_break() {
+    eval("
+        def test
+          yield 1
+          yield 2
+          :not_reached
+        end
+        def entry
+          test { |x| break x * 10 if x == 2 }
+        end
+        entry; entry
+    ");
+    assert_contains_opcode("test", YARVINSN_invokeblock);
+    assert_snapshot!(assert_compiles_allowing_exits("entry"), @"20");
+}
+
+/// `break` out of a directly dispatched block nested two `yield`s deep unwinds only out of
+/// the `yield` that owns that block, leaving the outer one to run to completion.
+#[test]
+fn test_invokeblock_direct_dispatch_with_nested_break() {
+    eval("
+        def test
+          yield 1
+          :after
+        end
+        def entry
+          test { |a| test { |b| break [:inner, a, b] } }
+        end
+        entry; entry
+    ");
+    assert_contains_opcode("test", YARVINSN_invokeblock);
+    assert_snapshot!(assert_compiles_allowing_exits("entry"), @":after");
+}
+
+/// `break` inside a lambda is a `return` from the lambda, not an unwind to the block owner.
+#[test]
+fn test_invokeblock_direct_dispatch_break_in_lambda() {
+    eval("
+        def test
+          yield 1
+        end
+        def entry
+          l = lambda { break :from_lambda }
+          [test { |x| x }, l.call]
+        end
+        entry; entry
+    ");
+    assert_contains_opcode("test", YARVINSN_invokeblock);
+    assert_snapshot!(assert_compiles_allowing_exits("entry"), @"[1, :from_lambda]");
+}
+
+/// A `break` whose block outlived the method that created it is still an orphan.
+#[test]
+fn test_invokeblock_direct_dispatch_orphan_break() {
+    eval("
+        def make(&b) = b
+        def test
+          yield 1
+        end
+        def entry
+          test { |x| x }
+          pr = make { break :nope }
+          begin
+            pr.call
+          rescue LocalJumpError
+            :orphan
+          end
+        end
+        entry; entry
+    ");
+    assert_contains_opcode("test", YARVINSN_invokeblock);
+    assert_snapshot!(assert_compiles_allowing_exits("entry"), @":orphan");
+}
+
+/// An `ensure` inside a block that `break`s still runs while the throw unwinds out of the
+/// JIT-pushed block frame.
+#[test]
+fn test_invokeblock_direct_dispatch_break_runs_ensure() {
+    eval("
+        def test
+          yield 1
+          :not_reached
+        end
+        def entry
+          ran = false
+          out = test { |x| begin; break :broke; ensure; ran = true; end }
+          [out, ran]
+        end
+        entry; entry
+    ");
+    assert_contains_opcode("test", YARVINSN_invokeblock);
+    assert_snapshot!(assert_compiles_allowing_exits("entry"), @"[:broke, true]");
+}
+
 /// The iterator is inlined into the caller and its `yield` reshapes the arguments for the
 /// block, whose body is then inlined at the yield too. The frame that push lays out has to
 /// follow the reshaped arguments, not the interpreter's stack, or the frame the block raises

zjit/src/hir.rs

@@ -3918,15 +3918,11 @@ fn direct_invoke_block_adapt(iseq: IseqPtr, argc: usize) -> Result<BlockArgAdapt
         std::cmp::Ordering::Less => BlockArgAdapt::NilFill(lead_num - argc),
         std::cmp::Ordering::Greater => BlockArgAdapt::Truncate(lead_num),
     };
-    // `break` out of a directly-invoked block frame does not unwind correctly:
-    // vm_throw_start() matches the block owner's `cfp->pc` against the CATCH_TYPE_BREAK
-    // entry's `cont`, and the PC the owner's frame reports after this dispatch does not
-    // match, so the break is reported as an orphan ("break from proc-closure").
-    // A plain non-local `return` is looked up by frame type and EP instead of by PC, so
-    // blocks that only throw TAG_RETURN -- what this dispatch is for -- are fine.
-    if crate::codegen::block_iseq_may_throw(iseq) && !block_iseq_throws_only_return(iseq) {
-        return Err(InvokeBlockMayThrow);
-    }
+    // A `throw` out of the block frame this dispatch pushes unwinds like any other: it
+    // longjmps out of every JIT native frame to the enclosing `vm_exec()`, which resumes at
+    // the catch entry the throw resolved to. `break` needed `vm_throw_start()` to read the
+    // throwing frame's ISEQ through `CFP_ISEQ()` rather than the raw `cfp->_iseq` a
+    // JIT-pushed frame never writes; without that it walked off a null ISEQ.
     Ok(adapt)
 }