Ensure context modules are handled in optimized defmodule

2 files changed, 10 insertions, 1 deletions

diff --git a/lib/elixir/src/elixir_compiler.erl b/lib/elixir/src/elixir_compiler.erl
index da27e0f91..b142442a8 100644
--- a/lib/elixir/src/elixir_compiler.erl
+++ b/lib/elixir/src/elixir_compiler.erl
@@ -137,7 +137,8 @@ fast_compile({defmodule, Meta, [Mod, [{do, TailBlock}]]}, NoLineE) ->
       'Elixir.Macro':expand(Mod, E)
   end,
 
-  elixir_module:compile(Expanded, Block, [], E).
+  ContextModules = [Expanded | ?key(E, context_modules)],
+  elixir_module:compile(Expanded, Block, [], E#{context_modules := ContextModules}).
 
 %% Bootstrapper
 
diff --git a/lib/elixir/test/elixir/macro_test.exs b/lib/elixir/test/elixir/macro_test.exs
index f01168ca9..1ad2702fc 100644
--- a/lib/elixir/test/elixir/macro_test.exs
+++ b/lib/elixir/test/elixir/macro_test.exs
@@ -815,6 +815,14 @@ defmodule MacroTest do
       defmodule Foo.Bar do
         assert __MODULE__ in __ENV__.context_modules
       end
+
+      assert Foo.Bar in __ENV__.context_modules
+
+      Code.compile_string("""
+      defmodule Foo.Bar.Compiled do
+        true = __MODULE__ in __ENV__.context_modules
+      end
+      """)
     end
 
     test "to_match/1" do