Add Enumerable#{flat_map,collect_concat}

mrbgems/mruby-enum-ext/mrblib/enum.rb

@@ -237,4 +237,34 @@ def count(v=NONE, &block)
     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

mrbgems/mruby-enum-ext/test/enum.rb

@@ -60,3 +60,9 @@
   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