Merge pull request #6765 from khasinski/fix-lazy-flat-map

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

@@ -291,11 +291,12 @@ class Enumerator
     #
     def flat_map(&block)
       Lazy.new(self){|yielder, val|
-        ary = block.call(val)
-        # TODO: check ary is an Array
-        ary.each {|x|
-          yielder << x
-        }
+        result = block.call(val)
+        if result.respond_to?(:each)
+          result.each {|x| yielder << x }
+        else
+          yielder << result
+        end
       }
     end
     alias collect_concat flat_map

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

@@ -46,6 +46,18 @@ assert("Enumerator::Lazy#to_enum") do
   assert_equal [0*1, 2*3, 4*5, 6*7], lazy_enum.map { |a| a.first * a.last }.first(4)
 end
 
+assert("Enumerator::Lazy#flat_map with array from block") do
+  assert_equal [1, 10, 2, 20, 3, 30], [1, 2, 3].lazy.flat_map {|x| [x, x*10]}.force
+end
+
+assert("Enumerator::Lazy#flat_map with non-enumerable from block") do
+  assert_equal [1, 2, 3], [1, 2, 3].lazy.flat_map {|x| x}.force
+end
+
+assert("Enumerator::Lazy#flat_map with nested array from block") do
+  assert_equal [[1, 2], [3, 4]], [1, 3].lazy.flat_map {|x| [[x, x+1]]}.force
+end
+
 assert("Enumerator::Lazy#grep_v") do
   lazy_grep_v = (0..).lazy.grep_v(2..4)
   assert_kind_of Enumerator::Lazy, lazy_grep_v