YJIT: Fix super() from a block method not rooted in a "def" method

Previously, we unconditionally passed to the callee
`GET_LEP(calling_frame)[VM_ENV_DATA_INDEX_SPECVAL]`, which in case the
block containing super() is in e.g. a `class`, did not resolve to a
block handler at all.

Properly locate the block handler using logic used in
Kernel#block_given? and defined?(yield). Rename BlockHandler::LEPSpecVal
to FindFromCurrentFrame to reflect this change, as the block handler is
not always in LEP.

[Bug #22116]

Author: XrXr

test/ruby/test_yjit.rb

@@ -1041,6 +1041,55 @@ class Target < Base; end
     RUBY
   end
 
+  def test_bmethod_super_block_forwarding
+    # super() in a block method forwards the block handler of environment
+    # it is in. In this test, the block is in a class, which never has
+    # a block handler, so Parent#foo gets no block.
+    assert_compiles(<<~RUBY, result: nil)
+      class Parent
+        def foo = defined?(yield)
+      end
+
+      class BmethodSuper < Parent
+        define_method(:foo) { super() }
+      end
+
+      BmethodSuper.new.foo {}
+    RUBY
+
+    # For contrast, the super() here forwards the block passed to add_then_override
+    # because block with super() is rooted in add_then_override, a "def" method.
+    assert_compiles(<<~RUBY, result: [1, 2])
+      def add_then_override(object)
+        object.define_singleton_method(:then) { super() }
+      end
+
+      add_then_override(one = Object.new) { 1 }
+      add_then_override(two = Object.new) { 2 }
+      [one.then {}, two.then {}]
+    RUBY
+  end
+
+  def test_bug_22116
+    # Regression test from report which used to crash during environment escape
+    # to pass block to Hash.new
+    assert_compiles(<<~RUBY, call_threshold: 1)
+      class C
+        def foo
+          Hash.new {|h, k| h[k] = [] }
+        end
+      end
+
+      class D < C
+        define_method(:foo) do
+          super()
+        end
+      end
+
+      D.new.foo
+    RUBY
+  end
+
   # Tests calling a variadic cfunc with many args
   def test_build_large_struct
     assert_compiles(<<~RUBY, insns: %i[opt_send_without_block], call_threshold: 2)

yjit/src/codegen.rs

@@ -6805,8 +6805,8 @@ enum SpecVal {
 pub enum BlockHandler {
     // send, invokesuper: blockiseq operand
     BlockISeq(IseqPtr),
-    // invokesuper: GET_BLOCK_HANDLER() (GET_LEP()[VM_ENV_DATA_INDEX_SPECVAL])
-    LEPSpecVal,
+    // For invokesuper; equivalent to GET_BLOCK_HANDLER()
+    FindFromCurrentFrame,
     // part of the allocate-free block forwarding scheme
     BlockParamProxy,
     // To avoid holding the block arg (e.g. proc and symbol) across C calls,
@@ -6870,9 +6870,20 @@ fn gen_push_frame(
                     let cfp_self = asm.lea(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SELF));
                     asm.or(cfp_self, Opnd::Imm(1))
                 }
-                BlockHandler::LEPSpecVal => {
-                    let lep_opnd = gen_get_lep(jit, asm);
-                    asm.load(Opnd::mem(64, lep_opnd, SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_SPECVAL))
+                BlockHandler::FindFromCurrentFrame => {
+                    // Find the calling frame's block handler. Similar shape to implementation for
+                    // defined?(yield). See also: vm_caller_setup_arg_block() with is_super, which
+                    // ends up in VM_CF_BLOCK_HANDLER().
+                    //
+                    // When the calling ISEQ is rooted in a ISEQ_TYPE_METHOD, the block handler is
+                    // at local_ep[VM_ENV_DATA_INDEX_SPECVAL]. Otherwise, the frame has no block
+                    // handler.
+                    if ISEQ_TYPE_METHOD == unsafe { rb_get_iseq_body_type(rb_get_iseq_body_local_iseq(jit.iseq)) } {
+                        let lep_opnd = gen_get_lep(jit, asm);
+                        asm.load(Opnd::mem(64, lep_opnd, SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_SPECVAL))
+                    } else {
+                        VM_BLOCK_HANDLER_NONE.into()
+                    }
                 }
                 BlockHandler::BlockParamProxy => {
                     let ep_opnd = gen_get_lep(jit, asm);
@@ -9880,7 +9891,7 @@ fn gen_invokesuper_specialized(
     let block = if let Some(iseq) = jit.get_arg(1).as_optional_ptr() {
         BlockHandler::BlockISeq(iseq)
     } else {
-        BlockHandler::LEPSpecVal
+        BlockHandler::FindFromCurrentFrame
     };
 
     // Fallback to dynamic dispatch if this callsite is megamorphic