Ensure we support counters in bitstring size vars

Signed-off-by: José Valim <[email protected]>

Author: José Valim

lib/elixir/src/elixir_bitstring.erl

@@ -253,13 +253,17 @@ expand_spec_arg(Expr, E) -> elixir_expand:expand(Expr, E).
 
 validate_spec_arg(Meta, size, Value, E, OriginalE) ->
   case Value of
-    {Var, _, Context} when is_atom(Var) and is_atom(Context) ->
-      case is_valid_spec_arg_var({Var, Context}, E, OriginalE) of
+    {Var, VarMeta, Context} when is_atom(Var) and is_atom(Context) ->
+      Tuple = {Var, elixir_utils:var_context(VarMeta, Context)},
+
+      case is_valid_spec_arg_var(Tuple, E, OriginalE) of
         true -> ok;
         false -> form_error(Meta, ?key(E, file), ?MODULE, {undefined_var_in_spec, Value})
       end;
+
     _ when is_integer(Value) ->
       ok;
+
     _ ->
       form_error(Meta, ?key(E, file), ?MODULE, {bad_size_argument, Value})
   end;

lib/elixir/src/elixir_erl_pass.erl

@@ -178,15 +178,15 @@ translate({with, Meta, [_ | _] = Args}, S) ->
 %% Variables
 
 translate({'^', Meta, [{Name, VarMeta, Kind}]}, #elixir_erl{context=match} = S) when is_atom(Name), is_atom(Kind) ->
-  Tuple = {Name, var_context(VarMeta, Kind)},
+  Tuple = {Name, elixir_utils:var_context(VarMeta, Kind)},
   {ok, {_Counter, Value}} = maps:find(Tuple, S#elixir_erl.backup_vars),
 
   PAnn = ?ann(?generated(Meta)),
   PVar = {var, PAnn, Value},
 
   case S#elixir_erl.extra of
     pin_guard ->
-      {TVar, TS} = elixir_erl_var:translate(VarMeta, Name, var_context(VarMeta, Kind), S),
+      {TVar, TS} = elixir_erl_var:translate(VarMeta, Name, elixir_utils:var_context(VarMeta, Kind), S),
       Guard = {op, PAnn, '=:=', PVar, TVar},
       {TVar, TS#elixir_erl{extra_guards=[Guard | TS#elixir_erl.extra_guards]}};
     _ ->
@@ -197,7 +197,7 @@ translate({'_', Meta, Kind}, #elixir_erl{context=match} = S) when is_atom(Kind)
   {{var, ?ann(Meta), '_'}, S};
 
 translate({Name, Meta, Kind}, S) when is_atom(Name), is_atom(Kind) ->
-  elixir_erl_var:translate(Meta, Name, var_context(Meta, Kind), S);
+  elixir_erl_var:translate(Meta, Name, elixir_utils:var_context(Meta, Kind), S);
 
 %% Local calls
 
@@ -283,12 +283,6 @@ build_list([H | T], Acc) ->
 build_list([], Acc) ->
   Acc.
 
-var_context(Meta, Kind) ->
-  case lists:keyfind(counter, 1, Meta) of
-    {counter, Counter} -> Counter;
-    false -> Kind
-  end.
-
 %% Pack a list of expressions from a block.
 unblock({'block', _, Exprs}) -> Exprs;
 unblock(Expr)                -> [Expr].

lib/elixir/src/elixir_expand.erl

@@ -353,7 +353,7 @@ expand({'_', Meta, Kind}, E) when is_atom(Kind) ->
 
 expand({Name, Meta, Kind} = Var, #{context := match} = E) when is_atom(Name), is_atom(Kind) ->
   #{unused_vars := Unused, current_vars := Current, prematch_vars := Prematch} = E,
-  Pair = {Name, var_context(Meta, Kind)},
+  Pair = {Name, elixir_utils:var_context(Meta, Kind)},
   PrematchVersion = var_version(Prematch, Pair),
 
   EE =
@@ -381,7 +381,7 @@ expand({Name, Meta, Kind} = Var, #{context := match} = E) when is_atom(Name), is
   {Var, EE};
 expand({Name, Meta, Kind} = Var, E) when is_atom(Name), is_atom(Kind) ->
   #{unused_vars := Unused, current_vars := Current} = E,
-  Pair = {Name, var_context(Meta, Kind)},
+  Pair = {Name, elixir_utils:var_context(Meta, Kind)},
 
   case Current of
     #{Pair := {Version, _}} ->
@@ -618,12 +618,6 @@ var_version(Map, Pair) ->
     _ -> -1
   end.
 
-var_context(Meta, Kind) ->
-  case lists:keyfind(counter, 1, Meta) of
-    {counter, Counter} -> Counter;
-    false -> Kind
-  end.
-
 maybe_warn_underscored_var_repeat(Meta, Name, Kind, E) ->
   case should_warn(Meta) andalso atom_to_list(Name) of
     "_" ++ _ ->

lib/elixir/src/elixir_utils.erl

@@ -1,7 +1,7 @@
 %% Convenience functions used throughout elixir source code
 %% for ast manipulation and querying.
 -module(elixir_utils).
--export([get_line/1, split_last/1, noop/0,
+-export([get_line/1, split_last/1, noop/0, var_context/2,
   characters_to_list/1, characters_to_binary/1, relative_to_cwd/1,
   macro_name/1, returns_boolean/1, caller/4, meta_keep/1,
   read_file_type/1, read_file_type/2, read_link_type/1, read_posix_mtime_and_size/1,
@@ -37,6 +37,12 @@ erlang_comparison_op_to_elixir('=:=') -> '===';
 erlang_comparison_op_to_elixir('=/=') -> '!==';
 erlang_comparison_op_to_elixir(Other) -> Other.
 
+var_context(Meta, Kind) ->
+  case lists:keyfind(counter, 1, Meta) of
+    {counter, Counter} -> Counter;
+    false -> Kind
+  end.
+
 % Extract guards
 
 extract_guards({'when', _, [Left, Right]}) -> {Left, extract_or_guards(Right)};

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

@@ -2160,6 +2160,17 @@ defmodule Kernel.ExpansionTest do
       assert expand(before_expansion) |> clean_meta([:alignment]) == after_expansion
     end
 
+    defmacro offset(size, binary) do
+      quote do
+        offset = unquote(size)
+        <<_::size(offset)>> = unquote(binary)
+      end
+    end
+
+    test "supports size from counters" do
+      assert offset(8, <<0>>)
+    end
+
     test "merges bitstrings" do
       import Kernel, except: [-: 2]
 
@@ -2233,6 +2244,44 @@ defmodule Kernel.ExpansionTest do
       assert_raise CompileError, message, fn ->
         expand(quote(do: <<"foo"::size(:oops)>>))
       end
+
+      assert_raise CompileError, ~r/undefined variable "foo"/, fn ->
+        code =
+          quote do
+            fn <<_::size(foo)>> -> :ok end
+          end
+
+        expand(code)
+      end
+
+      assert_raise CompileError, ~r/undefined variable "foo"/, fn ->
+        code =
+          quote do
+            fn <<_::size(foo), foo::size(8)>> -> :ok end
+          end
+
+        expand(code)
+      end
+
+      assert_raise CompileError, ~r/undefined variable "foo" in bitstring segment/, fn ->
+        code =
+          quote do
+            fn foo, <<_::size(foo)>> -> :ok end
+          end
+
+        expand(code)
+      end
+
+      message = ~r"size in bitstring expects an integer or a variable as argument, got: foo()"
+
+      assert_raise CompileError, message, fn ->
+        code =
+          quote do
+            fn <<_::size(foo())>> -> :ok end
+          end
+
+        expand(code)
+      end
     end
 
     test "raises for invalid unit" do
@@ -2384,44 +2433,6 @@ defmodule Kernel.ExpansionTest do
 
       expand(code)
     end
-
-    assert_raise CompileError, ~r/undefined variable "foo"/, fn ->
-      code =
-        quote do
-          fn <<_::size(foo)>> -> :ok end
-        end
-
-      expand(code)
-    end
-
-    assert_raise CompileError, ~r/undefined variable "foo"/, fn ->
-      code =
-        quote do
-          fn <<_::size(foo), foo::size(8)>> -> :ok end
-        end
-
-      expand(code)
-    end
-
-    assert_raise CompileError, ~r/undefined variable "foo" in bitstring segment/, fn ->
-      code =
-        quote do
-          fn foo, <<_::size(foo)>> -> :ok end
-        end
-
-      expand(code)
-    end
-
-    message = ~r"size in bitstring expects an integer or a variable as argument, got: foo()"
-
-    assert_raise CompileError, message, fn ->
-      code =
-        quote do
-          fn <<_::size(foo())>> -> :ok end
-        end
-
-      expand(code)
-    end
   end
 
   ## Helpers