* enumerator.c: Support for lazy.take.size

  [Feature #6636]

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@37523 b2dd03c8-39d4-4d8f-98ff-823fe69b080e

2 files changed, 15 insertions, 1 deletions

diff --git a/enumerator.c b/enumerator.c
index 1324c9284c..e5dd80a841 100644
--- a/enumerator.c
+++ b/enumerator.c
@@ -1625,6 +1625,15 @@ lazy_take_func(VALUE val, VALUE args, int argc, VALUE *argv)
 }
 
 static VALUE
+lazy_take_size(VALUE lazy) {
+    long len = NUM2LONG(RARRAY_PTR(rb_ivar_get(lazy, id_arguments))[0]);
+    VALUE receiver = lazy_receiver_size(lazy);
+    if (NIL_P(receiver) || (FIXNUM_P(receiver) && FIX2LONG(receiver) < len))
+	return receiver;
+    return LONG2NUM(len);
+}
+
+static VALUE
 lazy_take(VALUE obj, VALUE n)
 {
     NODE *memo;
@@ -1644,7 +1653,7 @@ lazy_take(VALUE obj, VALUE n)
     memo = NEW_MEMO(0, len, len);
     return lazy_set_method(rb_block_call(rb_cLazy, id_new, argc, argv,
 					 lazy_take_func, (VALUE) memo),
-			   rb_ary_new3(1, n), 0);
+			   rb_ary_new3(1, n), lazy_take_size);
 }
 
 static VALUE
diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb
index 8c03c700c3..e09c3af837 100644
--- a/test/ruby/test_lazy_enumerator.rb
+++ b/test/ruby/test_lazy_enumerator.rb
@@ -336,5 +336,10 @@ EOS
       assert_equal nil, lazy.send(m){}.size
     end
     assert_equal nil, lazy.grep(//).size
+
+    assert_equal 2, lazy.take(2).size
+    assert_equal 3, lazy.take(4).size
+    assert_equal 4, loop.lazy.take(4).size
+    assert_equal nil, lazy.select{}.take(4).size
   end
 end