diff options
author | matz <matz@b2dd03c8-39d4-4d8f-98ff-823fe69b080e> | 2007-02-27 09:19:21 +0000 |
---|---|---|
committer | matz <matz@b2dd03c8-39d4-4d8f-98ff-823fe69b080e> | 2007-02-27 09:19:21 +0000 |
commit | 60c3c901fe536de977368580f36576f2948e7b60 (patch) | |
tree | 0bb00b7d8b1bc06182d6495c301f2917ec462ebc /enum.c | |
parent | 8a1b1c2b41b475cc40a8013623b5b6fe5381d6da (diff) | |
download | ruby-60c3c901fe536de977368580f36576f2948e7b60.tar.gz |
* enum.c (enum_take): new method. [ruby-dev:30407]
* enum.c (enum_drop): ditto.
git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@11901 b2dd03c8-39d4-4d8f-98ff-823fe69b080e
Diffstat (limited to 'enum.c')
-rw-r--r-- | enum.c | 119 |
1 files changed, 119 insertions, 0 deletions
@@ -1221,6 +1221,123 @@ enum_zip(int argc, VALUE *argv, VALUE obj) return result; } +static VALUE +take_i(VALUE i, VALUE *arg) +{ + long n = arg[1]; + + if (arg[1]-- == 0) rb_iter_break(); + rb_ary_push(arg[0], i); + return Qnil; +} + +static VALUE +take_iter_i(VALUE i, VALUE *arg) +{ + if (!rb_yield(i)) rb_iter_break(); + rb_ary_push(arg[0], i); + return Qnil; +} + +/* + * call-seq: + * enum.take(n) => array + * enum.take {|arr| block } => array + * + * Without a block, returns first n elements from <i>enum</i> + * With a block, takes elements during block evaluation gives + * true. + * + * a = [1, 2, 3, 4, 5] + * + * a.take(3) # => [4, 5] + * a.take {|i| i < 3 } # => [3, 4, 5] + * + */ + +static VALUE +enum_take(int argc, VALUE *argv, VALUE obj) +{ + VALUE args[2]; + + if (!rb_block_given_p()) { + VALUE vlen; + + rb_scan_args(argc, argv, "1", &vlen); + args[1] = NUM2LONG(vlen); + args[0] = rb_ary_new2(args[1]); + rb_block_call(obj, id_each, 0, 0, take_i, (VALUE)args); + return args[0]; + } + rb_scan_args(argc, argv, "0"); + args[0] = rb_ary_new(); + rb_block_call(obj, id_each, 0, 0, take_iter_i, (VALUE)args); + return args[0]; +} + +static VALUE +drop_i(VALUE i, VALUE *arg) +{ + long n = arg[1]; + + if (arg[1] == 0) { + rb_ary_push(arg[0], i); + } + else { + arg[1]--; + } + return Qnil; +} + +static VALUE +drop_iter_i(VALUE i, VALUE *arg) +{ + if (!arg[1] && !rb_yield(i)) { + arg[1] = Qtrue; + } + if (arg[1]) { + rb_ary_push(arg[0], i); + } + return Qnil; +} + +/* + * call-seq: + * enum.drop(n) => array + * enum.drop {|arr| block } => array + * + * Without a block, drops first n elements from <i>enum</i>, and returns + * rest elements in an array. With a block, drops elements during block + * evaluation gives true. + * + * a = [1, 2, 3, 4, 5] + * + * a.drop(3) # => [4, 5] + * a.drop {|i| i < 3 } # => [3, 4, 5] + * + */ + +static VALUE +enum_drop(int argc, VALUE *argv, VALUE obj) +{ + VALUE args[2]; + + if (!rb_block_given_p()) { + VALUE vlen; + + rb_scan_args(argc, argv, "1", &vlen); + args[1] = NUM2LONG(vlen); + args[0] = rb_ary_new2(args[1]); + rb_block_call(obj, id_each, 0, 0, drop_i, (VALUE)args); + return args[0]; + } + rb_scan_args(argc, argv, "0"); + args[0] = rb_ary_new(); + args[1] = Qfalse; + rb_block_call(obj, id_each, 0, 0, drop_iter_i, (VALUE)args); + return args[0]; +} + /* * The <code>Enumerable</code> mixin provides collection classes with * several traversal and searching methods, and with the ability to @@ -1268,6 +1385,8 @@ Init_Enumerable(void) rb_define_method(rb_mEnumerable,"include?", enum_member, 1); rb_define_method(rb_mEnumerable,"each_with_index", enum_each_with_index, 0); rb_define_method(rb_mEnumerable, "zip", enum_zip, -1); + rb_define_method(rb_mEnumerable, "take", enum_take, -1); + rb_define_method(rb_mEnumerable, "drop", enum_drop, -1); id_eqq = rb_intern("==="); id_each = rb_intern("each"); |