From ded27bf5dc25d61c1483ae2a79570b1435c8348b Mon Sep 17 00:00:00 2001
From: shugo <shugo@b2dd03c8-39d4-4d8f-98ff-823fe69b080e>
Date: Mon, 19 Mar 2012 08:22:29 +0000
Subject: * enumerator.c (lazy_flat_map_func): convert the block value to  
 Array if it doesn't respond to each.  [ruby-core:43334]   [Bug #6155]

git-svn-id: svn+ssh://ci.ruby-lang.org/ruby/trunk@35092 b2dd03c8-39d4-4d8f-98ff-823fe69b080e
---
 enumerator.c | 32 +++++++++++++++++++++++++++++++-
 1 file changed, 31 insertions(+), 1 deletion(-)

(limited to 'enumerator.c')

diff --git a/enumerator.c b/enumerator.c
index 9484f495cc..865730912f 100644
--- a/enumerator.c
+++ b/enumerator.c
@@ -1299,6 +1299,32 @@ 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_each(VALUE obj)
+{
+    NODE *memo = RNODE(obj);
+    rb_block_call(memo->u1.value, id_each, 0, 0, lazy_flat_map_i,
+		  memo->u2.value);
+    return Qnil;
+}
+
+static VALUE
+lazy_flat_map_to_ary(VALUE obj)
+{
+    NODE *memo = RNODE(obj);
+    VALUE ary = rb_check_array_type(memo->u1.value);
+    if (NIL_P(ary)) {
+	rb_funcall(memo->u2.value, id_yield, 1, memo->u1.value);
+    }
+    else {
+	long i;
+	for (i = 0; i < RARRAY_LEN(ary); i++) {
+	    rb_funcall(memo->u2.value, id_yield, 1, RARRAY_PTR(ary)[i]);
+	}
+    }
+    return Qnil;
+}
+
 static VALUE
 lazy_flat_map_func(VALUE val, VALUE m, int argc, VALUE *argv)
 {
@@ -1310,7 +1336,11 @@ lazy_flat_map_func(VALUE val, VALUE m, int argc, VALUE *argv)
 	}
     }
     else {
-	rb_block_call(result, id_each, 0, 0, lazy_flat_map_i, argv[0]);
+	NODE *memo;
+	memo = NEW_MEMO(result, argv[0], 0);
+	rb_rescue2(lazy_flat_map_each, (VALUE) memo,
+		   lazy_flat_map_to_ary, (VALUE) memo,
+		   rb_eNoMethodError, (VALUE)0);
     }
     return Qnil;
 }
-- 
cgit v1.2.3