Track the current collectable accumulator during for ..., into: ...

lib/elixir/lib/enum.ex

@@ -1561,27 +1561,40 @@ defmodule Enum do
   defp into_protocol(enumerable, collectable) do
     {initial, fun} = Collectable.into(collectable)
 
+    reduce_into_protocol(enumerable, initial, fn entry, acc -> fun.(acc, {:cont, entry}) end, fun)
+    |> fun.(:done)
+  end
+
+  defp reduce_into_protocol([entry | entries], acc, callback, fun) do
+    next_acc = collect_into(entry, acc, callback, fun)
+    reduce_into_protocol(entries, next_acc, callback, fun)
+  end
+
+  defp reduce_into_protocol([], acc, _callback, _fun) do
+    acc
+  end
+
+  defp reduce_into_protocol(enumerable, acc, callback, fun) do
+    step = fn entry, acc -> {:suspend, {entry, acc}} end
+    reduce_into_protocol(&Enumerable.reduce(enumerable, &1, step), acc, callback, fun, true)
+  end
+
+  defp reduce_into_protocol(reduce, acc, callback, fun, _suspended?)
+       when is_function(reduce, 1) do
     try do
-      reduce_into_protocol(enumerable, initial, fun)
+      reduce.({:cont, acc})
     catch
       kind, reason ->
-        fun.(initial, :halt)
+        fun.(acc, :halt)
         :erlang.raise(kind, reason, __STACKTRACE__)
     else
-      acc -> fun.(acc, :done)
-    end
-  end
-
-  defp reduce_into_protocol(enumerable, initial, fun) when is_list(enumerable) do
-    :lists.foldl(fn x, acc -> fun.(acc, {:cont, x}) end, initial, enumerable)
-  end
+      {:suspended, {entry, acc}, continuation} ->
+        next_acc = collect_into(entry, acc, callback, fun, continuation)
+        reduce_into_protocol(continuation, next_acc, callback, fun, true)
 
-  defp reduce_into_protocol(enumerable, initial, fun) do
-    enumerable
-    |> Enumerable.reduce({:cont, initial}, fn x, acc ->
-      {:cont, fun.(acc, {:cont, x})}
-    end)
-    |> elem(1)
+      {_, acc} ->
+        acc
+    end
   end
 
   @doc """
@@ -1632,27 +1645,42 @@ defmodule Enum do
   defp into_protocol(enumerable, collectable, transform) do
     {initial, fun} = Collectable.into(collectable)
 
+    reduce_into_protocol(
+      enumerable,
+      initial,
+      fn entry, acc -> fun.(acc, {:cont, transform.(entry)}) end,
+      fun
+    )
+    |> fun.(:done)
+  end
+
+  defp collect_into(entry, acc, callback, fun) do
     try do
-      reduce_into_protocol(enumerable, initial, transform, fun)
+      callback.(entry, acc)
     catch
       kind, reason ->
-        fun.(initial, :halt)
+        fun.(acc, :halt)
         :erlang.raise(kind, reason, __STACKTRACE__)
-    else
-      acc -> fun.(acc, :done)
     end
   end
 
-  defp reduce_into_protocol(enumerable, initial, transform, fun) when is_list(enumerable) do
-    :lists.foldl(fn x, acc -> fun.(acc, {:cont, transform.(x)}) end, initial, enumerable)
+  defp collect_into(entry, acc, callback, fun, continuation) do
+    try do
+      callback.(entry, acc)
+    catch
+      kind, reason ->
+        safe_halt_continuation(continuation, acc)
+        fun.(acc, :halt)
+        :erlang.raise(kind, reason, __STACKTRACE__)
+    end
   end
 
-  defp reduce_into_protocol(enumerable, initial, transform, fun) do
-    enumerable
-    |> Enumerable.reduce({:cont, initial}, fn x, acc ->
-      {:cont, fun.(acc, {:cont, transform.(x)})}
-    end)
-    |> elem(1)
+  defp safe_halt_continuation(continuation, acc) do
+    try do
+      continuation.({:halt, acc})
+    catch
+      _, _ -> :ok
+    end
   end
 
   @doc """

lib/elixir/lib/stream.ex

@@ -587,33 +587,56 @@ defmodule Stream do
 
   defp do_into(enum, collectable, transform, acc, fun) do
     {initial, into} = Collectable.into(collectable)
+    step = fn x, acc -> {:suspend, {x, acc}} end
+    do_into_reduce(&Enumerable.reduce(enum, &1, step), initial, into, transform, acc, fun)
+  end
 
-    composed = fn x, [acc | collectable] ->
-      collectable = into.(collectable, {:cont, transform.(x)})
-      {reason, acc} = fun.(x, acc)
-      {reason, [acc | collectable]}
-    end
-
-    do_into(&Enumerable.reduce(enum, &1, composed), initial, into, acc)
+  defp do_into_reduce(reduce, collectable, into, transform, {:suspend, acc}, fun) do
+    {:suspended, acc, &do_into_reduce(reduce, collectable, into, transform, &1, fun)}
   end
 
-  defp do_into(reduce, collectable, into, {command, acc}) do
+  defp do_into_reduce(reduce, collectable, into, transform, {command, acc}, fun) do
     try do
-      reduce.({command, [acc | collectable]})
+      reduce.({command, acc})
     catch
       kind, reason ->
         into.(collectable, :halt)
         :erlang.raise(kind, reason, __STACKTRACE__)
     else
-      {:suspended, [acc | collectable], continuation} ->
-        {:suspended, acc, &do_into(continuation, collectable, into, &1)}
+      {:suspended, {x, acc}, continuation} ->
+        {reason, acc, collectable} =
+          next_into_step(x, acc, collectable, continuation, into, transform, fun)
 
-      {reason, [acc | collectable]} ->
+        do_into_reduce(continuation, collectable, into, transform, {reason, acc}, fun)
+
+      {reason, acc} ->
         into.(collectable, :done)
         {reason, acc}
     end
   end
 
+  defp next_into_step(x, acc, collectable, continuation, into, transform, fun) do
+    try do
+      collectable = into.(collectable, {:cont, transform.(x)})
+      {reason, acc} = fun.(x, acc)
+      {reason, acc, collectable}
+    catch
+      kind, reason ->
+        safe_halt_into(continuation, acc, collectable, into)
+        :erlang.raise(kind, reason, __STACKTRACE__)
+    end
+  end
+
+  defp safe_halt_into(continuation, acc, collectable, into) do
+    try do
+      continuation.({:halt, acc})
+    catch
+      _, _ -> :ok
+    after
+      into.(collectable, :halt)
+    end
+  end
+
   @doc """
   Creates a stream that will apply the given function on
   enumeration.

lib/elixir/src/elixir_erl_for.erl

@@ -166,42 +166,16 @@ build_into(Ann, Clauses, Expr, ?empty_map_set_pattern = _Into, Uniq, S) ->
   {ReduceExpr, SR} = build_reduce(Ann, Clauses, InnerFun, Expr, {nil, Ann}, Uniq, S),
   {?remote(Ann, 'Elixir.MapSet', new, [ReduceExpr]), SR};
 build_into(Ann, Clauses, Expr, Into, Uniq, S) ->
-  {Fun, SF}    = build_var(Ann, S),
-  {Acc, SA}    = build_var(Ann, SF),
-  {Kind, SK}   = build_var(Ann, SA),
-  {Reason, SR} = build_var(Ann, SK),
-  {Stack, ST}  = build_var(Ann, SR),
-  {Done, SD}   = build_var(Ann, ST),
-
-  InnerFun = fun(InnerExpr, InnerAcc) ->
-    {call, Ann, Fun, [InnerAcc, pair(Ann, cont, InnerExpr)]}
-  end,
+  {Fun, SF} = build_var(Ann, S),
+  {Acc, SA} = build_var(Ann, SF),
 
   MatchExpr = {match, Ann,
     {tuple, Ann, [Acc, Fun]},
     ?remote(Ann, 'Elixir.Collectable', into, [Into])
   },
 
-  {IntoReduceExpr, SN} = build_reduce(Ann, Clauses, InnerFun, Expr, Acc, Uniq, SD),
-
-  TryExpr =
-    {'try', Ann,
-      [IntoReduceExpr],
-      [{clause, Ann,
-        [Done],
-        [],
-        [{call, Ann, Fun, [Done, {atom, Ann, done}]}]}],
-      [stacktrace_clause(Ann, Fun, Acc, Kind, Reason, Stack)],
-      []},
-
-  {{block, Ann, [MatchExpr, TryExpr]}, SN}.
-
-stacktrace_clause(Ann, Fun, Acc, Kind, Reason, Stack) ->
-  {clause, Ann,
-    [{tuple, Ann, [Kind, Reason, Stack]}],
-    [],
-    [{call, Ann, Fun, [Acc, {atom, Ann, halt}]},
-     ?remote(Ann, erlang, raise, [Kind, Reason, Stack])]}.
+  {IntoReduceExpr, SN} = build_into_reduce(Ann, Clauses, Expr, Acc, Fun, Uniq, SA),
+  {{block, Ann, [MatchExpr, {call, Ann, Fun, [IntoReduceExpr, {atom, Ann, done}]}]}, SN}.
 
 %% Helpers
 
@@ -232,6 +206,53 @@ build_reduce(Ann, Clauses, InnerFun, Expr, Into, true, S) ->
   EnumReduceCall = build_reduce_each(Clauses, InnerExpr, NewInto, Acc, SU),
   {?remote(Ann, erlang, element, [{integer, Ann, 1}, EnumReduceCall]), SU}.
 
+build_into_reduce(Ann, Clauses, Expr, Into, Fun, false, S) ->
+  {Acc, SA} = build_var(Ann, S),
+  {ProtectedExpr, SP} = build_into_try(Ann, Expr, Fun, Acc, SA),
+  {build_reduce_each(Clauses, {call, Ann, Fun, [Acc, pair(Ann, cont, ProtectedExpr)]}, Into, Acc, SP), SP};
+build_into_reduce(Ann, Clauses, Expr, Into, Fun, true, S) ->
+  %% Those variables are used only inside the anonymous function
+  %% so we don't need to worry about returning the scope.
+  {Acc, SA} = build_var(Ann, S),
+  {Value, SV} = build_var(Ann, SA),
+  {IntoAcc, SI} = build_var(Ann, SV),
+  {UniqAcc, SU} = build_var(Ann, SI),
+
+  {ProtectedExpr, SP} = build_into_try(Ann, Expr, Fun, IntoAcc, SU),
+  NewInto = {tuple, Ann, [Into, {map, Ann, []}]},
+  AccTuple = {tuple, Ann, [IntoAcc, UniqAcc]},
+  PutUniqExpr = {map, Ann, UniqAcc, [{map_field_assoc, Ann, Value, {atom, Ann, true}}]},
+
+  InnerExpr = {block, Ann, [
+    {match, Ann, AccTuple, Acc},
+    {match, Ann, Value, ProtectedExpr},
+    {'case', Ann, UniqAcc, [
+      {clause, Ann, [{map, Ann, [{map_field_exact, Ann, Value, {atom, Ann, true}}]}], [], [AccTuple]},
+      {clause, Ann, [{map, Ann, []}], [], [{tuple, Ann, [{call, Ann, Fun, [IntoAcc, pair(Ann, cont, Value)]}, PutUniqExpr]}]}
+    ]}
+  ]},
+
+  EnumReduceCall = build_reduce_each(Clauses, InnerExpr, NewInto, Acc, SP),
+  {?remote(Ann, erlang, element, [{integer, Ann, 1}, EnumReduceCall]), SP}.
+
+build_into_try(Ann, Expr, Fun, Acc, S) ->
+  {Value, SV} = build_var(Ann, S),
+  {Kind, SK} = build_var(Ann, SV),
+  {Reason, SR} = build_var(Ann, SK),
+  {Stack, SS} = build_var(Ann, SR),
+
+  {{'try', Ann,
+    [Expr],
+    [{clause, Ann, [Value], [], [Value]}],
+    [{clause, Ann,
+      [{tuple, Ann, [Kind, Reason, Stack]}],
+      [],
+      [
+        {call, Ann, Fun, [Acc, {atom, Ann, halt}]},
+        ?remote(Ann, erlang, raise, [Kind, Reason, Stack])
+      ]}],
+    []}, SS}.
+
 build_reduce_each([{enum, Meta, Left, Right, Filters} | T], Expr, Arg, Acc, S) ->
   Ann = ?ann(Meta),
   True = build_reduce_each(T, Expr, Acc, Acc, S),

lib/elixir/test/elixir/enum_test.exs

@@ -8,6 +8,28 @@ defmodule EnumTest do
   use ExUnit.Case, async: true
   doctest Enum
 
+  defmodule IntoCollectable do
+    defstruct [:pid]
+
+    defimpl Collectable do
+      def into(%EnumTest.IntoCollectable{pid: pid}) do
+        {[],
+         fn
+           acc, {:cont, x} ->
+             [x | acc]
+
+           acc, :done ->
+             send(pid, {:done, Enum.reverse(acc)})
+             :ok
+
+           acc, :halt ->
+             send(pid, {:halt, Enum.reverse(acc)})
+             :ok
+         end}
+      end
+    end
+  end
+
   defp assert_runs_enumeration_only_once(enum_fun) do
     enumerator =
       Stream.map([:element], fn element ->
@@ -491,6 +513,23 @@ defmodule EnumTest do
     end
   end
 
+  test "into/2 halts with the current accumulator on enumerable exceptions" do
+    parent = self()
+
+    enumerable =
+      Stream.concat([
+        [1],
+        Stream.map([2], fn _ -> raise "boom" end)
+      ])
+
+    assert_raise RuntimeError, "boom", fn ->
+      Enum.into(enumerable, %IntoCollectable{pid: parent})
+    end
+
+    assert_received {:halt, [1]}
+    refute_received {:done, _}
+  end
+
   test "into/3" do
     assert Enum.into([1, 2, 3], [], fn x -> x * 2 end) == [2, 4, 6]
     assert Enum.into([1, 2, 3], "numbers: ", &to_string/1) == "numbers: 123"
@@ -503,6 +542,20 @@ defmodule EnumTest do
     end
   end
 
+  test "into/3 halts with the current accumulator on transform exceptions" do
+    parent = self()
+
+    assert_raise RuntimeError, "boom", fn ->
+      Enum.into([1, 2], %IntoCollectable{pid: parent}, fn
+        2 -> raise "boom"
+        x -> x
+      end)
+    end
+
+    assert_received {:halt, [1]}
+    refute_received {:done, _}
+  end
+
   test "join/2" do
     assert Enum.join([], " = ") == ""
     assert Enum.join([1, 2, 3], " = ") == "1 = 2 = 3"