Auto infer size of matched variable in bitstrings (#13106)

Author: sabiwara

lib/elixir/lib/kernel.ex

@@ -2064,9 +2064,6 @@ defmodule Kernel do
       {var, _, nil} when is_atom(var) ->
         invalid_concat_left_argument_error(Atom.to_string(var))
 
-      {:^, _, [{var, _, nil}]} when is_atom(var) ->
-        invalid_concat_left_argument_error("^#{Atom.to_string(var)}")
-
       _ ->
         expanded_arg
     end

lib/elixir/src/elixir_bitstring.erl

@@ -35,7 +35,12 @@ expand(BitstrMeta, Fun, [{'::', Meta, [Left, Right]} | T], Acc, S, E, Alignment,
 
   MatchOrRequireSize = RequireSize or is_match_size(T, EL),
   EType = expr_type(ELeft),
-  {ERight, EAlignment, SS, ES} = expand_specs(EType, Meta, Right, SL, OriginalS, EL, MatchOrRequireSize),
+  ExpectSize = case ELeft of
+    {'^', _, [{_, _, _}]} -> {infer, ELeft};
+    _ when MatchOrRequireSize -> required;
+    _ -> optional
+  end,
+  {ERight, EAlignment, SS, ES} = expand_specs(EType, Meta, Right, SL, OriginalS, EL, ExpectSize),
 
   EAcc = concat_or_prepend_bitstring(Meta, ELeft, ERight, Acc, ES, MatchOrRequireSize),
   expand(BitstrMeta, Fun, T, EAcc, {SS, OriginalS}, ES, alignment(Alignment, EAlignment), RequireSize);
@@ -147,7 +152,7 @@ expand_expr(Meta, Component, Fun, S, E) ->
 
 %% Expands and normalizes types of a bitstring.
 
-expand_specs(ExprType, Meta, Info, S, OriginalS, E, RequireSize) ->
+expand_specs(ExprType, Meta, Info, S, OriginalS, E, ExpectSize) ->
   Default =
     #{size => default,
       unit => default,
@@ -158,11 +163,17 @@ expand_specs(ExprType, Meta, Info, S, OriginalS, E, RequireSize) ->
     expand_each_spec(Meta, unpack_specs(Info, []), Default, S, OriginalS, E),
 
   MergedType = type(Meta, ExprType, Type, E),
-  validate_size_required(Meta, RequireSize, ExprType, MergedType, Size, ES),
+  validate_size_required(Meta, ExpectSize, ExprType, MergedType, Size, ES),
   SizeAndUnit = size_and_unit(Meta, ExprType, Size, Unit, ES),
   Alignment = compute_alignment(MergedType, Size, Unit),
 
-  [H | T] = build_spec(Meta, Size, Unit, MergedType, Endianness, Sign, SizeAndUnit, ES),
+  MaybeInferredSize = case {ExpectSize, MergedType, SizeAndUnit} of
+    {{infer, PinnedVar}, binary, []} -> [{size, Meta, [{{'.', Meta, [erlang, byte_size]}, Meta, [PinnedVar]}]}];
+    {{infer, PinnedVar}, bitstring, []} -> [{size, Meta, [{{'.', Meta, [erlang, bit_size]}, Meta, [PinnedVar]}]}];
+    _ -> SizeAndUnit
+  end,
+
+  [H | T] = build_spec(Meta, Size, Unit, MergedType, Endianness, Sign, MaybeInferredSize, ES),
   {lists:foldl(fun(I, Acc) -> {'-', Meta, [Acc, I]} end, H, T), Alignment, SS, ES}.
 
 type(_, default, default, _) ->
@@ -276,7 +287,7 @@ validate_spec_arg(Meta, unit, Value, _S, _OriginalS, E) when not is_integer(Valu
 validate_spec_arg(_Meta, _Key, _Value, _S, _OriginalS, _E) ->
   ok.
 
-validate_size_required(Meta, true, default, Type, default, E) when Type == binary; Type == bitstring ->
+validate_size_required(Meta, required, default, Type, default, E) when Type == binary; Type == bitstring ->
   function_error(Meta, E, ?MODULE, unsized_binary),
   ok;
 validate_size_required(_, _, _, _, _, _) ->

lib/elixir/test/elixir/kernel/binary_test.exs

@@ -128,20 +128,6 @@ defmodule Kernel.BinaryTest do
     assert_raise ArgumentError, message, fn ->
       Code.eval_string(~s["a" <> b <> "c" = "abc"])
     end
-
-    assert_raise ArgumentError, message, fn ->
-      Code.eval_string(~s[
-        a = "a"
-        ^a <> "b" = "ab"
-      ])
-    end
-
-    assert_raise ArgumentError, message, fn ->
-      Code.eval_string(~s[
-        b = "b"
-        "a" <> ^b <> "c" = "abc"
-      ])
-    end
   end
 
   test "hex" do
@@ -269,6 +255,28 @@ defmodule Kernel.BinaryTest do
     assert <<1::size((^foo).bar)>> = <<1::5>>
   end
 
+  test "automatic size computation of matched bitsyntax variable" do
+    var = "foo"
+    <<^var::binary, rest::binary>> = "foobar"
+    assert rest == "bar"
+
+    <<^var::bytes, rest::bytes>> = "foobar"
+    assert rest == "bar"
+
+    ^var <> rest = "foobar"
+    assert rest == "bar"
+
+    var = <<0, 1>>
+    <<^var::bitstring, rest::bitstring>> = <<0, 1, 2, 3>>
+    assert rest == <<2, 3>>
+
+    <<^var::bits, rest::bits>> = <<0, 1, 2, 3>>
+    assert rest == <<2, 3>>
+
+    ^var <> rest = <<0, 1, 2, 3>>
+    assert rest == <<2, 3>>
+  end
+
   defmacro signed_16 do
     quote do
       big - signed - integer - unit(16)