[fix] define_method captured block propagation in JIT mode (jruby#9310)

Author: kares

core/src/main/java/org/jruby/internal/runtime/methods/DefineMethodMethod.java

@@ -22,6 +22,10 @@ public DefineMethodMethod(IRScope method, Visibility visibility, RubyModule impl
         this.capturedBlock = capturedBlock;
     }
 
+    public Block getCapturedBlock() {
+        return capturedBlock;
+    }
+
     @Override
     public IRubyObject call(ThreadContext context, IRubyObject self, RubyModule clazz, String name, IRubyObject[] args, Block block) {
         return super.call(context, self, clazz, name, args, capturedBlock);

core/src/main/java/org/jruby/ir/IRClosure.java

@@ -196,7 +196,6 @@ public boolean isNestedClosuresSafeForMethodConversion() {
     public IRMethod convertToMethod(ByteList name) {
         // We want variable scoping to be the same as a method and not see outside itself.
         if (source == null ||
-            isWithinMethod() ||                 // def foo; define_method {}; end has potential to implicitly access parents var yield
             accessesParentsLocalVariables() ||  // Built methods cannot search down past method scope
             receivesClosureArg() ||             // we pass in captured block at define_method as block so explicits ones not supported
             usesZSuper() ||                     // methods defined from closures cannot use zsuper

core/src/main/java/org/jruby/ir/targets/indy/InvokeSite.java

@@ -21,6 +21,7 @@
 import org.jruby.internal.runtime.methods.AttrReaderMethod;
 import org.jruby.internal.runtime.methods.AttrWriterMethod;
 import org.jruby.internal.runtime.methods.CompiledIRMethod;
+import org.jruby.internal.runtime.methods.DefineMethodMethod;
 import org.jruby.internal.runtime.methods.DynamicMethod;
 import org.jruby.internal.runtime.methods.HandleMethod;
 import org.jruby.internal.runtime.methods.MixedModeIRMethod;
@@ -361,8 +362,8 @@ private MethodHandle createAttrWriterHandle(IRubyObject self, RubyClass cls, Var
         return setValue;
     }
 
-    MethodHandle buildJittedHandle(CacheEntry entry, boolean blockGiven) {
-        MethodHandle mh = null;
+    MethodHandle buildJittedHandle(final CacheEntry entry, final boolean blockGiven) {
+        MethodHandle mh;
         SmartBinder binder;
         CompiledIRMethod compiledIRMethod = null;
         DynamicMethod method = entry.method;
@@ -410,7 +411,12 @@ MethodHandle buildJittedHandle(CacheEntry entry, boolean blockGiven) {
                 }
             }
 
-            if (!blockGiven) {
+            if (method instanceof DefineMethodMethod defineMethod) {
+                // DefineMethodMethod captures the block from define_method's enclosing scope;
+                // always use that captured block for yield, regardless of caller's block
+                if (blockGiven) binder = binder.drop("block");
+                binder = binder.append("block", Block.class, defineMethod.getCapturedBlock());
+            } else if (!blockGiven) {
                 binder = binder.append("block", Block.class, Block.NULL_BLOCK);
             }
 
@@ -424,6 +430,8 @@ MethodHandle buildJittedHandle(CacheEntry entry, boolean blockGiven) {
             if (Options.INVOKEDYNAMIC_LOG_BINDING.load()) {
                 LOG.info(name() + "\tbound directly to jitted method " + Bootstrap.logMethod(method));
             }
+        } else {
+            mh = null;
         }
 
         return mh;

spec/regression/GH-8946_define_method_yield_with_jit_spec.rb

@@ -0,0 +1,45 @@
+# https://github.com/jruby/jruby/issues/8946
+#
+# When define_method is called within another method and the block uses yield,
+# the captured block must be used for yield even after JIT + invokedynamic
+# optimization (which previously bypassed DefineMethodMethod.call and lost
+# the captured block).
+
+describe 'define_method with yield inside another method' do
+  it 'yields to the captured block after JIT optimization' do
+    cls = Class.new do
+      class << self
+        def go(&blk)
+          define_method(:foo) { yield }
+        end
+      end
+    end
+
+    results = []
+    cls.go { results << :called }
+
+    obj = cls.new
+    # call enough times to trigger JIT and indy binding
+    10.times { obj.foo }
+
+    expect(results).to eq(Array.new(10, :called))
+  end
+
+  it 'yields to the captured block with arguments' do
+    cls = Class.new do
+      class << self
+        def go(&blk)
+          define_method(:bar) { |x| yield(x) }
+        end
+      end
+    end
+
+    results = []
+    cls.go { |v| results << v }
+
+    obj = cls.new
+    10.times { |i| obj.bar(i) }
+
+    expect(results).to eq((0..9).to_a)
+  end
+end