Merge pull request #5812 from jruby/enum-yield

[fix] Enumerable#first block signature (matches MRI)
jruby · Aug 5, 2019 · 96d8c16 · 96d8c16
2 parents dc64294 + d1048c4
commit 96d8c16
Show file tree

Hide file tree

Showing 7 changed files with 51 additions and 35 deletions.
diff --git a/core/src/main/java/org/jruby/RubyEnumerable.java b/core/src/main/java/org/jruby/RubyEnumerable.java
@@ -406,7 +406,7 @@ public static IRubyObject first(ThreadContext context, IRubyObject self) {
         final IRubyObject[] holder = new IRubyObject[]{ context.nil };
 
         try {
-            each(context, self, new JavaInternalBlockBody(context.runtime, context, null, Signature.ONE_REQUIRED) {
+            each(context, self, new JavaInternalBlockBody(context.runtime, context, "Enumerable#first", Signature.OPTIONAL) {
                 public IRubyObject yield(ThreadContext context, IRubyObject[] args) {
                     IRubyObject packedArg = packEnumValues(context.runtime, args);
                     holder[0] = packedArg;
@@ -428,7 +428,7 @@ public static IRubyObject first(ThreadContext context, IRubyObject self, final I
         final RubyArray result = RubyArray.newArray(runtime, firstCount);
 
         try {
-            each(context, self, new JavaInternalBlockBody(runtime, context, null, Signature.ONE_REQUIRED) {
+            each(context, self, new JavaInternalBlockBody(runtime, context, "Enumerable#first", Signature.OPTIONAL) {
                 private long iter = firstCount;
                 public IRubyObject yield(ThreadContext context, IRubyObject[] args) {
                     IRubyObject packedArg = packEnumValues(context.runtime, args);

diff --git a/core/src/main/java/org/jruby/RubyEnumerator.java b/core/src/main/java/org/jruby/RubyEnumerator.java
@@ -30,29 +30,18 @@
 
 import org.jruby.anno.JRubyMethod;
 import org.jruby.anno.JRubyModule;
-import org.jruby.exceptions.RaiseException;
 import org.jruby.exceptions.StopIteration;
-import org.jruby.exceptions.Unrescuable;
 import org.jruby.runtime.Arity;
-import org.jruby.runtime.Binding;
 import org.jruby.runtime.Block;
-import org.jruby.runtime.BlockCallback;
-import org.jruby.runtime.CallBlock;
-import org.jruby.runtime.Helpers;
 import org.jruby.runtime.JavaSites;
 import org.jruby.runtime.ObjectAllocator;
-import org.jruby.runtime.Signature;
 import org.jruby.runtime.ThreadContext;
 import org.jruby.runtime.builtin.IRubyObject;
 import org.jruby.util.ArraySupport;
 import org.jruby.util.ByteList;
-import org.jruby.util.cli.Options;
 
 import java.util.Arrays;
-import java.util.NoSuchElementException;
 import java.util.Spliterator;
-import java.util.concurrent.Future;
-import java.util.concurrent.SynchronousQueue;
 import java.util.stream.Stream;
 
 import static org.jruby.runtime.Visibility.PRIVATE;

diff --git a/core/src/main/java/org/jruby/RubyHash.java b/core/src/main/java/org/jruby/RubyHash.java
@@ -65,7 +65,6 @@
 import java.io.IOException;
 import java.util.AbstractCollection;
 import java.util.AbstractSet;
-import java.util.Arrays;
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.Map;

diff --git a/core/src/main/java/org/jruby/RubyYielder.java b/core/src/main/java/org/jruby/RubyYielder.java
@@ -28,7 +28,6 @@
 
 import org.jruby.anno.JRubyClass;
 import org.jruby.anno.JRubyMethod;
-import org.jruby.runtime.Arity;
 import org.jruby.runtime.Block;
 import org.jruby.runtime.BlockCallback;
 import org.jruby.runtime.CallBlock19;
@@ -113,17 +112,13 @@ public IRubyObject initialize(ThreadContext context, Block block) {
     }
 
     @JRubyMethod(rest = true)
-    public IRubyObject yield(ThreadContext context, IRubyObject[]args) {
+    public IRubyObject yield(ThreadContext context, IRubyObject[] args) {
         checkInit();
-        return proc.call(context, args, Block.NULL_BLOCK);
+        return proc.call(context, args);
     }
 
     @JRubyMethod(name = "<<", rest = true)
-    public IRubyObject op_lshift(ThreadContext context, IRubyObject[]args) {
-        if (args.length == 1 &&
-                args[0] instanceof RubyArray &&
-                ((RubyArray) args[0]).getLength() == 1)
-            args[0] = RubyArray.newArray(context.runtime, args[0]);
+    public IRubyObject op_lshift(ThreadContext context, IRubyObject[] args) {
         yield(context, args);
         return this;
     }

diff --git a/core/src/main/java/org/jruby/runtime/CallBlock.java b/core/src/main/java/org/jruby/runtime/CallBlock.java
@@ -59,10 +59,10 @@ private CallBlock(Signature signature, BlockCallback callback, ThreadContext con
         this.dummyScope = context.runtime.getStaticScopeFactory().getDummyScope();
     }
 
-    private IRubyObject[] adjustArgs(Block block, IRubyObject[] args) {
+    static IRubyObject[] adjustArgs(Block block, IRubyObject[] args) {
         Signature signature = block.getSignature();
         int required = signature.required();
-        if (signature.isFixed() && required  > 0 && required < args.length) args = ArraySupport.newCopy(args, required);
+        if (required > 0 && required < args.length && signature.isFixed()) args = ArraySupport.newCopy(args, required);
 
         return args;
     }

diff --git a/core/src/main/java/org/jruby/runtime/JavaInternalBlockBody.java b/core/src/main/java/org/jruby/runtime/JavaInternalBlockBody.java
@@ -47,22 +47,14 @@ private void threadCheck(ThreadContext yieldingContext) {
         }
     }
 
-    private IRubyObject[] adjustArgs(Block block, IRubyObject[] args) {
-        Signature signature = block.getSignature();
-        int required = signature.required();
-        if (signature.isFixed() && required  > 0 && required < args.length) args = ArraySupport.newCopy(args, required);
-
-        return args;
-    }
-
     @Override
     public IRubyObject call(ThreadContext context, Block block, IRubyObject[] args) {
-        return yield(context, block, adjustArgs(block, args), null);
+        return yield(context, block, CallBlock.adjustArgs(block, args), null);
     }
 
     @Override
     public IRubyObject call(ThreadContext context, Block block, IRubyObject[] args, Block blockArg) {
-        return yield(context, block, adjustArgs(block, args), null, blockArg);
+        return yield(context, block, CallBlock.adjustArgs(block, args), null, blockArg);
     }
 
     @Override
@@ -76,7 +68,7 @@ protected IRubyObject doYield(ThreadContext context, Block block, IRubyObject va
     protected IRubyObject doYield(ThreadContext context, Block block, IRubyObject[] args, IRubyObject self) {
         threadCheck(context);
 
-        return yield(context, adjustArgs(block, args));
+        return yield(context, CallBlock.adjustArgs(block, args));
     }
 
     public abstract IRubyObject yield(ThreadContext context, IRubyObject[] args);

diff --git a/test/jruby/test_enumerator.rb b/test/jruby/test_enumerator.rb
@@ -187,4 +187,45 @@ def assert_raise(klass = StandardError, &block)
     super
   end
 
+  public
+
+  def test_yielder
+    ary = Enumerator.new { |y| y.yield([1]) }.to_a
+    assert_equal [[1]], ary
+    ary = Enumerator.new { |y| y.yield([1]) }.lazy.to_a
+    assert_equal [[1]], ary
+    ary = Enumerator.new { |y| y << [1] }.to_a
+    assert_equal [[1]], ary
+
+    yields = []
+    y = Enumerator::Yielder.new { |args| yields << args }
+    y << [1]
+    assert_equal [[1]], yields
+
+    assert_equal 42, Enumerator.new { |y| y << 42 }.first
+    assert_equal [42], Enumerator.new { |y| y << [42] }.first
+    assert_equal [[42]], Enumerator.new { |y| y << [42] }.first(1)
+
+    assert_equal [], Enumerator.new { |y| y << [] }.first
+    assert_equal [], Enumerator.new { |y| y.yield [] }.first
+    assert_equal [1], Enumerator.new { |y| y.yield [1] }.first
+  end
+
+  def test_yield_map # GH-4108
+    ary = Enumerator.new { |y| y.yield([1]) }.map { |e| e }.to_a
+    assert_equal [[1]], ary
+    ary = Enumerator.new { |y| y.yield([1]) }.lazy.map { |e| e }.to_a
+    # NOTE: this one seems still failing (enumerable.lazy.map works)
+    # assert_equal [[1]], ary
+  end
+
+  def test_lazy_map_enumerable # GH-5044
+    a = [1, 2, 3].map { |a| [a + 1] }.to_a
+    assert_equal [[2], [3], [4]], a
+    a = [1, 2, 3].lazy.map { |a| [a + 1] }.to_a
+    assert_equal [[2], [3], [4]], a
+    a = [[1], [2], [3]].lazy.map(&:itself).to_a
+    assert_equal [[1], [2], [3]], a
+  end
+
 end