Add :reduce to for comprehensions (#8531)

Author: josevalim

lib/elixir/lib/kernel/special_forms.ex

@@ -1376,7 +1376,7 @@ defmodule Kernel.SpecialForms do
   filters or inside the block, are not reflected outside of the
   comprehension.
 
-  ## Into
+  ## The `:into` and `:uniq` options
 
   In the examples above, the result returned by the comprehension was
   always a list. The returned result can be configured by passing an
@@ -1396,18 +1396,57 @@ defmodule Kernel.SpecialForms do
         String.upcase(line)
       end
 
-  ## Uniq
-
-  `uniq: true` can also be given to comprehensions to guarantee that
-  that results are only added to the collection if they were not returned
+  Similarly, `uniq: true` can also be given to comprehensions to guarantee
+  the results are only added to the collection if they were not returned
   before. For example:
 
-      iex> for(x <- [1, 1, 2, 3], uniq: true, do: x * 2)
+      iex> for x <- [1, 1, 2, 3], uniq: true, do: x * 2
       [2, 4, 6]
 
-      iex> for(<<x <- "abcabc">>, uniq: true, into: "", do: <<x - 32>>)
+      iex> for <<x <- "abcabc">>, uniq: true, into: "", do: <<x - 32>>
       "ABC"
 
+  ## The `:reduce` option
+
+  While the `:into` option allows us to customize the comprehension behaviour
+  to a given data type, such as putting all of the values inside a map or inside
+  a binary, it is not always enough.
+
+  For example, imagine that you have a binary with letters where you want to
+  count how many times each lowercase letter happens, ignoring all uppercase
+  ones. For instance, for the string `"AbCabCABc"`, we want to return the map
+  `%{"a" => 1, "b" => 2, "c" => 1}`.
+
+  If we were to use `:into`, we would need a data type that computes the
+  frequency of each element it holds. While there is no such data type in
+  Elixir, you could implement one yourself.
+
+  A simpler option would be to use comprehensions for the mapping and
+  filtering of letters, and then we invoke `Enum.reduce/3` to build a map,
+  for example:
+
+      iex> letters = for <<x <- "AbCabCABc">>, x in ?a..?z, do: <<x>>
+      iex> Enum.reduce(letters, %{}, fn x, acc -> Map.update(acc, x, 1, & &1 + 1) end)
+      %{"a" => 1, "b" => 2, "c" => 1}
+
+  While the above is straight-forward, it has the downside of traversing the
+  data at least twice. If you are expecting long strings as inputs, this can
+  be quite expensive.
+
+  Luckily, comprehensions also support the `:reduce` option, which would allow
+  us to fuse both steps above into a single step:
+
+      iex> for <<x <- "AbCabCABc">>, x in ?a..?z, reduce: %{} do
+      ...>   acc -> Map.update(acc, <<x>>, 1, & &1 + 1)
+      ...> end
+      %{"a" => 1, "b" => 2, "c" => 1}
+
+  When the `:reduce` key is given, its value is used as the initial accumulator
+  and the `do` block must be changed to use `->` clauses, where the left side
+  of `->` receives the accumulated value of the previous iteration and the
+  expression on the right side must return the new accumulator value. Once there are no more
+  elements, the final accumulated value is returned. If there are no elements
+  at all, then the initial accumulator value is returned.
   """
   defmacro for(args), do: error!([args])
 

lib/elixir/src/elixir_erl_for.erl

@@ -3,9 +3,30 @@
 -include("elixir.hrl").
 
 translate(Meta, Args, Return, S) ->
-  Ann = ?ann(Meta),
   {Cases, [{do, Expr} | Opts]} = elixir_utils:split_last(Args),
 
+  case lists:keyfind(reduce, 1, Opts) of
+    {reduce, Reduce} -> translate_reduce(Meta, Cases, Expr, Reduce, S);
+    false -> translate_into(Meta, Cases, Expr, Opts, Return, S)
+  end.
+
+translate_reduce(Meta, Cases, Expr, Reduce, S) ->
+  Ann = ?ann(Meta),
+  {TReduce, SR} = elixir_erl_pass:translate(Reduce, S),
+  {TCases, SC} = translate_gen(Meta, Cases, [], SR),
+  CaseExpr = {'case', Meta, [ok, [{do, Expr}]]},
+  {TExpr, SE} = elixir_erl_pass:translate(CaseExpr, SC),
+
+  InnerFun = fun
+    ({'case', CaseAnn, _, CaseBlock}, InnerAcc) -> {'case', CaseAnn, InnerAcc, CaseBlock}
+  end,
+
+  SF = elixir_erl_var:mergec(SR, SE),
+  {build_reduce(Ann, TCases, InnerFun, TExpr, TReduce, false, SF), SF}.
+
+translate_into(Meta, Cases, Expr, Opts, Return, S) ->
+  Ann = ?ann(Meta),
+
   {TInto, SI} =
     case lists:keyfind(into, 1, Opts) of
       {into, Into} -> elixir_erl_pass:translate(Into, S);
@@ -16,7 +37,7 @@ translate(Meta, Args, Return, S) ->
   TUniq = lists:keyfind(uniq, 1, Opts) == {uniq, true},
 
   {TCases, SC} = translate_gen(Meta, Cases, [], SI),
-  {TExpr, SE}  = elixir_erl_pass:translate(wrap_expr(Expr, TInto), SC),
+  {TExpr, SE}  = elixir_erl_pass:translate(wrap_expr_if_unused(Expr, TInto), SC),
   SF = elixir_erl_var:mergec(SI, SE),
 
   case inline_or_into(TInto) of
@@ -26,8 +47,8 @@ translate(Meta, Args, Return, S) ->
 
 %% In case we have no return, we wrap the expression
 %% in a block that returns nil.
-wrap_expr(Expr, false) -> {'__block__', [], [Expr, nil]};
-wrap_expr(Expr, _)  -> Expr.
+wrap_expr_if_unused(Expr, false) -> {'__block__', [], [Expr, nil]};
+wrap_expr_if_unused(Expr, _)  -> Expr.
 
 translate_gen(ForMeta, [{'<-', Meta, [Left, Right]} | T], Acc, S) ->
   {TLeft, TRight, TFilters, TT, TS} = translate_gen(Meta, Left, Right, T, S),

lib/elixir/src/elixir_expand.erl

@@ -262,7 +262,7 @@ expand({for, Meta, [_ | _] = Args}, E) ->
         {Args, []}
     end,
 
-  validate_opts(Meta, for, [do, into, uniq], Block, E),
+  validate_opts(Meta, for, [do, into, uniq, reduce], Block, E),
 
   {Expr, Opts} =
     case lists:keytake(do, 1, Block) of
@@ -272,16 +272,16 @@ expand({for, Meta, [_ | _] = Args}, E) ->
         form_error(Meta, ?key(E, file), ?MODULE, {missing_option, for, [do]})
     end,
 
-  case lists:keyfind(uniq, 1, Opts) of
-    false -> ok;
-    {uniq, Value} when is_boolean(Value) -> ok;
-    {uniq, Value} -> form_error(Meta, ?key(E, file), ?MODULE, {for_invalid_uniq, Value})
-  end,
-
   {EOpts, EO} = expand(Opts, E),
   {ECases, EC} = lists:mapfoldl(fun expand_for/2, EO, Cases),
-  {EExpr, EE} = expand(Expr, EC),
   assert_generator_start(Meta, ECases, E),
+
+  {EExpr, EE} =
+    case validate_for_options(EOpts, false, false, false) of
+      {ok, MaybeReduce} -> expand_for_do_block(Meta, Expr, EC, MaybeReduce);
+      {error, Error} -> form_error(Meta, ?key(E, file), ?MODULE, Error)
+    end,
+
   {{for, Meta, ECases ++ [[{do, EExpr} | EOpts]]}, elixir_env:merge_and_check_unused_vars(E, EE)};
 
 %% With
@@ -697,6 +697,42 @@ generated_case_clauses([{do, Clauses}]) ->
   RClauses = [{'->', ?generated(Meta), Args} || {'->', Meta, Args} <- Clauses],
   [{do, RClauses}].
 
+%% Comprehensions
+
+validate_for_options([{into, _} = Pair | Opts], _Into, Uniq, Reduce) ->
+  validate_for_options(Opts, Pair, Uniq, Reduce);
+validate_for_options([{uniq, Boolean} = Pair | Opts], Into, _Uniq, Reduce) when is_boolean(Boolean) ->
+  validate_for_options(Opts, Into, Pair, Reduce);
+validate_for_options([{uniq, Value} | _], _, _, _) ->
+  {error, {for_invalid_uniq, Value}};
+validate_for_options([{reduce, _} = Pair | Opts], Into, Uniq, _Reduce) ->
+  validate_for_options(Opts, Into, Uniq, Pair);
+validate_for_options([], Into, Uniq, {reduce, _}) when Into /= false; Uniq /= false ->
+  {error, for_conflicting_reduce_into_uniq};
+validate_for_options([], _Into, _Uniq, Reduce) ->
+  {ok, Reduce}.
+
+expand_for_do_block(Meta, Expr, E, false) ->
+  case Expr of
+    [{'->', _, _} | _] -> form_error(Meta, ?key(E, file), ?MODULE, for_without_reduce_bad_block);
+    _ -> expand(Expr, E)
+  end;
+expand_for_do_block(Meta, Clauses, E, {reduce, _}) ->
+  Transformer = fun({_, _, [Left, _Right]} = Clause, Acc) ->
+    case Left of
+      [_] ->
+        {EClause, EAcc} = elixir_clauses:clause(Meta, fn, fun elixir_clauses:head/2, Clause, Acc),
+        {EClause, elixir_env:merge_and_check_unused_vars(Acc, EAcc)};
+      _ ->
+        form_error(Meta, ?key(E, file), ?MODULE, for_with_reduce_bad_block)
+    end
+  end,
+
+  case Clauses of
+    [{'->', _, _} | _] -> lists:mapfoldl(Transformer, E, Clauses);
+    _ -> form_error(Meta, ?key(E, file), ?MODULE, for_with_reduce_bad_block)
+  end.
+
 %% Locals
 
 assert_no_ambiguous_op(Name, Meta, [Arg], E) ->
@@ -1045,6 +1081,12 @@ format_error({invalid_args, Construct}) ->
   io_lib:format("invalid arguments for \"~ts\"", [Construct]);
 format_error({for_invalid_uniq, Value}) ->
   io_lib:format(":uniq option for comprehensions only accepts a boolean, got: ~ts", ['Elixir.Macro':to_string(Value)]);
+format_error(for_conflicting_reduce_into_uniq) ->
+  "cannot use :reduce alongside :into/:uniq in comprehension";
+format_error(for_with_reduce_bad_block) ->
+  "when using :reduce with comprehensions, the do block must be written using acc -> expr clauses, where each clause expects the accumulator as a single argument";
+format_error(for_without_reduce_bad_block) ->
+  "the do block was written using acc -> expr clauses but the :reduce option was not given";
 format_error(for_generator_start) ->
   "for comprehensions must start with a generator";
 format_error(unhandled_arrow_op) ->

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

@@ -226,6 +226,20 @@ defmodule Kernel.ComprehensionTest do
     assert IO.iodata_to_binary(Process.get(:into_cont)) == "roohkpmmfi"
   end
 
+  test "for comprehensions of map into map" do
+    enum = %{a: 2, b: 3}
+    assert for({k, v} <- enum, into: %{}, do: {k, v * v}) == %{a: 4, b: 9}
+  end
+
+  test "for comprehensions with reduce, generators and filters" do
+    acc =
+      for x <- 1..3, Integer.is_odd(x), <<y <- "hello">>, reduce: %{} do
+        acc -> Map.update(acc, x, [y], &[y | &1])
+      end
+
+    assert acc == %{1 => 'olleh', 3 => 'olleh'}
+  end
+
   ## List generators (inlined by the compiler)
 
   test "list for comprehensions" do
@@ -303,11 +317,6 @@ defmodule Kernel.ComprehensionTest do
             end) == <<7::size(1), 0::size(2), 1::size(1), 2::size(2), 3::size(1)>>
   end
 
-  test "map for comprehensions into map" do
-    enum = %{a: 2, b: 3}
-    assert for({k, v} <- enum, into: %{}, do: {k, v * v}) == %{a: 4, b: 9}
-  end
-
   test "list for comprehensions where value is not used" do
     enum = [1, 2, 3]
 
@@ -317,6 +326,15 @@ defmodule Kernel.ComprehensionTest do
            end) == "1\n2\n3\n"
   end
 
+  test "list for comprehensions with reduce, generators and filters" do
+    acc =
+      for x <- [1, 2, 3], Integer.is_odd(x), <<y <- "hello">>, reduce: %{} do
+        acc -> Map.update(acc, x, [y], &[y | &1])
+      end
+
+    assert acc == %{1 => 'olleh', 3 => 'olleh'}
+  end
+
   ## Binary generators (inlined by the compiler)
 
   test "binary for comprehensions" do
@@ -430,4 +448,15 @@ defmodule Kernel.ComprehensionTest do
              nil
            end) == "1\n2\n3\n"
   end
+
+  test "binary for comprehensions with reduce, generators and filters" do
+    bin = <<1, 2, 3>>
+
+    acc =
+      for <<x <- bin>>, Integer.is_odd(x), <<y <- "hello">>, reduce: %{} do
+        acc -> Map.update(acc, x, [y], &[y | &1])
+      end
+
+    assert acc == %{1 => 'olleh', 3 => 'olleh'}
+  end
 end

lib/elixir/test/elixir/kernel/expansion_test.exs

@@ -697,6 +697,26 @@ defmodule Kernel.ExpansionTest do
       end
     end
 
+    test "raise error on invalid reduce" do
+      assert_raise CompileError,
+                   ~r"cannot use :reduce alongside :into/:uniq in comprehension",
+                   fn ->
+                     expand(quote(do: for(x <- 1..3, reduce: %{}, into: %{}, do: (acc -> acc))))
+                   end
+
+      assert_raise CompileError,
+                   ~r"the do block was written using acc -> expr clauses but the :reduce option was not given",
+                   fn -> expand(quote(do: for(x <- 1..3, do: (acc -> acc)))) end
+
+      assert_raise CompileError,
+                   ~r"when using :reduce with comprehensions, the do block must be written using acc -> expr clauses",
+                   fn -> expand(quote(do: for(x <- 1..3, reduce: %{}, do: x))) end
+
+      assert_raise CompileError,
+                   ~r"when using :reduce with comprehensions, the do block must be written using acc -> expr clauses",
+                   fn -> expand(quote(do: for(x <- 1..3, reduce: %{}, do: (acc, x -> x)))) end
+    end
+
     test "raise error for unknown options" do
       assert_raise CompileError, ~r"unsupported option :else given to for", fn ->
         expand(quote(do: for(_ <- 1..2, do: 1, else: 1)))