Simplify and optimize flat_map [#84]

marcandre · Nov 12, 2014 · 62e3acd · 62e3acd
1 parent 26add80
commit 62e3acd
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 4 deletions.
diff --git a/lib/backports/1.9.2/enumerable/flat_map.rb b/lib/backports/1.9.2/enumerable/flat_map.rb
@@ -1,10 +1,13 @@
 unless Enumerable.method_defined? :flat_map
-  require 'backports/1.8.7/array/flatten'
-
   module Enumerable
-    def flat_map(&block)
+    def flat_map
       return to_enum(:flat_map) unless block_given?
-      map(&block).flatten(1)
+      r = []
+      each do |*args|
+        result = yield(*args)
+        result.respond_to?(:to_ary) ? r.concat(result) : r.push(result)
+      end
+      r
     end
     alias_method :collect_concat, :flat_map
   end

diff --git a/spec/tags/1.8.7/core/enumerable/flat_map_spec.rb b/spec/tags/1.8.7/core/enumerable/flat_map_spec.rb
@@ -0,0 +1 @@
+fails:Enumerable#flat_map calls to_ary but not to_a