Ensure suspending flat map does not consume items, closes #3751

Author: José Valim

lib/elixir/lib/stream.ex

@@ -687,6 +687,15 @@ defmodule Stream do
     do_transform(user_acc.(), user, fun, [], next, inner_acc, inner, after_fun)
   end
 
+  defp do_transform(user_acc, _user, _fun, _next_acc, _next, {:halt, inner_acc}, _inner, after_fun) do
+    do_after(after_fun, user_acc)
+    {:halted, inner_acc}
+  end
+
+  defp do_transform(user_acc, user, fun, next_acc, next, {:suspend, inner_acc}, inner, after_fun) do
+    {:suspended, inner_acc, &do_transform(user_acc, user, fun, next_acc, next, &1, inner, after_fun)}
+  end
+
   defp do_transform(user_acc, user, fun, next_acc, next, inner_acc, inner, after_fun) do
     case next.({:cont, next_acc}) do
       {:suspended, [val|next_acc], next} ->

lib/elixir/test/elixir/stream_test.exs

@@ -416,6 +416,22 @@ defmodule StreamTest do
     assert Stream.transform(nats, 0, &{[&1, &2], &1 + &2}) |> Enum.take(6) == [1, 0, 2, 1, 3, 3]
   end
 
+  test "transform/3 with early halt" do
+    stream = Stream.repeatedly(fn -> throw(:error) end)
+             |> Stream.transform(nil, &{[&1, &2], &1})
+
+    assert {:halted, nil} =
+      Enumerable.reduce(stream, {:halt, nil}, fn _, _ -> throw(:error) end)
+  end
+
+  test "transform/3 with early suspend" do
+    stream = Stream.repeatedly(fn -> throw(:error) end)
+             |> Stream.transform(nil, &{[&1, &2], &1})
+
+    assert {:suspended, nil, _} =
+      Enumerable.reduce(stream, {:suspend, nil}, fn _, _ -> throw(:error) end)
+  end
+
   test "transform/3 with halt" do
     stream = Stream.resource(fn -> 1 end,
                              fn acc -> {[acc], acc + 1} end,
@@ -570,6 +586,27 @@ defmodule StreamTest do
     assert Process.get(:stream_transform)
   end
 
+  test "transform/4 with early halt" do
+    stream = Stream.repeatedly(fn -> throw(:error) end)
+             |> Stream.transform(fn -> nil end, &{[&1, &2], &1},
+                                 fn nil -> Process.put(:stream_transform, true) end)
+
+    Process.put(:stream_transform, false)
+    assert {:halted, nil} =
+      Enumerable.reduce(stream, {:halt, nil}, fn _, _ -> throw(:error) end)
+    assert Process.get(:stream_transform)
+  end
+
+  test "transform/4 with early suspend" do
+    stream = Stream.repeatedly(fn -> throw(:error) end)
+             |> Stream.transform(fn -> nil end, &{[&1, &2], &1},
+                                 fn nil -> Process.put(:stream_transform, true) end)
+
+    refute Process.get(:stream_transform)
+    assert {:suspended, nil, _} =
+      Enumerable.reduce(stream, {:suspend, nil}, fn _, _ -> throw(:error) end)
+  end
+
   test "transform/4 closes on outer errors" do
     stream = Stream.transform(1..10, fn -> 0 end,
                               fn 3, _ -> throw(:error)