* 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

1 files changed, 13 insertions, 8 deletions

diff --git a/enumerator.c b/enumerator.c
index 65e750e7ba..bc7f89522d 100644
--- a/enumerator.c
+++ b/enumerator.c
@@ -1294,20 +1294,25 @@ lazy_map(VALUE obj)
 }
 
 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);
-    }
-    else {
+    if (TYPE(result) == T_ARRAY) {
 	int i;
-	for (i = 0; i < RARRAY_LEN(ary); i++) {
-	    rb_funcall(argv[0], id_yield, 1, RARRAY_PTR(ary)[i]);
+	for (i = 0; i < RARRAY_LEN(result); i++) {
+	    rb_funcall(argv[0], id_yield, 1, RARRAY_PTR(result)[i]);
 	}
-	return Qnil;
     }
+    else {
+	rb_block_call(result, id_each, 0, 0, lazy_flat_map_i, argv[0]);
+    }
+    return Qnil;
 }
 
 static VALUE