From a21d0f72c280a3effe15554279ae006918cd97ce Mon Sep 17 00:00:00 2001
From: shugo <shugo@b2dd03c8-39d4-4d8f-98ff-823fe69b080e>
Date: Fri, 9 Mar 2012 05:34:41 +0000
Subject: [PATCH] * enumerator.c (lazy_flat_map): add
 Enumerable::Lazy#flat_map.

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

diff --git a/ChangeLog b/ChangeLog
index 45ac98fee8..0949d17e32 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,3 +1,7 @@
+Fri Mar  9 14:29:32 2012  Shugo Maeda  <shugo@ruby-lang.org>
+
+	* enumerator.c (lazy_flat_map): add Enumerable::Lazy#flat_map.
+
 Fri Mar  9 06:29:22 2012  Aaron Patterson <aaron@tenderlovemaking.com>
 
 	* ext/psych/lib/psych.rb (load, parse): stop parsing or loading after
diff --git a/enumerator.c b/enumerator.c
index c1a553a705..9e3da15dd3 100644
--- a/enumerator.c
+++ b/enumerator.c
@@ -1272,6 +1272,32 @@ lazy_map(VALUE obj)
     return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_map_func, 0);
 }
 
+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 {
+	int i;
+	for (i = 0; i < RARRAY_LEN(ary); i++) {
+	    rb_funcall(argv[0], id_yield, 1, RARRAY_PTR(ary)[i]);
+	}
+	return Qnil;
+    }
+}
+
+static VALUE
+lazy_flat_map(VALUE obj)
+{
+    if (!rb_block_given_p()) {
+	rb_raise(rb_eArgError, "tried to call lazy flat_map without a block");
+    }
+
+    return rb_block_call(rb_cLazy, id_new, 1, &obj, lazy_flat_map_func, 0);
+}
 
 static VALUE
 lazy_select_func(VALUE val, VALUE m, int argc, VALUE *argv)
@@ -1377,11 +1403,13 @@ Init_Enumerator(void)
     rb_define_method(rb_mEnumerable, "lazy", enumerable_lazy, 0);
     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);
     rb_define_method(rb_cLazy, "reject", lazy_reject, 0);
     rb_define_method(rb_cLazy, "grep", lazy_grep, 1);
 
     rb_define_alias(rb_cLazy, "collect", "map");
+    rb_define_alias(rb_cLazy, "collect_concat", "flat_map");
     rb_define_alias(rb_cLazy, "find_all", "select");
 
     rb_eStopIteration = rb_define_class("StopIteration", rb_eIndexError);
diff --git a/test/ruby/test_lazy_enumerator.rb b/test/ruby/test_lazy_enumerator.rb
index eeede6600c..9b74429a7f 100644
--- a/test/ruby/test_lazy_enumerator.rb
+++ b/test/ruby/test_lazy_enumerator.rb
@@ -80,6 +80,14 @@ class TestLazyEnumerator < Test::Unit::TestCase
     assert_equal(1, a.current)
   end
 
+  def test_flat_map
+    a = Step.new(1..3)
+    assert_equal(2, a.flat_map {|x| [x * 2]}.first)
+    assert_equal(3, a.current)
+    assert_equal(2, a.lazy.flat_map {|x| [x * 2]}.first)
+    assert_equal(1, a.current)
+  end
+
   def test_reject
     a = Step.new(1..6)
     assert_equal(4, a.reject {|x| x < 4}.first)