1
0
Fork 0
mirror of https://github.com/ruby/ruby.git synced 2022-11-09 12:17:21 -05:00

* 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
This commit is contained in:
marcandre 2012-11-06 17:16:29 +00:00
parent 5dbbfc3b34
commit 9aafa954aa
2 changed files with 15 additions and 1 deletions

View file

@ -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

View file

@ -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