Inline comprehensions into non-empty bitstrings

lib/elixir/src/elixir_erl_for.erl

@@ -140,7 +140,7 @@ build_inline_each(Ann, Clauses, Expr, {nil, _} = Into, Uniq, InitVars, S) ->
   InnerFun = fun(InnerExpr, InnerAcc) -> {cons, Ann, InnerExpr, InnerAcc} end,
   {ReduceExpr, SR} = build_reduce(Ann, Clauses, InnerFun, Expr, Into, Uniq, InitVars, S),
   {?remote(Ann, lists, reverse, [ReduceExpr]), SR};
-build_inline_each(Ann, Clauses, Expr, {bin, _, []}, Uniq, InitVars, S) ->
+build_inline_each(Ann, Clauses, Expr, {bin, _, _} = Into, Uniq, InitVars, S) ->
   {InnerValue, SV} = build_var(Ann, S),
   Generated = erl_anno:set_generated(true, Ann),
 
@@ -158,7 +158,8 @@ build_inline_each(Ann, Clauses, Expr, {bin, _, []}, Uniq, InitVars, S) ->
     ]}
   end,
 
-  {ReduceExpr, SR} = build_reduce(Ann, Clauses, InnerFun, Expr, {nil, Ann}, Uniq, InitVars, SV),
+  InitAcc = {cons, Generated, Into, {nil, Generated}},
+  {ReduceExpr, SR} = build_reduce(Ann, Clauses, InnerFun, Expr, InitAcc, Uniq, InitVars, SV),
   {?remote(Ann, erlang, list_to_bitstring, [ReduceExpr]), SR}.
 
 build_into(Ann, Clauses, Expr, {map, _, []}, Uniq, InitVars, S) ->
@@ -335,6 +336,10 @@ valid_var_size({op, _Ann, _Op, Left, Right}, Vars) ->
   valid_var_size(Left, Vars) andalso valid_var_size(Right, Vars);
 valid_var_size(_Size, _vars) -> false.
 
+build_comprehension(Ann, Clauses, Expr, {bin, IntoAnn, [_ | _] = Elements}) ->
+  Generated = erl_anno:set_generated(true, Ann),
+  Comprehension = {bc, Ann, Expr, comprehension_clause(Clauses)},
+  {bin, IntoAnn, Elements ++ [{bin_element, Generated, Comprehension, default, [bitstring]}]};
 build_comprehension(Ann, Clauses, Expr, Into) ->
   {comprehension_kind(Into), Ann, Expr, comprehension_clause(Clauses)}.
 
@@ -350,7 +355,7 @@ comprehension_kind(false) -> lc;
 comprehension_kind({nil, _}) -> lc;
 comprehension_kind({bin, _, []}) -> bc.
 
-inline_or_into({bin, _, []}) -> inline;
+inline_or_into({bin, _, _}) -> inline;
 inline_or_into({nil, _}) -> inline;
 inline_or_into(false) -> inline;
 inline_or_into(_) -> into.

lib/elixir/test/elixir/fixtures/dialyzer/for_bitstring.ex

@@ -2,4 +2,8 @@ defmodule Dialyzer.ForBitstring do
   def foo() do
     for a <- 1..3, into: "", do: <<a>>
   end
+
+  def bar() do
+    for a <- 1..3, into: "prefix", do: <<a>>
+  end
 end

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

@@ -81,7 +81,9 @@ defmodule Kernel.ComprehensionTest do
     assert for(x <- list, uniq: true, into: [], do: x * 2) == [2, 4, 6]
     assert for(x <- list, uniq: true, into: %{}, do: {x, 1}) == %{1 => 1, 2 => 1, 3 => 1}
     assert for(x <- list, uniq: true, into: "", do: to_bin(x * 2)) == <<2, 4, 6>>
+    assert for(x <- list, uniq: true, into: <<0>>, do: to_bin(x * 2)) == <<0, 2, 4, 6>>
     assert for(<<x <- "abcabc">>, uniq: true, into: "", do: to_bin(x)) == "abc"
+    assert for(<<x <- "abcabc">>, uniq: true, into: "0", do: to_bin(x)) == "0abc"
 
     Process.put(:into_cont, [])
     Process.put(:into_done, false)
@@ -212,6 +214,10 @@ defmodule Kernel.ComprehensionTest do
     assert (for x <- enum, into: "" do
               if Integer.is_even(x), do: <<x::size(2)>>, else: <<x::size(1)>>
             end) == <<0::size(2), 1::size(1), 2::size(2), 3::size(1)>>
+
+    assert (for x <- enum, into: <<7::size(1)>> do
+              if Integer.is_even(x), do: <<x::size(2)>>, else: <<x::size(1)>>
+            end) == <<7::size(1), 0::size(2), 1::size(1), 2::size(2), 3::size(1)>>
   end
 
   test "for comprehensions into dynamic binary" do
@@ -363,6 +369,10 @@ defmodule Kernel.ComprehensionTest do
     assert (for x <- enum, into: "" do
               if Integer.is_even(x), do: <<x::size(2)>>, else: <<x::size(1)>>
             end) == <<0::size(2), 1::size(1), 2::size(2), 3::size(1)>>
+
+    assert (for x <- enum, into: <<7::size(1)>> do
+              if Integer.is_even(x), do: <<x::size(2)>>, else: <<x::size(1)>>
+            end) == <<7::size(1), 0::size(2), 1::size(1), 2::size(2), 3::size(1)>>
   end
 
   test "list for comprehensions into dynamic binary" do
@@ -437,6 +447,10 @@ defmodule Kernel.ComprehensionTest do
     assert (for <<x <- bin>>, into: "" do
               if Integer.is_even(x), do: <<x::size(2)>>, else: <<x::size(1)>>
             end) == <<0::size(2), 1::size(1), 2::size(2), 3::size(1)>>
+
+    assert (for <<x <- bin>>, into: <<7::size(1)>> do
+              if Integer.is_even(x), do: <<x::size(2)>>, else: <<x::size(1)>>
+            end) == <<7::size(1), 0::size(2), 1::size(1), 2::size(2), 3::size(1)>>
   end
 
   test "binary for comprehensions into dynamic binary" do