shugo	2012-03-09 14:34:41 +0900 (Fri, 09 Mar 2012)

  New Revision: 34956

  http://svn.ruby-lang.org/cgi-bin/viewvc.cgi?view=rev&revision=34956

  Log:
    * enumerator.c (lazy_flat_map): add Enumerable::Lazy#flat_map.

  Modified files:
    trunk/ChangeLog
    trunk/enumerator.c
    trunk/test/ruby/test_lazy_enumerator.rb

Index: ChangeLog
===================================================================
--- ChangeLog	(revision 34955)
+++ ChangeLog	(revision 34956)
@@ -1,3 +1,7 @@
+Fri Mar  9 14:29:32 2012  Shugo Maeda  <shugo@r...>
+
+	* enumerator.c (lazy_flat_map): add Enumerable::Lazy#flat_map.
+
 Fri Mar  9 06:29:22 2012  Aaron Patterson <aaron@t...>
 
 	* ext/psych/lib/psych.rb (load, parse): stop parsing or loading after
Index: enumerator.c
===================================================================
--- enumerator.c	(revision 34955)
+++ enumerator.c	(revision 34956)
@@ -1272,8 +1272,34 @@
     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)
 {
     VALUE element = argv[1];
@@ -1377,11 +1403,13 @@
     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);
Index: test/ruby/test_lazy_enumerator.rb
===================================================================
--- test/ruby/test_lazy_enumerator.rb	(revision 34955)
+++ test/ruby/test_lazy_enumerator.rb	(revision 34956)
@@ -80,6 +80,14 @@
     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)

--
ML: ruby-changes@q...
Info: http://www.atdot.net/~ko1/quickml/