Only consider outer halts in flat_map when the op is not halt itself

Signed-off-by: José Valim <jose.valim@plataformatec.com.br>

Author: José Valim

lib/elixir/lib/stream.ex

@@ -699,9 +699,11 @@ defmodule Stream do
         next.({:halt, next_acc})
         :erlang.raise(kind, reason, stacktrace)
     else
-      {:halted, [:outer|acc]} ->
+      # Only take into account outer halts when the op is not halt itself.
+      # Otherwise, we were the ones wishing to halt, so we should just stop.
+      {:halted, [:outer|acc]} when op != :halt ->
         do_transform(user_acc, user, fun, next_acc, next, {:cont, acc}, inner)
-      {:halted, [:inner|acc]} ->
+      {:halted, [_|acc]} ->
         next.({:halt, next_acc})
         {:halted, acc}
       {:done, [_|acc]} ->

lib/elixir/test/elixir/stream_test.exs

@@ -326,6 +326,18 @@ defmodule StreamTest do
     assert Enum.take(stream, 5) == [1, 2, 3, 4, 3]
   end
 
+  test "flat_map/2 properly halts both inner and outer stream when inner stream is halted" do
+    # Fixes a bug that, when the inner stream was done,
+    # sending it a halt would cause it to return the
+    # inner stream was halted, forcing flat_map to get
+    # the next value from the outer stream, evaluate it,
+    # get annother inner stream, just to halt it.
+    assert [1, 2] # 2 should never be used
+           |> Stream.flat_map(fn 1 -> Stream.repeatedly(fn -> 1 end) end)
+           |> Stream.flat_map(fn 1 -> Stream.repeatedly(fn -> 1 end) end)
+           |> Enum.take(1) == [1]
+  end
+
   test "interval/1" do
     stream = Stream.interval(10)
     now = :erlang.now