diff options
author | shugo <shugo@b2dd03c8-39d4-4d8f-98ff-823fe69b080e> | 2012-03-14 10:29:25 +0000 |
---|---|---|
committer | shugo <shugo@b2dd03c8-39d4-4d8f-98ff-823fe69b080e> | 2012-03-14 10:29:25 +0000 |
commit | 0b35e9e5ddb329dac027d72851db015a738852d7 (patch) | |
tree | 001a9eb630f1b96357a951357e9b3b2607149a58 /enumerator.c | |
parent | 0d94434eca76abeb13dda88d2695777bbca7f31a (diff) |
* enumerator.c (lazy_take): add Enumerable::Lazy#take.
git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@35018 b2dd03c8-39d4-4d8f-98ff-823fe69b080e
Diffstat (limited to 'enumerator.c')
-rw-r--r-- | enumerator.c | 34 |
1 files changed, 33 insertions, 1 deletions
diff --git a/enumerator.c b/enumerator.c index 3176b5f756..45fa4f89e6 100644 --- a/enumerator.c +++ b/enumerator.c @@ -13,6 +13,7 @@ ************************************************/ #include "ruby/ruby.h" +#include "node.h" #include "internal.h" /* @@ -1169,7 +1170,10 @@ lazy_init_iterator(VALUE val, VALUE m, int argc, VALUE *argv) static VALUE lazy_init_yielder(VALUE val, VALUE m, int argc, VALUE *argv) { - return rb_funcall2(m, id_yield, 1, &val); + VALUE result; + result = rb_funcall2(m, id_yield, 1, &val); + if (result == Qundef) rb_iter_break(); + return result; } static VALUE @@ -1399,6 +1403,33 @@ lazy_zip(int argc, VALUE *argv, VALUE obj) } static VALUE +lazy_take_func(VALUE val, VALUE args, int argc, VALUE *argv) +{ + NODE *memo = RNODE(args); + + if (memo->u3.cnt == 0) { + return Qundef; + } + rb_funcall2(argv[0], id_yield, argc - 1, argv + 1); + memo->u3.cnt--; + return Qnil; +} + +static VALUE +lazy_take(VALUE obj, VALUE n) +{ + NODE *memo; + long len = NUM2LONG(n); + + if (len < 0) { + rb_raise(rb_eArgError, "attempt to take negative size"); + } + memo = NEW_MEMO(0, 0, len); + return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_take_func, + (VALUE) memo); +} + +static VALUE lazy_lazy(VALUE obj) { return obj; @@ -1492,6 +1523,7 @@ InitVM_Enumerator(void) rb_define_method(rb_cLazy, "reject", lazy_reject, 0); rb_define_method(rb_cLazy, "grep", lazy_grep, 1); rb_define_method(rb_cLazy, "zip", lazy_zip, -1); + rb_define_method(rb_cLazy, "take", lazy_take, 1); rb_define_method(rb_cLazy, "lazy", lazy_lazy, 0); rb_define_alias(rb_cLazy, "collect", "map"); |