ruby-changes:22969
From: shugo <ko1@a...>
Date: Wed, 14 Mar 2012 19:29:37 +0900 (JST)
Subject: [ruby-changes:22969] shugo:r35018 (trunk): * enumerator.c (lazy_take): add Enumerable::Lazy#take.
shugo 2012-03-14 19:29:25 +0900 (Wed, 14 Mar 2012) New Revision: 35018 http://svn.ruby-lang.org/cgi-bin/viewvc.cgi?view=rev&revision=35018 Log: * enumerator.c (lazy_take): add Enumerable::Lazy#take. Modified files: trunk/ChangeLog trunk/enum.c trunk/enumerator.c trunk/node.h trunk/test/ruby/test_lazy_enumerator.rb Index: ChangeLog =================================================================== --- ChangeLog (revision 35017) +++ ChangeLog (revision 35018) @@ -1,3 +1,7 @@ +Wed Mar 14 19:28:40 2012 Shugo Maeda <shugo@r...> + + * enumerator.c (lazy_take): add Enumerable::Lazy#take. + Wed Mar 14 18:40:36 2012 Shugo Maeda <shugo@r...> * enumerator.c: use long for array indices. Index: enumerator.c =================================================================== --- enumerator.c (revision 35017) +++ enumerator.c (revision 35018) @@ -13,6 +13,7 @@ ************************************************/ #include "ruby/ruby.h" +#include "node.h" #include "internal.h" /* @@ -1169,7 +1170,10 @@ 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 @@ } 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 @@ 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"); Index: enum.c =================================================================== --- enum.c (revision 35017) +++ enum.c (revision 35018) @@ -17,15 +17,6 @@ #define STATIC_ASSERT(name, expr) typedef int static_assert_##name##_check[1 - 2*!(expr)] -#define NEW_MEMO(a, b, c) rb_node_newnode(NODE_MEMO, (a), (b), (c)) - -#define roomof(x, y) ((sizeof(x) + sizeof(y) - 1) / sizeof(y)) -#define MEMO_FOR(type, value) ((type *)RARRAY_PTR(value)) -#define NEW_MEMO_FOR(type, value) \ - (rb_ary_set_len(((value) = rb_ary_tmp_new(roomof(type, VALUE))), \ - roomof(type, VALUE)), \ - MEMO_FOR(type, value)) - VALUE rb_mEnumerable; static ID id_next; #define id_each idEach Index: test/ruby/test_lazy_enumerator.rb =================================================================== --- test/ruby/test_lazy_enumerator.rb (revision 35017) +++ test/ruby/test_lazy_enumerator.rb (revision 35018) @@ -143,4 +143,12 @@ assert_equal(["a", 1], a.lazy.zip("a".."c") {|x, y| [y, x]}.first) assert_equal(1, a.current) end + + def test_take + a = Step.new(1..3) + assert_equal(1, a.take(2).first) + assert_equal(2, a.current) + assert_equal(1, a.lazy.take(2).first) + assert_equal(1, a.current) + end end Index: node.h =================================================================== --- node.h (revision 35017) +++ node.h (revision 35018) @@ -451,7 +451,15 @@ #define NEW_ATTRASGN(r,m,a) NEW_NODE(NODE_ATTRASGN,r,m,a) #define NEW_PRELUDE(p,b) NEW_NODE(NODE_PRELUDE,p,b,0) #define NEW_OPTBLOCK(a) NEW_NODE(NODE_OPTBLOCK,a,0,0) +#define NEW_MEMO(a,b,c) NEW_NODE(NODE_MEMO,a,b,c) +#define roomof(x, y) ((sizeof(x) + sizeof(y) - 1) / sizeof(y)) +#define MEMO_FOR(type, value) ((type *)RARRAY_PTR(value)) +#define NEW_MEMO_FOR(type, value) \ + (rb_ary_set_len(((value) = rb_ary_tmp_new(roomof(type, VALUE))), \ + roomof(type, VALUE)), \ + MEMO_FOR(type, value)) + #if defined __GNUC__ && __GNUC__ >= 4 #pragma GCC visibility push(default) #endif -- ML: ruby-changes@q... Info: http://www.atdot.net/~ko1/quickml/