Do not invoke after_fun in chunk_while if emitted element halts, closes #7112

Author: José Valim

lib/elixir/lib/stream.ex

@@ -231,16 +231,37 @@ defmodule Stream do
   def chunk_while(enum, acc, chunk_fun, after_fun) do
     lazy(
       enum,
-      acc,
-      fn f1 -> R.chunk_while(chunk_fun, f1) end,
-      &after_chunk_while(&1, &2, after_fun)
+      [acc | after_fun],
+      fn f1 -> chunk_while_fun(chunk_fun, f1) end,
+      &after_chunk_while/2
     )
   end
 
-  defp after_chunk_while(acc(h, acc, t), f1, after_fun) do
+  defp chunk_while_fun(callback, fun) do
+    fn entry, acc(head, [acc | after_fun], tail) ->
+      case callback.(entry, acc) do
+        {:cont, emit, acc} ->
+          # If we emit an item and then we have to halt,
+          # we need to disable the after_fun callback to
+          # avoid emitting even more items.
+          case next(fun, emit, [head | tail]) do
+            {:halt, [head | tail]} -> {:halt, acc(head, [acc | &{:cont, &1}], tail)}
+            {command, [head | tail]} -> {command, acc(head, [acc | after_fun], tail)}
+          end
+
+        {:cont, acc} ->
+          skip(acc(head, [acc | after_fun], tail))
+
+        {:halt, acc} ->
+          {:halt, acc(head, [acc | after_fun], tail)}
+      end
+    end
+  end
+
+  defp after_chunk_while(acc(h, [acc | after_fun], t), f1) do
     case after_fun.(acc) do
-      {:cont, emit, acc} -> next_with_acc(f1, emit, h, acc, t)
-      {:cont, acc} -> {:cont, acc(h, acc, t)}
+      {:cont, emit, acc} -> next_with_acc(f1, emit, h, [acc | after_fun], t)
+      {:cont, acc} -> {:cont, acc(h, [acc | after_fun], t)}
     end
   end
 

lib/elixir/lib/stream/reducers.ex

@@ -55,18 +55,6 @@ defmodule Stream.Reducers do
     chunk_by.(enumerable, nil, chunk_fun, after_fun)
   end
 
-  defmacro chunk_while(callback, fun \\ nil) do
-    quote do
-      fn entry, acc(head, acc, tail) ->
-        case unquote(callback).(entry, acc) do
-          {:cont, emit, acc} -> next_with_acc(unquote(fun), emit, head, acc, tail)
-          {:cont, acc} -> skip(acc(head, acc, tail))
-          {:halt, acc} -> {:halt, acc(head, acc, tail)}
-        end
-      end
-    end
-  end
-
   defmacro dedup(callback, fun \\ nil) do
     quote do
       fn entry, acc(head, prev, tail) = acc ->

lib/elixir/test/elixir/stream_test.exs

@@ -157,6 +157,27 @@ defmodule StreamTest do
              [[5, 7, 9]]
   end
 
+  test "chunk_while/4 with inner halt" do
+    chunk_fun = fn
+      i, [] ->
+        {:cont, [i]}
+
+      i, chunk ->
+        if rem(i, 2) == 0 do
+          {:cont, Enum.reverse(chunk), [i]}
+        else
+          {:cont, [i | chunk]}
+        end
+    end
+
+    after_fun = fn
+      [] -> {:cont, []}
+      chunk -> {:cont, Enum.reverse(chunk), []}
+    end
+
+    assert Stream.chunk_while([1, 2, 3, 4, 5], [], chunk_fun, after_fun) |> Enum.at(0) == [1]
+  end
+
   test "concat/1" do
     stream = Stream.concat([1..3, [], [4, 5, 6], [], 7..9])
     assert is_function(stream)