Ensure Stream.zip/2 and Stream.flat_map/2 follow up suspended calls

Author: José Valim

lib/elixir/lib/enum.ex

@@ -72,6 +72,9 @@ defprotocol Enumerable do
   `:suspended` tuple must be explicitly handled by the caller and
   never leak. In practice, this means regular enumeration functions
   just need to concern about `:done` and `:halted` results.
+
+  Furthermore, a `:suspend` call must always be followed by another call,
+  eventually halting or continuing until the end.
   """
   @type result :: { :done, term } | { :halted, term } | { :suspended, term, continuation }
 

lib/elixir/lib/stream.ex

@@ -450,10 +450,13 @@ defmodule Stream do
     try do
       reduce.(acc)
     catch
-      { :stream_flat_map, h } -> { :halted, h }
+      { :stream_flat_map, h } ->
+        reduce.({ :halt, elem(acc, 1) })
+        next.({ :halt, next_acc })
+        { :halted, h }
     else
-      { _, acc }              -> do_flat_map(next_acc, next, mapper, { :cont, acc }, fun)
-      { :suspended, acc, c }  -> { :suspended, acc, &do_flat_map(next_acc, next, mapper, &1, fun, c) }
+      { _, acc }             -> do_flat_map(next_acc, next, mapper, { :cont, acc }, fun)
+      { :suspended, acc, c } -> { :suspended, acc, &do_flat_map(next_acc, next, mapper, &1, fun, c) }
     end
   end
 
@@ -677,31 +680,43 @@ defmodule Stream do
     right_fun = &Enumerable.reduce(right, &1, step)
 
     # Return a function as a lazy enumerator.
-    &do_zip(left_fun, [], right_fun, [], &1, &2)
+    &do_zip([{ left_fun, [] }, { right_fun, [] }], &1, &2)
   end
 
-  defp do_zip(_left_fun, _left_acc, _right_fun, _right_acc, { :halt, acc }, _fun) do
+  defp do_zip(zips, { :halt, acc }, _fun) do
+    do_zip_close(zips)
     { :halted, acc }
   end
 
-  defp do_zip(left_fun, left_acc, right_fun, right_acc, { :suspend, acc }, fun) do
-    { :suspended, acc, &do_zip(left_fun, left_acc, right_fun, right_acc, &1, fun) }
+  defp do_zip(zips, { :suspend, acc }, fun) do
+    { :suspended, acc, &do_zip(zips, &1, fun) }
   end
 
-  defp do_zip(left_fun, left_acc, right_fun, right_acc, { :cont, acc }, callback) do
-    case left_fun.({ :cont, left_acc }) do
-      { :suspended, [x|left_acc], left_fun } ->
-        case right_fun.({ :cont, right_acc }) do
-          { :suspended, [y|right_acc], right_fun } ->
-            do_zip(left_fun, left_acc, right_fun, right_acc, callback.({ x, y }, acc), callback)
-          { reason, _ } ->
-            { reason, acc }
-        end
-      { reason, _ } ->
-        { reason, acc }
+  defp do_zip(zips, { :cont, acc }, callback) do
+    do_zip(zips, acc, callback, [], [])
+  end
+
+  defp do_zip([{ fun, fun_acc }|t], acc, callback, list, buffer) do
+    case fun.({ :cont, fun_acc }) do
+      { :suspended, [i|fun_acc], fun } ->
+        do_zip(t, acc, callback, [i|list], [{ fun, fun_acc }|buffer])
+      { _, _ } ->
+        do_zip_close(:lists.reverse(buffer) ++ t)
+        { :done, acc }
     end
   end
 
+  defp do_zip([], acc, callback, list, buffer) do
+    zipped = list_to_tuple(:lists.reverse(list))
+    do_zip(:lists.reverse(buffer), callback.(zipped, acc), callback)
+  end
+
+  defp do_zip_close([]), do: :ok
+  defp do_zip_close([{ fun, acc }|t]) do
+    fun.({ :halt, acc })
+    do_zip_close(t)
+  end
+
   defp do_zip_step(x, acc) do
     { :suspend, [x|acc] }
   end
@@ -829,7 +844,8 @@ defmodule Stream do
     { :suspended, acc, &do_resource(next_acc, next_fun, &1, fun, after_fun) }
   end
 
-  defp do_resource(_next_acc, _next_fun, { :halt, acc }, _fun, _after_fun) do
+  defp do_resource(next_acc, _next_fun, { :halt, acc }, _fun, after_fun) do
+    after_fun.(next_acc)
     { :halted, acc }
   end
 
@@ -841,8 +857,11 @@ defmodule Stream do
         after_fun.(next_acc)
         :erlang.raise(kind, reason, :erlang.get_stacktrace)
     else
-      nil             -> { :done, acc }
-      { v, next_acc } -> do_resource(next_acc, next_fun, fun.(v, acc), fun, after_fun)
+      nil ->
+        after_fun.(next_acc)
+        { :done, acc }
+      { v, next_acc } ->
+        do_resource(next_acc, next_fun, fun.(v, acc), fun, after_fun)
     end
   end
 

lib/elixir/test/elixir/stream_test.exs

@@ -242,6 +242,28 @@ defmodule StreamTest do
     assert Enum.zip(list, list) == Enum.zip(stream, stream)
   end
 
+  test "flat_map does not leave stream suspended" do
+    stream = Stream.flat_map [1,2,3],
+      fn i ->
+        Stream.resource(fn -> i end,
+                        fn acc -> { acc, acc + 1 } end,
+                        fn _ -> Process.put(:stream_flat_map, true) end)
+      end
+
+    Process.put(:stream_flat_map, false)
+    assert stream |> Enum.take(3) == [1,2,3]
+    assert Process.get(:stream_flat_map)
+
+    stream = Stream.resource(fn -> 1 end,
+                             fn acc -> { acc, acc + 1 } end,
+                             fn _ -> Process.put(:stream_flat_map, true) end)
+    stream = Stream.flat_map(stream, fn i -> [i, i + 1, i + 2] end)
+
+    Process.put(:stream_flat_map, false)
+    assert stream |> Enum.take(3) == [1,2,3]
+    assert Process.get(:stream_flat_map)
+  end
+
   test "iterate" do
     stream = Stream.iterate(0, &(&1+2))
     assert Enum.take(stream, 5) == [0,2,4,6,8]
@@ -402,6 +424,31 @@ defmodule StreamTest do
            [{1,:a},{2,:b},{3,:c},{4,:a},{5,:b},{6,:c}]
   end
 
+  test "zip/2 does not leave streams suspended" do
+    stream = Stream.resource(fn -> 1 end,
+                             fn acc -> { acc, acc + 1 } end,
+                             fn _ -> Process.put(:stream_zip, true) end)
+
+    Process.put(:stream_zip, false)
+    assert Stream.zip([:a, :b, :c], stream) |> Enum.to_list == [a: 1, b: 2, c: 3]
+    assert Process.get(:stream_zip)
+
+    Process.put(:stream_zip, false)
+    assert Stream.zip(stream, [:a, :b, :c]) |> Enum.to_list == [{ 1, :a }, { 2, :b }, { 3, :c }]
+    assert Process.get(:stream_zip)
+  end
+
+  test "zip/2 does not leave streams suspended on halt" do
+    stream = Stream.resource(fn -> 1 end,
+                             fn acc -> { acc, acc + 1 } end,
+                             fn _ -> Process.put(:stream_zip, :done) end)
+
+    assert Stream.zip([:a, :b, :c, :d, :e], stream) |> Enum.take(3) ==
+           [a: 1, b: 2, c: 3]
+
+    assert Process.get(:stream_zip) == :done
+  end
+
   test "with_index" do
     stream = Stream.with_index([1,2,3])
     assert is_lazy(stream)