[fix] Enumerable#take does not require an arg

... `[[]].to_enum.take(1)` regressed in 9.2.9 (GH-5968) but the required signature was incorrect along the way
jruby · Dec 28, 2019 · c7bd47c · c7bd47c
1 parent be2da73
commit c7bd47c
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 3 deletions.
diff --git a/core/src/main/java/org/jruby/RubyEnumerable.java b/core/src/main/java/org/jruby/RubyEnumerable.java
@@ -314,14 +314,14 @@ public static IRubyObject take(ThreadContext context, IRubyObject self, IRubyObj
         final RubyArray result = runtime.newArray();
 
         try {
-            each(context, self, new JavaInternalBlockBody(runtime, Signature.ONE_REQUIRED) {
+            each(context, self, new JavaInternalBlockBody(runtime, Signature.OPTIONAL) {
                 long i = len; // Atomic ?
                 @Override
                 public IRubyObject yield(ThreadContext context, IRubyObject[] args) {
-                    return doYield(context, null, packEnumValues(context, args));
+                    return yield(context, packEnumValues(context, args));
                 }
                 @Override
-                protected IRubyObject doYield(ThreadContext context, Block unused, IRubyObject value) {
+                public IRubyObject yield(ThreadContext context, IRubyObject value) {
                     synchronized (result) {
                         result.append(value);
                         if (--i == 0) throw JumpException.SPECIAL_JUMP;

diff --git a/test/jruby/test_enumerator.rb b/test/jruby/test_enumerator.rb
@@ -287,4 +287,17 @@ def ary.to_ary; [5, 6]; end
     # assert_raise_with_message(TypeError, /C\u{1f5ff}/) {(1..1).zip(obj)}
   end
 
+  def test_take
+    assert_equal [[]], [[]].to_enum.take(1)
+    assert_equal [[]], [[]].to_enum.take(3)
+    assert_equal [], [[]].to_enum.take(0)
+
+    a = [[1], 2, [[3]], [4]]
+    assert_equal [[1]], a.to_enum.take(1)
+    assert_equal [[1], 2, [[3]]], a.to_enum.take(3)
+    assert_equal a, a.to_enum.take(10)
+
+    assert_equal [], [].to_enum.take(1)
+  end
+
 end