diff options
author | marcandre <marcandre@b2dd03c8-39d4-4d8f-98ff-823fe69b080e> | 2012-11-06 17:16:44 +0000 |
---|---|---|
committer | marcandre <marcandre@b2dd03c8-39d4-4d8f-98ff-823fe69b080e> | 2012-11-06 17:16:44 +0000 |
commit | 7a31096255cd226d16c918fa859e3e4a654e64b9 (patch) | |
tree | c8a2ff5a6613b2742ce82fad6f20d7f6231942aa | |
parent | 9aafa954aa9ddb54f943c029100c03c4d4a1701b (diff) |
* enumerator.c: Add support for lazy.drop.size
[Feature #6636]
git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@37524 b2dd03c8-39d4-4d8f-98ff-823fe69b080e
-rw-r--r-- | enumerator.c | 15 | ||||
-rw-r--r-- | test/ruby/test_lazy_enumerator.rb | 5 |
2 files changed, 19 insertions, 1 deletions
diff --git a/enumerator.c b/enumerator.c index e5dd80a841..0a51709ac1 100644 --- a/enumerator.c +++ b/enumerator.c @@ -1674,6 +1674,19 @@ lazy_take_while(VALUE obj) } static VALUE +lazy_drop_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)) + return receiver; + if (FIXNUM_P(receiver)) { + len = FIX2LONG(receiver) - len; + return LONG2FIX(len < 0 ? 0 : len); + } + return rb_funcall(receiver, '-', 1, LONG2NUM(len)); +} + +static VALUE lazy_drop_func(VALUE val, VALUE args, int argc, VALUE *argv) { NODE *memo = RNODE(args); @@ -1699,7 +1712,7 @@ lazy_drop(VALUE obj, VALUE n) memo = NEW_MEMO(0, 0, len); return lazy_set_method(rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_drop_func, (VALUE) memo), - rb_ary_new3(1, n), 0); + rb_ary_new3(1, n), lazy_drop_size); } static VALUE diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb index e09c3af837..9647590136 100644 --- a/test/ruby/test_lazy_enumerator.rb +++ b/test/ruby/test_lazy_enumerator.rb @@ -341,5 +341,10 @@ EOS assert_equal 3, lazy.take(4).size assert_equal 4, loop.lazy.take(4).size assert_equal nil, lazy.select{}.take(4).size + + assert_equal 1, lazy.drop(2).size + assert_equal 0, lazy.drop(4).size + assert_equal Float::INFINITY, loop.lazy.drop(4).size + assert_equal nil, lazy.select{}.drop(4).size end end |