Support nested modules

lib/gradient/ast_specifier.ex

@@ -64,16 +64,8 @@ defmodule Gradient.AstSpecifier do
   """
   @spec specify(nonempty_list(:erl_parse.abstract_form())) :: [:erl_parse.abstract_form()]
   def specify(forms) do
-    with {:attribute, line, :file, {path, _}} <- hd(forms),
-         path <- to_string(path),
-         {:ok, code} <- File.read(path),
-         {:ok, tokens} <- :elixir.string_to_tokens(String.to_charlist(code), line, line, path, []) do
-      run_mappers(forms, tokens)
-    else
-      error ->
-        IO.puts("Error occurred when specifying forms : #{inspect(error)}")
-        forms
-    end
+    tokens = Gradient.ElixirFileUtils.load_tokens(forms)
+    run_mappers(forms, tokens)
   end
 
   @doc """

lib/gradient/elixir_file_utils.ex

@@ -3,6 +3,8 @@ defmodule Gradient.ElixirFileUtils do
   Module used to load beam files generated from Elixir.
   """
 
+  alias Gradient.Types
+
   @type path() :: :file.filename() | String.t()
 
   @type abstract_forms() :: [:erl_parse.abstract_form() | :erl_parse.form_info()]
@@ -39,4 +41,18 @@ defmodule Gradient.ElixirFileUtils do
         {:forms_error, reason}
     end
   end
+
+  @spec load_tokens([:erl_parse.abstract_form()]) :: Types.tokens()
+  def load_tokens(forms) do
+    with [{:attribute, _, :file, {path, _}} | _] <- forms,
+         path <- to_string(path),
+         {:ok, code} <- File.read(path),
+         {:ok, tokens} <- :elixir.string_to_tokens(String.to_charlist(code), 1, 1, path, []) do
+      tokens
+    else
+      error ->
+        IO.puts("Cannot load tokens: #{inspect(error)}")
+        []
+    end
+  end
 end

test/examples/nested_modules.ex

@@ -0,0 +1,17 @@
+defmodule NestedModules do
+  defmodule ModuleA do
+    def name do
+      :module_a
+    end
+  end
+
+  defmodule ModuleB do
+    def name do
+      :module_b
+    end
+  end
+
+  def name do
+    :module
+  end
+end

test/gradient/ast_specifier_test.exs

@@ -1578,6 +1578,19 @@ defmodule Gradient.AstSpecifierTest do
     assert [_] = AstSpecifier.run_mappers(forms, [])
   end
 
+  test "nested modules" do
+    {tokensA, astA} = load("Elixir.NestedModules.ModuleA.beam", "nested_modules.ex")
+    {tokensB, astB} = load("Elixir.NestedModules.ModuleB.beam", "nested_modules.ex")
+    {tokens, ast} = load("Elixir.NestedModules.beam", "nested_modules.ex")
+
+    assert {:function, 3, :name, 0, [{:clause, 3, [], [], [{:atom, 4, :module_a}]}]} =
+             List.last(AstSpecifier.run_mappers(astA, tokensA))
+    assert {:function, 9, :name, 0, [{:clause, 9, [], [], [{:atom, 10, :module_b}]}]} =
+             List.last(AstSpecifier.run_mappers(astB, tokensB))
+    assert {:function, 14, :name, 0, [{:clause, 14, [], [], [{:atom, 15, :module}]}]} =
+             List.last(AstSpecifier.run_mappers(ast, tokens))
+  end
+
   # Helpers
 
   def filter_attributes(ast, type) do

test/support/helpers.ex

@@ -8,18 +8,13 @@ defmodule Gradient.TestHelpers do
     beam_file = String.to_charlist(@examples_path <> beam_file)
     ex_file = @examples_path <> ex_file
 
-    code =
-      File.read!(ex_file)
-      |> String.to_charlist()
-
-    {:ok, tokens} =
-      code
-      |> :elixir.string_to_tokens(1, 1, ex_file, [])
-
     {:ok, {_, [abstract_code: {:raw_abstract_v1, ast}]}} =
       :beam_lib.chunks(beam_file, [:abstract_code])
 
     ast = replace_file_path(ast, ex_file)
+
+    [_ | _] = tokens = Gradient.ElixirFileUtils.load_tokens(ast)
+
     {tokens, ast}
   end