Fix auto-splatting when passing methods as blocks

marshalium · marshalium · commit efdb8a98dee5 · 2020-10-27T11:52:56.000-07:00
This fixes issue #5896 by always splatting out the arguments if possible.
diff --git a/core/src/main/java/org/jruby/runtime/Helpers.java b/core/src/main/java/org/jruby/runtime/Helpers.java
@@ -2439,15 +2439,17 @@ public static RubyString appendAsString(RubyString target, IRubyObject other) {
     // . Array with multiple values and NO rest should extract args if there are more than one argument
 
     static IRubyObject[] restructureBlockArgs(ThreadContext context,
-        IRubyObject value, Signature signature, Block.Type type, boolean needsSplat) {
+        IRubyObject value, Signature signature, Block.Type type) {
 
         if (!type.checkArity && signature == Signature.NO_ARGUMENTS) return IRubyObject.NULL_ARRAY;
 
         if (value == null) return IRubyObject.NULL_ARRAY;
 
-        if (needsSplat) {
+        if (signature != Signature.ONE_ARGUMENT) {
             IRubyObject ary = Helpers.aryToAry(context, value);
-            if (ary instanceof RubyArray) return ((RubyArray) ary).toJavaArrayMaybeUnsafe();
+            if (ary instanceof RubyArray) {
+                return ((RubyArray) ary).toJavaArrayMaybeUnsafe();
+            }
         }
 
         return new IRubyObject[] { value };
diff --git a/core/src/main/java/org/jruby/runtime/MethodBlockBody.java b/core/src/main/java/org/jruby/runtime/MethodBlockBody.java
@@ -64,7 +64,7 @@ public IRubyObject call(ThreadContext context, Block block, IRubyObject[] args,
 
     @Override
     protected IRubyObject doYield(ThreadContext context, Block block, IRubyObject value) {
-        IRubyObject[] realArgs = Helpers.restructureBlockArgs(context, value, getSignature(), block.type, false);
+        IRubyObject[] realArgs = Helpers.restructureBlockArgs(context, value, getSignature(), block.type);
         return entry.method.call(context, receiver, entry.sourceModule, originName, realArgs, Block.NULL_BLOCK);
     }
 
diff --git a/test/jruby/test_hash.rb b/test/jruby/test_hash.rb
@@ -108,4 +108,25 @@ def test_compare_by_identity
     assert_equal 2, hash[arr2]
   end
 
+  def one_arg_method(a)
+    a
+  end
+
+  def two_arg_method(a, b)
+    [a, b]
+  end
+
+  def test_each_yield_to_method
+    hsh = { a: 1, b: 2}
+    hsh.each { |a| a }
+    hsh.each { |a, b| [a, b]}
+
+    one_arg_lambda = lambda { |a| a }
+    two_arg_lambda = lambda { |a, b| [a, b] }
+    hsh.each(&one_arg_lambda)
+    hsh.each(&two_arg_lambda)
+
+    hsh.each(&method(:one_arg_method))
+    hsh.each(&method(:two_arg_method))
+  end
 end

Original file line number	Diff line number	Diff line change
`@@ -64,7 +64,7 @@ public IRubyObject call(ThreadContext context, Block block, IRubyObject[] args,`
`64`	`64`
`65`	`65`	`@Override`
`66`	`66`	`protected IRubyObject doYield(ThreadContext context, Block block, IRubyObject value) {`
`67`		`- IRubyObject[] realArgs = Helpers.restructureBlockArgs(context, value, getSignature(), block.type, false);`
	`67`	`+ IRubyObject[] realArgs = Helpers.restructureBlockArgs(context, value, getSignature(), block.type);`
`68`	`68`	`return entry.method.call(context, receiver, entry.sourceModule, originName, realArgs, Block.NULL_BLOCK);`
`69`	`69`	`}`
`70`	`70`