From d18dc746986e1611a11c46c0f728fd9f08bbdfc7 Mon Sep 17 00:00:00 2001 From: shugo Date: Thu, 15 Mar 2012 02:00:30 +0000 Subject: * enumerator.c (lazy_cycle): add Enumerable::Lazy#cycle. git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@35028 b2dd03c8-39d4-4d8f-98ff-823fe69b080e --- ChangeLog | 4 ++++ enumerator.c | 46 ++++++++++++++++++++++++++++++++++++--- test/ruby/test_lazy_enumerator.rb | 10 +++++++++ 3 files changed, 57 insertions(+), 3 deletions(-) diff --git a/ChangeLog b/ChangeLog index 8ce14b7164..76bbe927b1 100644 --- a/ChangeLog +++ b/ChangeLog @@ -1,3 +1,7 @@ +Thu Mar 15 10:57:27 2012 Shugo Maeda + + * enumerator.c (lazy_cycle): add Enumerable::Lazy#cycle. + Thu Mar 15 10:31:40 2012 Nobuyoshi Nakada * test/ruby/test_arity.rb (TestArity#err_mess): use assert_raise. diff --git a/enumerator.c b/enumerator.c index 28abc74987..23d7b4122b 100644 --- a/enumerator.c +++ b/enumerator.c @@ -1191,15 +1191,31 @@ lazy_init_block(VALUE val, VALUE m, int argc, VALUE *argv) } static VALUE -lazy_initialize(VALUE self, VALUE obj) +lazy_initialize(int argc, VALUE *argv, VALUE self) { + VALUE obj, meth; VALUE generator; + int offset; + if (argc < 1) { + rb_raise(rb_eArgError, "wrong number of arguments (%d for 1..)", argc); + } + else { + obj = argv[0]; + if (argc == 1) { + meth = sym_each; + offset = 1; + } + else { + meth = argv[1]; + offset = 2; + } + } generator = generator_allocate(rb_cGenerator); rb_block_call(generator, id_initialize, 0, 0, (rb_block_given_p() ? lazy_init_block_i: lazy_init_block), obj); - enumerator_init(self, generator, sym_each, 0, 0); + enumerator_init(self, generator, meth, argc - offset, argv + offset); return self; } @@ -1500,6 +1516,29 @@ lazy_drop_while(VALUE obj) (VALUE) memo); } +static VALUE +lazy_cycle_func(VALUE val, VALUE m, int argc, VALUE *argv) +{ + return rb_funcall2(argv[0], id_yield, argc - 1, argv + 1); +} + +static VALUE +lazy_cycle(int argc, VALUE *argv, VALUE obj) +{ + VALUE args; + int i; + + args = rb_ary_new2(argc + 1); + rb_ary_push(args, obj); + rb_ary_push(args, ID2SYM(rb_intern("cycle"))); + for (i = 0; i < argc; i++) { + rb_ary_push(args, argv[i]); + } + return rb_block_call(rb_cLazy, id_new, RARRAY_LEN(args), RARRAY_PTR(args), + rb_block_given_p() ? lazy_map_func : lazy_cycle_func, + 0); +} + static VALUE lazy_lazy(VALUE obj) { @@ -1587,7 +1626,7 @@ InitVM_Enumerator(void) /* Enumerable::Lazy */ rb_cLazy = rb_define_class_under(rb_mEnumerable, "Lazy", rb_cEnumerator); rb_define_method(rb_mEnumerable, "lazy", enumerable_lazy, 0); - rb_define_method(rb_cLazy, "initialize", lazy_initialize, 1); + rb_define_method(rb_cLazy, "initialize", lazy_initialize, -1); rb_define_method(rb_cLazy, "map", lazy_map, 0); rb_define_method(rb_cLazy, "flat_map", lazy_flat_map, 0); rb_define_method(rb_cLazy, "select", lazy_select, 0); @@ -1598,6 +1637,7 @@ InitVM_Enumerator(void) rb_define_method(rb_cLazy, "take_while", lazy_take_while, 0); rb_define_method(rb_cLazy, "drop", lazy_drop, 1); rb_define_method(rb_cLazy, "drop_while", lazy_drop_while, 0); + rb_define_method(rb_cLazy, "cycle", lazy_cycle, -1); rb_define_method(rb_cLazy, "lazy", lazy_lazy, 0); rb_define_alias(rb_cLazy, "collect", "map"); diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb index 6b4d49a041..1c0096f88c 100644 --- a/test/ruby/test_lazy_enumerator.rb +++ b/test/ruby/test_lazy_enumerator.rb @@ -184,6 +184,16 @@ class TestLazyEnumerator < Test::Unit::TestCase assert_equal([4, 5], (1..Float::INFINITY).lazy.drop(3).take(2).to_a) end + def test_cycle + a = Step.new(1..3) + assert_equal("1", a.cycle(2).map(&:to_s).first) + assert_equal(3, a.current) + assert_equal("1", a.lazy.cycle(2).map(&:to_s).first) + assert_equal(1, a.current) + assert_equal("1", a.lazy.cycle(2, &:to_s).first) + assert_equal(1, a.current) + end + def test_force assert_equal([1, 2, 3], (1..Float::INFINITY).lazy.take(3).force) end -- cgit v1.2.3