Add Enumerable#{flat_map,collect_concat}

2 files changed, 36 insertions, 0 deletions

diff --git a/mrbgems/mruby-enum-ext/mrblib/enum.rb b/mrbgems/mruby-enum-ext/mrblib/enum.rb
index c0f0ea3f2..6413198d4 100644
--- a/mrbgems/mruby-enum-ext/mrblib/enum.rb
+++ b/mrbgems/mruby-enum-ext/mrblib/enum.rb
@@ -237,4 +237,34 @@ module Enumerable
     end
     count
   end
+
+  ##
+  # call-seq:
+  #    enum.flat_map       { |obj| block } -> array
+  #    enum.collect_concat { |obj| block } -> array
+  #    enum.flat_map                       -> an_enumerator
+  #    enum.collect_concat                 -> an_enumerator
+  #
+  # Returns a new array with the concatenated results of running
+  # <em>block</em> once for every element in <i>enum</i>.
+  #
+  # If no block is given, an enumerator is returned instead.
+  #
+  #    [1, 2, 3, 4].flat_map { |e| [e, -e] } #=> [1, -1, 2, -2, 3, -3, 4, -4]
+  #    [[1, 2], [3, 4]].flat_map { |e| e + [100] } #=> [1, 2, 100, 3, 4, 100]
+  def flat_map(&block)
+    return to_enum :flat_map unless block_given?
+
+    ary = []
+    self.each do |e|
+      e2 = block.call(e)
+      if e2.respond_to? :each
+        e2.each { |e3| ary.push(e3) }
+      else
+        ary.push(e2)
+      end
+    end
+    ary
+  end
+  alias collect_concat flat_map
 end
diff --git a/mrbgems/mruby-enum-ext/test/enum.rb b/mrbgems/mruby-enum-ext/test/enum.rb
index 065ef7f5f..9523b497e 100644
--- a/mrbgems/mruby-enum-ext/test/enum.rb
+++ b/mrbgems/mruby-enum-ext/test/enum.rb
@@ -60,3 +60,9 @@ assert("Enumerable#count") do
   assert_equal 2, a.count(2)
   assert_equal 3, a.count{|x| x % 2 == 0}
 end
+
+assert("Enumerable#flat_map") do
+  assert_equal [1, 2, 3, 4], [1, 2, 3, 4].flat_map { |e| e }
+  assert_equal [1, -1, 2, -2, 3, -3, 4, -4], [1, 2, 3, 4].flat_map { |e| [e, -e] }
+  assert_equal [1, 2, 100, 3, 4, 100], [[1, 2], [3, 4]].flat_map { |e| e + [100] }
+end