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

Pass the current accumulator to final :halt invocation on exception

Fixes #15265

Author: lukaszsamson

lib/elixir/src/elixir_erl_for.erl

@@ -172,35 +172,48 @@ build_into(Ann, Clauses, Expr, Into, Uniq, S) ->
   {Reason, SR} = build_var(Ann, SK),
   {Stack, ST}  = build_var(Ann, SR),
   {Done, SD}   = build_var(Ann, ST),
+  {Ref, SRef}  = build_var(Ann, SD),
+  {Current, SC} = build_var(Ann, SRef),
 
   InnerFun = fun(InnerExpr, InnerAcc) ->
-    {call, Ann, Fun, [InnerAcc, pair(Ann, cont, InnerExpr)]}
+    {block, Ann, [
+      {match, Ann, Current, {call, Ann, Fun, [InnerAcc, pair(Ann, cont, InnerExpr)]}},
+      ?remote(Ann, erlang, put, [Ref, Current]),
+      Current
+    ]}
   end,
 
   MatchExpr = {match, Ann,
     {tuple, Ann, [Acc, Fun]},
     ?remote(Ann, 'Elixir.Collectable', into, [Into])
   },
 
-  {IntoReduceExpr, SN} = build_reduce(Ann, Clauses, InnerFun, Expr, Acc, Uniq, SD),
+  {IntoReduceExpr, SN} = build_reduce(Ann, Clauses, InnerFun, Expr, Acc, Uniq, SC),
+  RefExpr = {match, Ann, Ref, ?remote(Ann, erlang, make_ref, [])},
 
   TryExpr =
     {'try', Ann,
       [IntoReduceExpr],
       [{clause, Ann,
         [Done],
         [],
-        [{call, Ann, Fun, [Done, {atom, Ann, done}]}]}],
-      [stacktrace_clause(Ann, Fun, Acc, Kind, Reason, Stack)],
+        [
+          ?remote(Ann, erlang, erase, [Ref]),
+          {call, Ann, Fun, [Done, {atom, Ann, done}]}
+        ]}],
+      [stacktrace_clause(Ann, Fun, Ref, Kind, Reason, Stack)],
       []},
 
-  {{block, Ann, [MatchExpr, TryExpr]}, SN}.
+  {{block, Ann, [RefExpr, MatchExpr, ?remote(Ann, erlang, put, [Ref, Acc]), TryExpr]}, SN}.
 
-stacktrace_clause(Ann, Fun, Acc, Kind, Reason, Stack) ->
+stacktrace_clause(Ann, Fun, Ref, Kind, Reason, Stack) ->
+  CurrentAcc = ?remote(Ann, erlang, get, [Ref]),
   {clause, Ann,
     [{tuple, Ann, [Kind, Reason, Stack]}],
     [],
-    [{call, Ann, Fun, [Acc, {atom, Ann, halt}]},
+    [
+     {call, Ann, Fun, [CurrentAcc, {atom, Ann, halt}]},
+     ?remote(Ann, erlang, erase, [Ref]),
      ?remote(Ann, erlang, raise, [Kind, Reason, Stack])]}.
 
 %% Helpers

lib/elixir/test/elixir/kernel/comprehension_test.exs

@@ -26,6 +26,31 @@ defmodule Kernel.ComprehensionTest do
     end
   end
 
+  defmodule EvolvingCollectable do
+    defstruct []
+
+    defimpl Collectable do
+      def into(struct) do
+        fun = fn
+          acc, {:cont, x} ->
+            next = acc ++ [x]
+            Process.put(:into_trace, [{acc, x, next} | Process.get(:into_trace)])
+            next
+
+          acc, :done ->
+            Process.put(:into_done_acc, acc)
+            acc
+
+          acc, :halt ->
+            Process.put(:into_halt_acc, acc)
+            acc
+        end
+
+        {[struct], fun}
+      end
+    end
+  end
+
   defp to_bin(x) do
     <<x>>
   end
@@ -102,6 +127,27 @@ defmodule Kernel.ComprehensionTest do
     assert Process.get(:into_halt)
   end
 
+  test "for into halts with the current accumulator" do
+    Process.put(:into_trace, [])
+    Process.put(:into_halt_acc, nil)
+    Process.put(:into_done_acc, nil)
+
+    assert_raise RuntimeError, "oops", fn ->
+      for x <- [1, 2, 3], into: %EvolvingCollectable{} do
+        if x == 3, do: raise("oops")
+        x
+      end
+    end
+
+    assert Process.get(:into_trace) |> Enum.reverse() == [
+             {[%EvolvingCollectable{}], 1, [%EvolvingCollectable{}, 1]},
+             {[%EvolvingCollectable{}, 1], 2, [%EvolvingCollectable{}, 1, 2]}
+           ]
+
+    assert Process.get(:into_halt_acc) == [%EvolvingCollectable{}, 1, 2]
+    assert Process.get(:into_done_acc) == nil
+  end
+
   test "nested for comprehensions with unique values" do
     assert for(x <- [1, 1, 2], uniq: true, do: for(y <- [3, 3], uniq: true, do: x * y)) == [
              [3],