Do not include module variables in prune_binding

Author: josevalim

lib/elixir/lib/code.ex

@@ -863,7 +863,10 @@ defmodule Code do
   ## Options
 
     * `:prune_binding` - (since v1.14.2) prune binding to keep only
-      variables read or written by the evaluated code
+      variables read or written by the evaluated code. Note that
+      variables used by modules are always pruned, even if later used
+      by the modules. You can submit to the `:on_module` tracer event
+      and access the variables used by the module from its environment.
 
   """
   @doc since: "1.14.0"

lib/elixir/lib/kernel.ex

@@ -4800,11 +4800,27 @@ defmodule Kernel do
           :elixir_quote.escape(block, :none, false)
       end
 
-    module_vars = :lists.map(&module_var/1, :maps.keys(env.versioned_vars))
+    versioned_vars = env.versioned_vars
+    prune = :erlang.is_map_key({:elixir, :prune_binding}, versioned_vars)
+
+    var_meta =
+      case prune do
+        true -> [generated: true, keep_unused: true]
+        false -> [generated: false]
+      end
+
+    module_vars = :lists.map(&module_var(&1, var_meta), :maps.keys(versioned_vars))
 
     quote do
       unquote(with_alias)
-      :elixir_module.compile(unquote(expanded), unquote(escaped), unquote(module_vars), __ENV__)
+
+      :elixir_module.compile(
+        unquote(expanded),
+        unquote(escaped),
+        unquote(module_vars),
+        unquote(prune),
+        __ENV__
+      )
     end
   end
 
@@ -4853,8 +4869,8 @@ defmodule Kernel do
     {module, module, nil}
   end
 
-  defp module_var({name, kind}) when is_atom(kind), do: {name, [generated: true], kind}
-  defp module_var({name, kind}), do: {name, [counter: kind, generated: true], nil}
+  defp module_var({name, kind}, meta) when is_atom(kind), do: {name, meta, kind}
+  defp module_var({name, kind}, meta), do: {name, [counter: kind] ++ meta, nil}
 
   @doc ~S"""
   Defines a public function with the given name and body.

lib/elixir/lib/module.ex

@@ -822,7 +822,7 @@ defmodule Module do
     next = :elixir_module.next_counter(nil)
     meta = Keyword.take(opts, [:line, :generated])
     quoted = :elixir_quote.linify_with_context_counter(meta, {module, next}, quoted)
-    :elixir_module.compile(module, quoted, [], :elixir.env_for_eval(opts))
+    :elixir_module.compile(module, quoted, [], false, :elixir.env_for_eval(opts))
   end
 
   @doc """

lib/elixir/src/elixir.erl

@@ -270,7 +270,7 @@ eval_forms(Tree, Binding, OrigE) ->
   eval_forms(Tree, Binding, OrigE, []).
 eval_forms(Tree, Binding, OrigE, Opts) ->
   Prune = proplists:get_value(prune_binding, Opts, false),
-  {ExVars, ErlVars, ErlBinding} = elixir_erl_var:load_binding(Binding),
+  {ExVars, ErlVars, ErlBinding} = elixir_erl_var:load_binding(Binding, Prune),
   E = elixir_env:with_vars(OrigE, ExVars),
   ExS = elixir_env:env_to_ex(E),
   ErlS = elixir_erl_var:from_env(E, ErlVars),
@@ -292,8 +292,8 @@ eval_forms(Tree, Binding, OrigE, Opts) ->
 
       ExternalHandler = eval_external_handler(NewE),
       {value, Value, NewBinding} = erl_eval:exprs(Exprs, ErlBinding, none, ExternalHandler),
+      PruneBefore = if Prune -> length(Binding); true -> -1 end,
 
-      PruneBefore = if Prune -> length(Binding); true -> 0 end,
       {DumpedBinding, DumpedVars} =
         elixir_erl_var:dump_binding(NewBinding, NewErlS, NewExS, PruneBefore),
 

lib/elixir/src/elixir_bootstrap.erl

@@ -20,7 +20,7 @@
 
 'MACRO-defmodule'(_Caller, Alias, [{do, Block}]) ->
   Escaped = elixir_quote:escape(Block, none, false),
-  Args = [Alias, Escaped, [], env()],
+  Args = [Alias, Escaped, [], false, env()],
   {{'.', [], [elixir_module, compile]}, [], Args}.
 
 '__info__'(functions) ->

lib/elixir/src/elixir_compiler.erl

@@ -48,7 +48,7 @@ compile(Quoted, ArgsList, E) ->
     elixir_erl_compiler:spawn(fun() -> spawned_compile(Expanded, E) end),
 
   Args = list_to_tuple(ArgsList),
-  {dispatch(Module, Fun, Args, Purgeable), EE}.
+  {dispatch(Module, Fun, Args, Purgeable), SE, EE}.
 
 spawned_compile(ExExprs, #{line := Line, file := File} = E) ->
   {Vars, S} = elixir_erl_var:from_env(E),
@@ -127,7 +127,7 @@ fast_compile({defmodule, Meta, [Mod, [{do, TailBlock}]]}, NoLineE) ->
   end,
 
   ContextModules = [Expanded | ?key(E, context_modules)],
-  elixir_module:compile(Expanded, Block, [], E#{context_modules := ContextModules}).
+  elixir_module:compile(Expanded, Block, [], false, E#{context_modules := ContextModules}).
 
 %% Bootstrapper
 

lib/elixir/src/elixir_erl_var.erl

@@ -2,7 +2,7 @@
 -module(elixir_erl_var).
 -export([
   translate/4, assign/2, build/2,
-  load_binding/1, dump_binding/4,
+  load_binding/2, dump_binding/4,
   from_env/1, from_env/2
 ]).
 -include("elixir.hrl").
@@ -65,16 +65,16 @@ to_erl_vars([], _Counter) ->
 to_erl_var(Counter) ->
   list_to_atom("_@" ++ integer_to_list(Counter)).
 
-load_binding(Binding) ->
-  load_binding(Binding, #{}, [], [], 0).
+load_binding(Binding, Prune) ->
+  load_binding(Binding, #{}, [], [], 0, Prune).
 
-load_binding([Binding | NextBindings], ExVars, ErlVars, Normalized, Counter) ->
+load_binding([Binding | NextBindings], ExVars, ErlVars, Normalized, Counter, Prune) ->
   {Pair, Value} = load_pair(Binding),
 
   case ExVars of
     #{Pair := VarCounter} ->
       ErlVar = to_erl_var(VarCounter),
-      load_binding(NextBindings, ExVars, ErlVars, [{ErlVar, Value} | Normalized], Counter);
+      load_binding(NextBindings, ExVars, ErlVars, [{ErlVar, Value} | Normalized], Counter, Prune);
 
     #{} ->
       ErlVar = to_erl_var(Counter),
@@ -84,10 +84,13 @@ load_binding([Binding | NextBindings], ExVars, ErlVars, Normalized, Counter) ->
         ExVars#{Pair => Counter},
         [{Counter, ErlVar} | ErlVars],
         [{ErlVar, Value} | Normalized],
-        Counter + 1
+        Counter + 1,
+        Prune
       )
   end;
-load_binding([], ExVars, ErlVars, Normalized, _Counter) ->
+load_binding([], ExVars, ErlVars, Normalized, Counter, true) ->
+  load_binding([{{elixir, prune_binding}, true}], ExVars, ErlVars, Normalized, Counter, false);
+load_binding([], ExVars, ErlVars, Normalized, _Counter, false) ->
   %% TODO: Remove me once we require Erlang/OTP 24+
   %% Also revisit dump_binding below and remove the vars field for simplicity.
   Mod =
@@ -109,7 +112,7 @@ dump_binding(Binding, ErlS, ExS, PruneBefore) ->
     %% If the variable is part of the pruning (usually the input binding)
     %% and is unused, we removed it from vars.
     (Pair, Version, {B, V})
-    when Version < PruneBefore, not is_map_key({Pair, Version}, Unused) ->
+    when Version =< PruneBefore, not is_map_key({Pair, Version}, Unused) ->
       {B, maps:remove(Pair, V)};
 
     ({Var, Kind} = Pair, Version, {B, V}) when is_atom(Kind) ->

lib/elixir/src/elixir_expand.erl

@@ -307,7 +307,7 @@ expand({Name, Meta, Kind}, S, #{context := match} = E) when is_atom(Name), is_at
     %% Variable was already overridden
     #{Pair := VarVersion} when VarVersion >= PrematchVersion ->
       maybe_warn_underscored_var_repeat(Meta, Name, Kind, E),
-      NewUnused = var_used(Pair, VarVersion, Unused),
+      NewUnused = var_used(Meta, Pair, VarVersion, Unused),
       Var = {Name, [{version, VarVersion} | Meta], Kind},
       {Var, S#elixir_ex{unused={NewUnused, Version}}, E};
 
@@ -354,7 +354,7 @@ expand({Name, Meta, Kind}, S, E) when is_atom(Name), is_atom(Kind) ->
     {ok, PairVersion} ->
       maybe_warn_underscored_var_access(Meta, Name, Kind, E),
       Var = {Name, [{version, PairVersion} | Meta], Kind},
-      {Var, S#elixir_ex{unused={var_used(Pair, PairVersion, Unused), Version}}, E};
+      {Var, S#elixir_ex{unused={var_used(Meta, Pair, PairVersion, Unused), Version}}, E};
 
     Error ->
       case lists:keyfind(if_undefined, 1, Meta) of
@@ -613,8 +613,11 @@ var_unused({_, Kind} = Pair, Meta, Version, Unused, Override) ->
     false -> Unused
   end.
 
-var_used({_, Kind} = Pair, Version, Unused) ->
+var_used(Meta, {_, Kind} = Pair, Version, Unused) ->
+  KeepUnused = lists:keymember(keep_unused, 1, Meta),
+
   if
+    KeepUnused -> Unused;
     is_atom(Kind) -> Unused#{{Pair, Version} => false};
     true -> Unused
   end.

lib/elixir/src/elixir_module.erl

@@ -1,6 +1,6 @@
 -module(elixir_module).
 -export([file/1, data_tables/1, is_open/1, mode/1, delete_definition_attributes/6,
-         compile/4, expand_callback/6, format_error/1, compiler_modules/0,
+         compile/5, expand_callback/6, format_error/1, compiler_modules/0,
          write_cache/3, read_cache/2, next_counter/1]).
 -include("elixir.hrl").
 -define(counter_attr, {elixir, counter}).
@@ -64,7 +64,7 @@ next_counter(Module) ->
 
 %% Compilation hook
 
-compile(Module, Block, Vars, Env) when is_atom(Module) ->
+compile(Module, Block, Vars, Prune, Env) when is_atom(Module) ->
   #{line := Line, function := Function, versioned_vars := OldVerVars} = Env,
 
   {VerVars, _} =
@@ -82,16 +82,16 @@ compile(Module, Block, Vars, Env) when is_atom(Module) ->
     #{lexical_tracker := nil} ->
       elixir_lexical:run(
         MaybeLexEnv,
-        fun(LexEnv) -> compile(Line, Module, Block, Vars, LexEnv) end,
+        fun(LexEnv) -> compile(Line, Module, Block, Vars, Prune, LexEnv) end,
         fun(_LexEnv) -> ok end
       );
     _ ->
-      compile(Line, Module, Block, Vars, MaybeLexEnv)
+      compile(Line, Module, Block, Vars, Prune, MaybeLexEnv)
   end;
-compile(Module, _Block, _Vars, #{line := Line, file := File}) ->
+compile(Module, _Block, _Vars, _Prune, #{line := Line, file := File}) ->
   elixir_errors:form_error([{line, Line}], File, ?MODULE, {invalid_module, Module}).
 
-compile(Line, Module, Block, Vars, E) ->
+compile(Line, Module, Block, Vars, Prune, E) ->
   File = ?key(E, file),
   check_module_availability(Line, File, Module),
   ModuleAsCharlist = validate_module_name(Line, File, Module),
@@ -102,7 +102,7 @@ compile(Line, Module, Block, Vars, E) ->
 
   try
     put_compiler_modules([Module | CompilerModules]),
-    {Result, NE} = eval_form(Line, Module, DataBag, Block, Vars, E),
+    {Result, ModuleE, CallbackE} = eval_form(Line, Module, DataBag, Block, Vars, Prune, E),
     CheckerInfo = checker_info(),
 
     {Binary, PersistedAttributes, Autoload} =
@@ -133,7 +133,7 @@ compile(Line, Module, Block, Vars, E) ->
         CompileOpts = validate_compile_opts(RawCompileOpts, AllDefinitions, Unreachable, File, Line),
 
         AfterVerify = bag_lookup_element(DataBag, {accumulate, after_verify}, 2),
-        [elixir_env:trace({remote_function, [], VerifyMod, VerifyFun, 1}, E) ||
+        [elixir_env:trace({remote_function, [], VerifyMod, VerifyFun, 1}, CallbackE) ||
          {VerifyMod, VerifyFun} <- AfterVerify],
 
         ModuleMap = #{
@@ -158,8 +158,8 @@ compile(Line, Module, Block, Vars, E) ->
       end),
 
     Autoload andalso code:load_binary(Module, beam_location(ModuleAsCharlist), Binary),
-    eval_callbacks(Line, DataBag, after_compile, [NE, Binary], NE),
-    elixir_env:trace({on_module, Binary, none}, E),
+    eval_callbacks(Line, DataBag, after_compile, [CallbackE, Binary], CallbackE),
+    elixir_env:trace({on_module, Binary, none}, ModuleE),
     warn_unused_attributes(File, DataSet, DataBag, PersistedAttributes),
     make_module_available(Module, Binary),
     (CheckerInfo == undefined) andalso
@@ -375,13 +375,27 @@ build(Line, File, Module) ->
 
 %% Handles module and callback evaluations.
 
-eval_form(Line, Module, DataBag, Block, Vars, E) ->
-  {Value, EE} = elixir_compiler:compile(Block, Vars, E),
+eval_form(Line, Module, DataBag, Block, Vars, Prune, E) ->
+  {Value, ExS, EE} = elixir_compiler:compile(Block, Vars, E),
   elixir_overridable:store_not_overridden(Module),
   EV = (elixir_env:reset_vars(EE))#{line := Line},
   EC = eval_callbacks(Line, DataBag, before_compile, [EV], EV),
   elixir_overridable:store_not_overridden(Module),
-  {Value, EC}.
+  {Value, maybe_prune_versioned_vars(Prune, Vars, ExS, E), EC}.
+
+maybe_prune_versioned_vars(false, _Vars, _Exs, E) ->
+  E;
+maybe_prune_versioned_vars(true, Vars, ExS, E) ->
+  PruneBefore = length(Vars),
+  #elixir_ex{vars={ExVars, _}, unused={Unused, _}} = ExS,
+
+  VersionedVars =
+    maps:filter(fun
+      (Pair, Version) when Version < PruneBefore, not is_map_key({Pair, Version}, Unused) -> false;
+      (_, _) -> true
+    end, ExVars),
+
+  E#{versioned_vars := VersionedVars}.
 
 eval_callbacks(Line, DataBag, Name, Args, E) ->
   Callbacks = bag_lookup_element(DataBag, {accumulate, Name}, 2),

lib/elixir/test/elixir/code_test.exs

@@ -206,6 +206,34 @@ defmodule CodeTest do
              {[{{:x, :foo}, 2}], [x: :foo]}
   end
 
+  defmodule Tracer do
+    def trace(event, env) do
+      send(self(), {:trace, event, env})
+      :ok
+    end
+  end
+
+  test "eval_quoted_with_env/3 with tracing and pruning" do
+    env = %{Code.env_for_eval(__ENV__) | tracers: [Tracer], function: nil}
+    binding = [x: 1, y: 2, z: 3]
+
+    quoted =
+      quote do
+        defmodule Elixir.CodeTest.TracingPruning do
+          var!(y) = :updated
+          var!(y)
+          var!(x)
+        end
+      end
+
+    {_, binding, env} = Code.eval_quoted_with_env(quoted, binding, env, prune_binding: true)
+    assert Enum.sort(binding) == []
+    assert env.versioned_vars == %{}
+
+    assert_receive {:trace, {:on_module, _, _}, %{module: CodeTest.TracingPruning} = trace_env}
+    assert trace_env.versioned_vars == %{{:result, Kernel} => 5, {:x, nil} => 1, {:y, nil} => 4}
+  end
+
   test "compile_file/1" do
     assert Code.compile_file(fixture_path("code_sample.exs")) == []
     refute fixture_path("code_sample.exs") in Code.required_files()