Add Macro.expand_literals/2 and do partial literal expansion in @

Closes #12174.

Author: josevalim

lib/elixir/lib/kernel.ex

@@ -3590,33 +3590,31 @@ defmodule Kernel do
   defp collect_traces(:after_verify, arg, _env), do: {arg, []}
   defp collect_traces(:on_definition, arg, _env), do: {arg, []}
 
-  defp collect_traces(_name, arg, env) do
-    case is_pid(env.lexical_tracker) and Macro.quoted_literal?(arg) do
-      true ->
-        env = %{env | function: {:__info__, 1}}
-
-        {arg, aliases} =
-          Macro.prewalk(arg, %{}, fn
-            {:__aliases__, _, _} = alias, acc ->
-              case Macro.expand(alias, env) do
-                atom when is_atom(atom) -> {atom, Map.put(acc, atom, [])}
-                _ -> {alias, acc}
-              end
-
-            node, acc ->
-              {node, acc}
-          end)
+  defp collect_traces(_name, arg, %{lexical_tracker: pid} = env) when is_pid(pid) do
+    env = %{env | function: {:__info__, 1}}
+
+    {arg, aliases} =
+      Macro.expand_literals(arg, %{}, fn
+        {:__aliases__, _, _} = alias, acc ->
+          case Macro.expand(alias, env) do
+            atom when is_atom(atom) -> {atom, Map.put(acc, atom, [])}
+            _ -> {alias, acc}
+          end
 
-        case map_size(aliases) do
-          0 -> {arg, []}
-          _ -> {arg, [{env.line, env.lexical_tracker, env.tracers, :maps.keys(aliases)}]}
-        end
+        node, acc ->
+          {Macro.expand(node, env), acc}
+      end)
 
-      false ->
-        {arg, []}
+    case map_size(aliases) do
+      0 -> {arg, []}
+      _ -> {arg, [{env.line, env.lexical_tracker, env.tracers, :maps.keys(aliases)}]}
     end
   end
 
+  defp collect_traces(_name, arg, _env) do
+    {arg, []}
+  end
+
   defp typespec?(:type), do: true
   defp typespec?(:typep), do: true
   defp typespec?(:opaque), do: true

lib/elixir/lib/macro.ex

@@ -1966,42 +1966,95 @@ defmodule Macro do
 
   def quoted_literal?({:%{}, _, args}), do: quoted_literal?(args)
   def quoted_literal?({:{}, _, args}), do: quoted_literal?(args)
+  def quoted_literal?({:__MODULE__, _, ctx}) when is_atom(ctx), do: true
   def quoted_literal?({left, right}), do: quoted_literal?(left) and quoted_literal?(right)
   def quoted_literal?(list) when is_list(list), do: :lists.all(&quoted_literal?/1, list)
   def quoted_literal?(term), do: is_atom(term) or is_number(term) or is_binary(term)
 
-  @doc """
-  Expands the `ast` representing a quoted literal within the given `env`.
+  @doc false
+  @deprecated "Use Macro.expand_literals/2 instead"
+  def expand_literal(ast, env) do
+    expand_literals(ast, env)
+  end
 
-  This function checks if the given AST represents a quoted literal
-  (using `quoted_literal?/1`) and then expands all relevant nodes.
-  If a literal node is not given, then it returns the AST as is.
-  At the moment, the only expandable literal nodes in an AST are
-  aliases, so this function only expands aliases.
+  @doc """
+  Expands all literals in `ast` with the given `env`.
 
   This function is mostly used to remove compile-time dependencies
   from AST nodes. In such cases, the given environment is usually
-  manipulate to represent a function:
+  manipulated to represent a function:
 
-      Macro.expand_literal(ast, %{env | function: {:my_code, 1}})
+      Macro.expand_literals(ast, %{env | function: {:my_code, 1}})
+
+  At the moment, the only expandable literal nodes in an AST are
+  aliases, so this function only expands aliases.
 
   However, be careful when removing compile-time dependencies between
   modules. If you remove them but you still invoke the module at
-  compile-time, you may end-up with broken behaviour.
+  compile-time, Elixir will be unable to properly recompile modules
+  when they change.
   """
-  @doc since: "1.14.0"
-  @spec expand_literal(t(), Macro.Env.t()) :: t()
-  def expand_literal(ast, env) do
-    if quoted_literal?(ast) do
-      prewalk(ast, fn
-        {:__aliases__, _, _} = alias -> expand(alias, env)
-        other -> other
-      end)
+  @doc since: "1.14.1"
+  @spec expand_literals(t(), Macro.Env.t()) :: t()
+  def expand_literals(ast, env) do
+    {ast, :ok} = expand_literals(ast, :ok, fn node, :ok -> {expand(node, env), :ok} end)
+    ast
+  end
+
+  @doc """
+  Expands all literals in `ast` with the given `acc` and `fun`.
+
+  `fun` will be invoked with an expandable AST node and `acc` and
+  must return a new node with `acc`. This is a general version of
+  `expand_literals/2` which supports a custom expansion function.
+  Please check `expand_literals/2` for use cases and pitfalls.
+  """
+  @doc since: "1.14.1"
+  @spec expand_literals(t(), acc, (t(), acc -> {t(), acc})) :: t() when acc: term()
+  def expand_literals(ast, acc, fun)
+
+  def expand_literals({:__aliases__, meta, args}, acc, fun) do
+    {args, acc} = expand_literals(args, acc, fun)
+
+    if :lists.all(&is_atom/1, args) do
+      fun.({:__aliases__, meta, args}, acc)
     else
-      ast
+      {{:__aliases__, meta, args}, acc}
     end
   end
 
+  def expand_literals({:__MODULE__, _meta, ctx} = node, acc, fun) when is_atom(ctx) do
+    fun.(node, acc)
+  end
+
+  def expand_literals({:%, meta, [left, right]}, acc, fun) do
+    {left, acc} = expand_literals(left, acc, fun)
+    {right, acc} = expand_literals(right, acc, fun)
+    {{:%, meta, [left, right]}, acc}
+  end
+
+  def expand_literals({:%{}, meta, args}, acc, fun) do
+    {args, acc} = expand_literals(args, acc, fun)
+    {{:%{}, meta, args}, acc}
+  end
+
+  def expand_literals({:{}, meta, args}, acc, fun) do
+    {args, acc} = expand_literals(args, acc, fun)
+    {{:{}, meta, args}, acc}
+  end
+
+  def expand_literals({left, right}, acc, fun) do
+    {left, acc} = expand_literals(left, acc, fun)
+    {right, acc} = expand_literals(right, acc, fun)
+    {{left, right}, acc}
+  end
+
+  def expand_literals(list, acc, fun) when is_list(list) do
+    :lists.mapfoldl(&expand_literals(&1, &2, fun), acc, list)
+  end
+
+  def expand_literals(term, acc, _fun), do: {term, acc}
+
   @doc """
   Receives an AST node and expands it until it can no longer
   be expanded.

lib/elixir/test/elixir/kernel/lexical_tracker_test.exs

@@ -162,7 +162,7 @@ defmodule Kernel.LexicalTrackerTest do
       {{compile, _, _, _}, _binding} =
         Code.eval_string("""
         defmodule Kernel.LexicalTrackerTest.Attribute1 do
-          @example [String, Enum]
+          @example [String, Enum, 3 + 10]
           def foo(atom) when atom in @example, do: atom
           Kernel.LexicalTracker.references(__ENV__.lexical_tracker)
         end |> elem(3)

lib/elixir/test/elixir/macro_test.exs

@@ -269,9 +269,26 @@ defmodule MacroTest do
     assert expand_and_clean(quote(do: oror(1, false)), __ENV__) == quoted
   end
 
-  test "expand_literal" do
-    assert Macro.expand_literal(quote(do: Foo), __ENV__) == Foo
-    assert Macro.expand_literal(quote(do: Foo + Bar), __ENV__) == quote(do: Foo + Bar)
+  test "expand_literals/2" do
+    assert Macro.expand_literals(quote(do: Foo), __ENV__) == Foo
+    assert Macro.expand_literals(quote(do: Foo + Bar), __ENV__) == quote(do: Foo + Bar)
+    assert Macro.expand_literals(quote(do: __MODULE__), __ENV__) == __MODULE__
+    assert Macro.expand_literals(quote(do: __MODULE__.Foo), __ENV__) == __MODULE__.Foo
+    assert Macro.expand_literals(quote(do: [Foo, 1 + 2]), __ENV__) == [Foo, quote(do: 1 + 2)]
+  end
+
+  test "expand_literals/3" do
+    fun = fn node, acc ->
+      expanded = Macro.expand(node, __ENV__)
+      {expanded, [expanded | acc]}
+    end
+
+    assert Macro.expand_literals(quote(do: Foo), [], fun) == {Foo, [Foo]}
+    assert Macro.expand_literals(quote(do: Foo + Bar), [], fun) == {quote(do: Foo + Bar), []}
+    assert Macro.expand_literals(quote(do: __MODULE__), [], fun) == {__MODULE__, [__MODULE__]}
+
+    assert Macro.expand_literals(quote(do: __MODULE__.Foo), [], fun) ==
+             {__MODULE__.Foo, [__MODULE__.Foo, __MODULE__]}
   end
 
   test "var/2" do