* enumerator.c (enumerator_each, generator_each): pass arguments to

  the block with yielder.


git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@34950 b2dd03c8-39d4-4d8f-98ff-823fe69b080e

4 files changed, 51 insertions, 9 deletions

diff --git a/ChangeLog b/ChangeLog
index ec1f6c5eb7..64a71a7376 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,3 +1,8 @@
+Fri Mar  9 00:25:59 2012  Nobuyoshi Nakada  <nobu@ruby-lang.org>
+
+	* enumerator.c (enumerator_each, generator_each): pass arguments to
+	  the block with yielder.
+
 Fri Mar  9 00:25:02 2012  Nobuyoshi Nakada  <nobu@ruby-lang.org>
 
 	* array.c (rb_ary_cat): new function to concat objects into array.
diff --git a/common.mk b/common.mk
index a350525a55..72791d9be1 100644
--- a/common.mk
+++ b/common.mk
@@ -621,7 +621,8 @@ encoding.$(OBJEXT): {$(VPATH)}encoding.c $(RUBY_H_INCLUDES) \
   {$(VPATH)}internal.h
 enum.$(OBJEXT): {$(VPATH)}enum.c $(RUBY_H_INCLUDES) {$(VPATH)}node.h \
   {$(VPATH)}util.h $(ID_H_INCLUDES)
-enumerator.$(OBJEXT): {$(VPATH)}enumerator.c $(RUBY_H_INCLUDES)
+enumerator.$(OBJEXT): {$(VPATH)}enumerator.c $(RUBY_H_INCLUDES) \
+  {$(VPATH)}internal.h
 error.$(OBJEXT): {$(VPATH)}error.c {$(VPATH)}known_errors.inc \
   $(RUBY_H_INCLUDES) $(VM_CORE_H_INCLUDES) $(ENCODING_H_INCLUDES) \
   {$(VPATH)}debug.h \
diff --git a/enumerator.c b/enumerator.c
index d61d152c8b..5de0ac5cb7 100644
--- a/enumerator.c
+++ b/enumerator.c
@@ -13,6 +13,7 @@
 ************************************************/
 
 #include "ruby/ruby.h"
+#include "internal.h"
 
 /*
  * Document-class: Enumerator
@@ -421,8 +422,20 @@ enumerator_block_call(VALUE obj, rb_block_call_func *func, VALUE arg)
  *
  */
 static VALUE
-enumerator_each(VALUE obj)
+enumerator_each(int argc, VALUE *argv, VALUE obj)
 {
+    if (argc > 0) {
+	struct enumerator *e = enumerator_ptr(obj = rb_obj_dup(obj));
+	VALUE args = e->args;
+	if (args) {
+	    args = rb_ary_dup(args);
+	    rb_ary_cat(args, argv, argc);
+	}
+	else {
+	    args = rb_ary_new4(argc, argv);
+	}
+	e->args = args;
+    }
     if (!rb_block_given_p()) return obj;
     return enumerator_block_call(obj, 0, obj);
 }
@@ -1090,7 +1103,8 @@ generator_initialize(int argc, VALUE *argv, VALUE obj)
 	rb_need_block();
 
 	proc = rb_block_proc();
-    } else {
+    }
+    else {
 	rb_scan_args(argc, argv, "1", &proc);
 
 	if (!rb_obj_is_proc(proc))
@@ -1127,14 +1141,17 @@ generator_init_copy(VALUE obj, VALUE orig)
 
 /* :nodoc: */
 static VALUE
-generator_each(VALUE obj)
+generator_each(int argc, VALUE *argv, VALUE obj)
 {
     struct generator *ptr = generator_ptr(obj);
-    VALUE yielder;
+    VALUE args = rb_ary_new2(argc + 1);
 
-    yielder = yielder_new();
+    rb_ary_push(args, yielder_new());
+    if (argc > 0) {
+	rb_ary_cat(args, argv, argc);
+    }
 
-    return rb_proc_call(ptr->proc, rb_ary_new3(1, yielder));
+    return rb_proc_call(ptr->proc, args);
 }
 
 /*
@@ -1201,7 +1218,7 @@ Init_Enumerator(void)
     rb_define_alloc_func(rb_cEnumerator, enumerator_allocate);
     rb_define_method(rb_cEnumerator, "initialize", enumerator_initialize, -1);
     rb_define_method(rb_cEnumerator, "initialize_copy", enumerator_init_copy, 1);
-    rb_define_method(rb_cEnumerator, "each", enumerator_each, 0);
+    rb_define_method(rb_cEnumerator, "each", enumerator_each, -1);
     rb_define_method(rb_cEnumerator, "each_with_index", enumerator_each_with_index, 0);
     rb_define_method(rb_cEnumerator, "each_with_object", enumerator_with_object, 1);
     rb_define_method(rb_cEnumerator, "with_index", enumerator_with_index, -1);
@@ -1223,7 +1240,7 @@ Init_Enumerator(void)
     rb_define_alloc_func(rb_cGenerator, generator_allocate);
     rb_define_method(rb_cGenerator, "initialize", generator_initialize, -1);
     rb_define_method(rb_cGenerator, "initialize_copy", generator_init_copy, 1);
-    rb_define_method(rb_cGenerator, "each", generator_each, 0);
+    rb_define_method(rb_cGenerator, "each", generator_each, -1);
 
     /* Yielder */
     rb_cYielder = rb_define_class_under(rb_cEnumerator, "Yielder", rb_cObject);
diff --git a/test/ruby/test_enumerator.rb b/test/ruby/test_enumerator.rb
index 10ba2f184e..a1ad22160c 100644
--- a/test/ruby/test_enumerator.rb
+++ b/test/ruby/test_enumerator.rb
@@ -238,6 +238,18 @@ class TestEnumerator < Test::Unit::TestCase
     assert_equal([1,2], e.next_values)
   end
 
+  def test_each_arg
+    o = Object.new
+    def o.each(ary)
+      ary << 1
+      yield
+    end
+    ary = []
+    e = o.to_enum.each(ary)
+    e.next
+    assert_equal([1], ary)
+  end
+
   def test_feed
     o = Object.new
     def o.each(ary)
@@ -358,6 +370,13 @@ class TestEnumerator < Test::Unit::TestCase
     assert_equal([1, 2, 3], a)
   end
 
+  def test_generator_args
+    g = Enumerator::Generator.new {|y, x| y << 1 << 2 << 3; x }
+    a = []
+    assert_equal(:bar, g.each(:bar) {|x| a << x })
+    assert_equal([1, 2, 3], a)
+  end
+
   def test_yielder
     # note: Enumerator::Yielder is a class just for internal
     a = []