From 9aafa954aa9ddb54f943c029100c03c4d4a1701b Mon Sep 17 00:00:00 2001 From: marcandre Date: Tue, 6 Nov 2012 17:16:29 +0000 Subject: * 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 --- enumerator.c | 11 ++++++++++- test/ruby/test_lazy_enumerator.rb | 5 +++++ 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/enumerator.c b/enumerator.c index 1324c9284c..e5dd80a841 100644 --- a/enumerator.c +++ b/enumerator.c @@ -1624,6 +1624,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) { @@ -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 -- cgit v1.2.3