From 6b4525912b67b085b1c6d958fdca8aef87a6c28d Mon Sep 17 00:00:00 2001
From: shugo <shugo@b2dd03c8-39d4-4d8f-98ff-823fe69b080e>
Date: Tue, 13 Mar 2012 23:16:37 +0000
Subject: [PATCH] * enumerator.c (lazy_zip_func): use each for non-Array
 objects.

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@35004 b2dd03c8-39d4-4d8f-98ff-823fe69b080e
---
 ChangeLog                         |  4 ++++
 enumerator.c                      | 21 +++++++++++++--------
 test/ruby/test_lazy_enumerator.rb | 10 ++++++++++
 3 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/ChangeLog b/ChangeLog
index a7544fa261..8666549e49 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,3 +1,7 @@
+Wed Mar 14 08:15:54 2012  Shugo Maeda  <shugo@ruby-lang.org>
+
+	* enumerator.c (lazy_zip_func): use each for non-Array objects.
+
 Wed Mar 14 08:06:35 2012  Shugo Maeda  <shugo@ruby-lang.org>
 
 	* enumerator.c (lazy_zip): add Enumerable::Lazy#flat_map.
diff --git a/enumerator.c b/enumerator.c
index 65e750e7ba..bc7f89522d 100644
--- a/enumerator.c
+++ b/enumerator.c
@@ -1293,21 +1293,26 @@ lazy_map(VALUE obj)
     return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_map_func, 0);
 }
 
+static VALUE
+lazy_flat_map_i(VALUE i, VALUE yielder, int argc, VALUE *argv)
+{
+    return rb_funcall2(yielder, id_yield, argc, argv);
+}
+
 static VALUE
 lazy_flat_map_func(VALUE val, VALUE m, int argc, VALUE *argv)
 {
     VALUE result = rb_yield_values2(argc - 1, &argv[1]);
-    VALUE ary = rb_check_array_type(result);
-    if (NIL_P(ary)) {
-	return rb_funcall(argv[0], id_yield, 1, result);
+    if (TYPE(result) == T_ARRAY) {
+	int i;
+	for (i = 0; i < RARRAY_LEN(result); i++) {
+	    rb_funcall(argv[0], id_yield, 1, RARRAY_PTR(result)[i]);
+	}
     }
     else {
-	int i;
-	for (i = 0; i < RARRAY_LEN(ary); i++) {
-	    rb_funcall(argv[0], id_yield, 1, RARRAY_PTR(ary)[i]);
-	}
-	return Qnil;
+	rb_block_call(result, id_each, 0, 0, lazy_flat_map_i, argv[0]);
     }
+    return Qnil;
 }
 
 static VALUE
diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb
index 6d96438454..9dbabafaa9 100644
--- a/test/ruby/test_lazy_enumerator.rb
+++ b/test/ruby/test_lazy_enumerator.rb
@@ -88,6 +88,16 @@ class TestLazyEnumerator < Test::Unit::TestCase
     assert_equal(1, a.current)
   end
 
+  def test_flat_map_nested
+    a = Step.new(1..3)
+    assert_equal([1, "a"],
+                 a.flat_map {|x| ("a".."c").map {|y| [x, y]}}.first)
+    assert_equal(3, a.current)
+    assert_equal([1, "a"],
+                 a.lazy.flat_map {|x| ("a".."c").lazy.map {|y| [x, y]}}.first)
+    assert_equal(1, a.current)
+  end
+
   def test_reject
     a = Step.new(1..6)
     assert_equal(4, a.reject {|x| x < 4}.first)