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/