From 38558dd95e6443a8412304718ad77a331c185f5d Mon Sep 17 00:00:00 2001
From: Alan Wu <XrXr@users.noreply.github.com>
Date: Wed, 13 Aug 2025 12:39:06 -0400
Subject: [PATCH] YJIT: Fix `defined?(yield)` and `block_given?` at top level

Previously, YJIT returned truthy for the block given query at the top
level. That's incorrect because the top level script never receives a
block, and `yield` is a syntax error there.

Inside methods, the number of hops to get from `iseq` to
`iseq->body->local_iseq` is the same as the number of
`VM_ENV_PREV_EP(ep)` hops to get to an environment with
`VM_ENV_FLAG_LOCAL`. YJIT and the interpreter both rely on this as can
be seen in get_lvar_level(). However, this identity does not hold for
the top level frame because of vm_set_eval_stack(), which sets up
`TOPLEVEL_BINDING`.

Since only methods can take a block that `yield` goes to, have ISEQs
that are the child of a non-method ISEQ return falsy for the block given
query. This fixes the issue for the top level script and is an
optimization for non-method contexts such as inside `ISEQ_TYPE_CLASS`.
---
 bootstraptest/test_yjit.rb | 27 +++++++++++++++++++++++++++
 yjit/src/codegen.rs        | 26 +++++++++++++++++---------
 2 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 3e3936942d..f76af3633d 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -5342,3 +5342,30 @@ assert_equal 'nil', %{
 
   test_local_fill_in_forwardable.inspect
 }
+
+# Test defined?(yield) and block_given? in non-method context.
+# It's good that the body of this runs at true top level and isn't wrapped in a block.
+assert_equal 'false', %{
+  RESULT = []
+  RESULT << defined?(yield)
+  RESULT << block_given?
+
+  1.times do
+    RESULT << defined?(yield)
+    RESULT << block_given?
+  end
+
+  module ModuleContext
+    1.times do
+      RESULT << defined?(yield)
+      RESULT << block_given?
+    end
+  end
+
+  class << self
+    RESULT << defined?(yield)
+    RESULT << block_given?
+  end
+
+  RESULT.any?
+}
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index b8b15adc8b..9644b948d7 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -6656,16 +6656,24 @@ fn gen_block_given(
 ) {
     asm_comment!(asm, "block_given?");
 
-    // Same as rb_vm_frame_block_handler
-    let ep_opnd = gen_get_lep(jit, asm);
-    let block_handler = asm.load(
-        Opnd::mem(64, ep_opnd, SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_SPECVAL)
-    );
+    // `yield` goes to the block handler stowed in the "local" iseq which is
+    // the current iseq or a parent. Only the "method" iseq type can be passed a
+    // block handler. (e.g. `yield` in the top level script is a syntax error.)
+    let local_iseq = unsafe { rb_get_iseq_body_local_iseq(jit.iseq) };
+    if unsafe { rb_get_iseq_body_type(local_iseq) } == ISEQ_TYPE_METHOD {
+        // Same as rb_vm_frame_block_handler
+        let ep_opnd = gen_get_lep(jit, asm);
+        let block_handler = asm.load(
+            Opnd::mem(64, ep_opnd, SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_SPECVAL)
+        );
 
-    // Return `block_handler != VM_BLOCK_HANDLER_NONE`
-    asm.cmp(block_handler, VM_BLOCK_HANDLER_NONE.into());
-    let block_given = asm.csel_ne(true_opnd, false_opnd);
-    asm.mov(out_opnd, block_given);
+        // Return `block_handler != VM_BLOCK_HANDLER_NONE`
+        asm.cmp(block_handler, VM_BLOCK_HANDLER_NONE.into());
+        let block_given = asm.csel_ne(true_opnd, false_opnd);
+        asm.mov(out_opnd, block_given);
+    } else {
+        asm.mov(out_opnd, false_opnd);
+    }
 }
 
 // Codegen for rb_class_superclass()