Skip to content

HTTPS clone URL

Subversion checkout URL

You can clone with HTTPS or Subversion.

Download ZIP
Browse files

* 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
  • Loading branch information...
commit 0b35e9e5ddb329dac027d72851db015a738852d7 1 parent 0d94434
@shugo shugo authored
View
4 ChangeLog
@@ -1,3 +1,7 @@
+Wed Mar 14 19:28:40 2012 Shugo Maeda <shugo@ruby-lang.org>
+
+ * enumerator.c (lazy_take): add Enumerable::Lazy#take.
+
Wed Mar 14 18:40:36 2012 Shugo Maeda <shugo@ruby-lang.org>
* enumerator.c: use long for array indices.
View
9 enum.c
@@ -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
View
34 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");
View
8 node.h
@@ -451,6 +451,14 @@ typedef struct RNode {
#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)
View
8 test/ruby/test_lazy_enumerator.rb
@@ -143,4 +143,12 @@ def test_zip_with_block
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
Please sign in to comment.
Something went wrong with that request. Please try again.