diff options
author | marcandre <marcandre@b2dd03c8-39d4-4d8f-98ff-823fe69b080e> | 2012-11-06 17:15:59 +0000 |
---|---|---|
committer | marcandre <marcandre@b2dd03c8-39d4-4d8f-98ff-823fe69b080e> | 2012-11-06 17:15:59 +0000 |
commit | c8426ce840bf5329fd31e8696d4f225a207b6550 (patch) | |
tree | 1b549e6415558dfa4908a633535b8e429f7dd077 | |
parent | 44374034268913246d517ff25365e0bccb453e02 (diff) | |
download | ruby-c8426ce840bf5329fd31e8696d4f225a207b6550.tar.gz |
* enumerator.c: Support for lazy.size
[Feature #6636]
git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@37521 b2dd03c8-39d4-4d8f-98ff-823fe69b080e
-rw-r--r-- | enumerator.c | 23 | ||||
-rw-r--r-- | test/ruby/test_lazy_enumerator.rb | 6 |
2 files changed, 27 insertions, 2 deletions
diff --git a/enumerator.c b/enumerator.c index a5eb999d8d..582ec69111 100644 --- a/enumerator.c +++ b/enumerator.c @@ -104,7 +104,7 @@ */ VALUE rb_cEnumerator; VALUE rb_cLazy; -static ID id_rewind, id_each, id_new, id_initialize, id_yield, id_call; +static ID id_rewind, id_each, id_new, id_initialize, id_yield, id_call, id_size; static ID id_eqq, id_next, id_result, id_lazy, id_receiver, id_arguments, id_method; static VALUE sym_each, sym_cycle; @@ -1240,6 +1240,23 @@ generator_each(int argc, VALUE *argv, VALUE obj) /* Lazy Enumerator methods */ static VALUE +lazy_receiver_size(VALUE self) +{ + VALUE r = rb_check_funcall(rb_ivar_get(self, id_receiver), id_size, 0, 0); + return (r == Qundef) ? Qnil : r; +} + +static VALUE +lazy_size(VALUE self) +{ + struct enumerator *e = enumerator_ptr(self); + if (e->size_fn) { + return (*e->size_fn)(self); + } + return Qnil; +} + +static VALUE lazy_init_iterator(VALUE val, VALUE m, int argc, VALUE *argv) { VALUE result; @@ -1313,7 +1330,7 @@ lazy_initialize(int argc, VALUE *argv, VALUE self) rb_block_call(generator, id_initialize, 0, 0, (rb_block_given_p() ? lazy_init_block_i : lazy_init_block), obj); - enumerator_init(self, generator, meth, argc - offset, argv + offset, 0, Qnil); + enumerator_init(self, generator, meth, argc - offset, argv + offset, lazy_receiver_size, Qnil); rb_ivar_set(self, id_receiver, obj); return self; @@ -1820,6 +1837,7 @@ InitVM_Enumerator(void) rb_define_method(rb_cLazy, "flat_map", lazy_flat_map, 0); rb_define_method(rb_cLazy, "collect_concat", lazy_flat_map, 0); rb_define_method(rb_cLazy, "select", lazy_select, 0); + rb_define_method(rb_cLazy, "size", lazy_size, 0); rb_define_method(rb_cLazy, "find_all", lazy_select, 0); rb_define_method(rb_cLazy, "reject", lazy_reject, 0); rb_define_method(rb_cLazy, "grep", lazy_grep, 1); @@ -1860,6 +1878,7 @@ Init_Enumerator(void) id_rewind = rb_intern("rewind"); id_each = rb_intern("each"); id_call = rb_intern("call"); + id_size = rb_intern("size"); id_yield = rb_intern("yield"); id_new = rb_intern("new"); id_initialize = rb_intern("initialize"); diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb index 106f8006f9..754b6aa8e0 100644 --- a/test/ruby/test_lazy_enumerator.rb +++ b/test/ruby/test_lazy_enumerator.rb @@ -323,4 +323,10 @@ class TestLazyEnumerator < Test::Unit::TestCase #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: #<Enumerator::Lazy: 1..10>:map>:collect>:flat_map>:collect_concat>:select>:find_all>:reject>:grep(1)>:zip("a".."c")>:take(10)>:take_while>:drop(3)>:drop_while>:cycle(3)> EOS end + + def test_size + lazy = [1, 2, 3].lazy + assert_equal 3, lazy.size + assert_equal 42, Enumerator.new(42){}.lazy.size + end end |