Repository: adobe/elixir-styler
Pull request #237 (2025-07-14 18:56:01)

Add apply alias code
Closes #235

todos
- [x] test for an alias being removed from inside a quote
- [x] test for alias as


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index e505fe7..1650284 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ they can and will change without that change being reflected in Styler's semanti
 
 ### Improvements
 
+- apply aliases to code. if a module is aliased, and then later referenced with its full name, Styler will now shorten it to its alias. (#235, h/t me)
 - added `:minimum_supported_elixir_version` configuration to better support libraries using Styler (#231, h/t @maennchen)
 
 ### Fixes
diff --git a/docs/module_directives.md b/docs/module_directives.md
index 4bfcd59..cf32f1b 100644
--- a/docs/module_directives.md
+++ b/docs/module_directives.md
@@ -189,3 +189,29 @@ You can specify additional modules to exclude from lifting via the `:alias_lifti
   styler: [alias_lifting_exclude: [:C]],
 ]
 ```
+
+## Alias Application
+
+Styler applies aliases in those cases where a developer wrote out a full module name without realizing that the module is already aliased.
+
+```elixir
+# Given
+alias A.B
+alias A.B.C
+alias A.B.C.D, as: X
+
+A.B.foo()
+A.B.C.foo()
+A.B.C.D.woo()
+C.D.woo()
+
+# Styled
+alias A.B
+alias A.B.C
+alias A.B.C.D, as: X
+
+B.foo()
+C.foo()
+X.woo()
+X.woo()
+```
diff --git a/lib/alias_env.ex b/lib/alias_env.ex
index 6da3615..23357ab 100644
--- a/lib/alias_env.ex
+++ b/lib/alias_env.ex
@@ -14,7 +14,7 @@ defmodule Styler.AliasEnv do
 
   Not anywhere as correct as what the compiler gives us, but close enough for open source work.
 
-  A alias env is a map from an alias's `as` to its resolution in a context.
+  An alias env is a map from an alias's `as` to its resolution in a context.
 
   Given the ast for
 
@@ -25,40 +25,79 @@ defmodule Styler.AliasEnv do
       %{:Bar => [:Foo, :Bar]}
   """
   def define(env \\ %{}, ast)
-
   def define(env, asts) when is_list(asts), do: Enum.reduce(asts, env, &define(&2, &1))
+  def define(env, {:alias, _, [{:__aliases__, _, aliases}]}), do: define(env, aliases, List.last(aliases))
+  def define(env, {:alias, _, [{:__aliases__, _, aliases}, [{_, {:__aliases__, _, [as]}}]]}), do: define(env, aliases, as)
+  # `alias __MODULE__` or other oddities i'm not bothering to get right
+  def define(env, {:alias, _, _}), do: env
 
-  def define(env, {:alias, _, aliases}) do
-    case aliases do
-      [{:__aliases__, _, aliases}] -> define(env, aliases, List.last(aliases))
-      [{:__aliases__, _, aliases}, [{_as, {:__aliases__, _, [as]}}]] -> define(env, aliases, as)
-      # `alias __MODULE__` or other oddities i'm not bothering to get right
-      _ -> env
-    end
-  end
+  defp define(env, modules, as), do: Map.put(env, as, expand(env, modules))
+
+  @doc """
+  Lengthens an alias to its full name, if its first name is defined in the environment"
+
+  Useful for transforming the ast for code like:
 
-  defp define(env, modules, as), do: Map.put(env, as, do_expand(env, modules))
+      alias Bar.Baz.Foo #<- given the env with this alias
+      Foo.Woo.Cool # <- ast
 
+  to the ast for code like:
+
+      alias Bar.Baz.Foo
+      Bar.Baz.Foo.Woo.Cool
+  """
   # no need to traverse ast if there are no aliases
-  def expand(env, ast) when map_size(env) == 0, do: ast
+  def expand_ast(env, ast) when map_size(env) == 0, do: ast
 
-  def expand(env, ast) do
+  def expand_ast(env, ast) do
     Macro.prewalk(ast, fn
-      {:__aliases__, meta, modules} -> {:__aliases__, meta, do_expand(env, modules)}
+      {:__aliases__, meta, modules} -> {:__aliases__, meta, expand(env, modules)}
       ast -> ast
     end)
   end
 
-  # if the list of modules is itself already aliased, dealias it with the compound alias
-  # given:
-  #   alias Foo.Bar
-  #   Bar.Baz.Bop.baz()
-  #
-  # lifting Bar.Baz.Bop should result in:
-  #   alias Foo.Bar
-  #   alias Foo.Bar.Baz.Bop
-  #   Bop.baz()
-  defp do_expand(env, [first | rest] = modules) do
+  @doc """
+  Expands modules from env (wow that was helpful).
+
+  Using the examples from `expand_ast`, this works roughly like so:
+
+     > expand(%{Foo: [Bar, Baz, Foo]}, [Foo, Woo, Cool])
+     => [Bar, Baz, Foo, Woo, Cool]
+     > expand(%{}, [No, Alias, For, Me])
+     => [No, Alias, For, Me]
+  """
+  def expand(env, [first | rest] = modules) do
     if dealias = env[first], do: dealias ++ rest, else: modules
   end
+
+  @doc """
+  An inverted AliasEnv is useful for translating a module to its alias, if one existed in the env
+
+  In the case that a module is aliased multiple times, the inverted env will only keep the final alias as lexically sorted
+  """
+  def invert(env) do
+    # It's a bit of a bummer to do the extra group_by out of caution that this 1-off mistake happens,
+    # but ultimately we're usually working with a small list so performance costs are negligible
+    env
+    |> Enum.group_by(fn {_, v} -> v end, fn {k, _} -> k end)
+    |> Map.new(fn
+      {modules, [as]} ->
+        {modules, as}
+
+      # someone has something goofy going on, aliasing the same module with multiple names
+      # alias A.B.C
+      # alias A.B.C, as: Bar
+      # alias A.B.C, as: Foo
+      # we'll choose the one that comes last lexically, which will be the alpha-sorted last entry that isn't the default as
+      # "bug": if the last happens to be `alias A.B.C, as: C`, well, they shouldn't've written such crazy code
+      {modules, multiple_as} ->
+        default_as = List.last(modules)
+        # being clever - rather than rejecting the default up front and doing an extra list-traversal,
+        # just sort things and if the default comes first, grab the second element
+        case Enum.sort(multiple_as, :desc) do
+          [^default_as, last_as | _] -> {modules, last_as}
+          [last_as | _] -> {modules, last_as}
+        end
+    end)
+  end
 end
diff --git a/lib/style/module_directives.ex b/lib/style/module_directives.ex
index b88f3c6..1072531 100644
--- a/lib/style/module_directives.ex
+++ b/lib/style/module_directives.ex
@@ -152,13 +152,13 @@ defmodule Styler.Style.ModuleDirectives do
     alias: [],
     require: [],
     nondirectives: [],
-    dealiases: %{},
+    alias_env: %{},
     attrs: MapSet.new(),
     attr_lifts: []
   }
 
   defp lift_module_attrs({node, _, _} = ast, %{attrs: attrs} = acc) do
-    if Enum.empty?(attrs) do
+    if MapSet.size(attrs) == 0 do
       {ast, acc}
     else
       use? = node == :use
@@ -188,8 +188,8 @@ defmodule Styler.Style.ModuleDirectives do
       |> Zipper.children()
       |> Enum.reduce(@acc, fn
         {:@, _, [{attr_directive, _, _}]} = ast, acc when attr_directive in @attr_directives ->
-          # attr_directives are moved above aliases, so we need to dealias them
-          {ast, acc} = acc.dealiases |> AliasEnv.expand(ast) |> lift_module_attrs(acc)
+          # attr_directives are moved above aliases, so we need to expand them
+          {ast, acc} = acc.alias_env |> AliasEnv.expand_ast(ast) |> lift_module_attrs(acc)
           %{acc | attr_directive => [ast | acc[attr_directive]]}
 
         {:@, _, [{attr, _, _}]} = ast, acc ->
@@ -198,12 +198,12 @@ defmodule Styler.Style.ModuleDirectives do
         {directive, _, _} = ast, acc when directive in @directives ->
           {ast, acc} = lift_module_attrs(ast, acc)
           ast = expand(ast)
-          # import and used get hoisted above aliases, so need to dealias
-          ast = if directive in ~w(import use)a, do: AliasEnv.expand(acc.dealiases, ast), else: ast
-          dealiases = if directive == :alias, do: AliasEnv.define(acc.dealiases, ast), else: acc.dealiases
+          # import and used get hoisted above aliases, so need to expand them
+          ast = if directive in ~w(import use)a, do: AliasEnv.expand_ast(acc.alias_env, ast), else: ast
+          alias_env = if directive == :alias, do: AliasEnv.define(acc.alias_env, ast), else: acc.alias_env
 
           # the reverse accounts for `expand` putting things in reading order, whereas we're accumulating in reverse
-          %{acc | directive => Enum.reverse(ast, acc[directive]), dealiases: dealiases}
+          %{acc | directive => Enum.reverse(ast, acc[directive]), alias_env: alias_env}
 
         ast, acc ->
           %{acc | nondirectives: [ast | acc.nondirectives]}
@@ -213,10 +213,12 @@ defmodule Styler.Style.ModuleDirectives do
         {:moduledoc, []} -> {:moduledoc, List.wrap(moduledoc)}
         {:use, uses} -> {:use, uses |> Enum.reverse() |> Style.reset_newlines()}
         {directive, to_sort} when directive in ~w(behaviour import alias require)a -> {directive, sort(to_sort)}
-        {:dealiases, d} -> {:dealiases, d}
+        {:alias_env, d} -> {:alias_env, d}
         {k, v} -> {k, Enum.reverse(v)}
       end)
+      |> redefine_alias_env()
       |> lift_aliases()
+      |> apply_aliases()
 
     # Not happy with it, but this does the work to move module attribute assignments above the module or quote or whatever
     # Given that it'll only be run once and not again, i'm okay with it being inefficient
@@ -278,42 +280,45 @@ defmodule Styler.Style.ModuleDirectives do
     end
   end
 
-  defp lift_aliases(%{alias: aliases, require: requires, nondirectives: nondirectives} = acc) do
-    # we can't use the dealias map built into state as that's what things look like before sorting
-    # now that we've sorted, it could be different!
-    dealiases = AliasEnv.define(aliases)
-    liftable = find_liftable_aliases(requires ++ nondirectives, dealiases)
+  # alias_env have to be recomputed after we've sorted our `alias` nodes
+  defp redefine_alias_env(%{alias: aliases} = acc), do: %{acc | alias_env: AliasEnv.define(aliases)}
+
+  defp lift_aliases(%{alias: aliases, require: requires, nondirectives: nondirectives, alias_env: alias_env} = acc) do
+    liftable = find_liftable_aliases(requires ++ nondirectives, alias_env)
 
     if Enum.any?(liftable) do
       # This is a silly hack that helps comments stay put.
       # The `cap_line` algo was designed to handle high-line stuff moving up into low line territory, so we set our
       # new node to have an arbitrarily high line annnnd comments behave! i think.
       m = [line: 999_999]
+      aliases_m = [last: m] ++ m
 
       aliases =
         liftable
-        |> Enum.map(&AliasEnv.expand(dealiases, {:alias, m, [{:__aliases__, [{:last, m} | m], &1}]}))
+        |> Enum.map(&{:alias, m, [{:__aliases__, aliases_m, AliasEnv.expand(alias_env, &1)}]})
         |> Enum.concat(aliases)
         |> sort()
 
-      # lifting could've given us a new order
-      requires = requires |> do_lift_aliases(liftable) |> sort()
-      nondirectives = do_lift_aliases(nondirectives, liftable)
-      %{acc | alias: aliases, require: requires, nondirectives: nondirectives}
+      # aliases to be lifted have to be applied immediately; yes, these means we're doing multiple apply_alias traversals,
+      # but we're only doing the duplicate traversals when we're updating the file, so the extra cost of walking
+      # is negligible vs doing disk I/O
+      %{acc | alias: aliases}
+      |> apply_aliases(Map.new(liftable, fn modules -> {modules, List.last(modules)} end))
+      |> redefine_alias_env()
     else
       acc
     end
   end
 
-  defp find_liftable_aliases(ast, dealiases) do
-    excluded = dealiases |> Map.keys() |> Enum.into(Styler.Config.get(:lifting_excludes))
+  defp find_liftable_aliases(ast, alias_env) do
+    excluded = alias_env |> Map.keys() |> Enum.into(Styler.Config.get(:lifting_excludes))
 
-    firsts = MapSet.new(dealiases, fn {_last, [first | _]} -> first end)
+    firsts = MapSet.new(alias_env, fn {_last, [first | _]} -> first end)
 
     ast
     |> Zipper.zip()
     # we're reducing a datastructure that looks like
-    # %{last => {aliases, seen_before?} | :some_collision_probelm}
+    # %{last => {aliases, seen_before?} | :some_collision_problem}
     |> Zipper.reduce_while(%{}, fn
       # we don't want to rewrite alias name `defx Aliases ... do` of these three keywords
       {{defx, _, args}, _} = zipper, lifts when defx in ~w(defmodule defimpl defprotocol)a ->
@@ -328,7 +333,7 @@ defmodule Styler.Style.ModuleDirectives do
               lifts
           end
 
-        # move the focus to the body block, zkipping over the alias (and the `for` keyword for `defimpl`)
+        # move the focus to the body block, skipping over the alias (and the `for` keyword for `defimpl`)
         {:skip, zipper |> Zipper.down() |> Zipper.rightmost() |> Zipper.down() |> Zipper.down(), lifts}
 
       {{:quote, _, _}, _} = zipper, lifts ->
@@ -340,8 +345,9 @@ defmodule Styler.Style.ModuleDirectives do
         lifts =
           cond do
             # this alias already exists, they just wrote it out fully and are leaving it up to us to shorten it down!
-            dealiases[last] == aliases ->
-              Map.put(lifts, last, {aliases, true})
+            # we'll get it when we do the apply-aliases scan
+            alias_env[last] == aliases ->
+              lifts
 
             last in excluded or Enum.any?(aliases, &(not is_atom(&1))) ->
               lifts
@@ -433,27 +439,72 @@ defmodule Styler.Style.ModuleDirectives do
     |> MapSet.new(fn {_, {aliases, true}} -> aliases end)
   end
 
-  defp do_lift_aliases(ast, to_alias) do
+  defp apply_aliases(acc) do
+    apply_aliases(acc, AliasEnv.invert(acc.alias_env))
+  end
+
+  defp apply_aliases(acc, inverted_env) when map_size(inverted_env) == 0, do: acc
+
+  defp apply_aliases(%{require: requires, nondirectives: nondirectives, alias_env: alias_env} = acc, inverted_env) do
+    # applying aliases to requires can change their ordering again
+    requires = requires |> apply_aliases(inverted_env, alias_env) |> sort()
+    nondirectives = apply_aliases(nondirectives, inverted_env, alias_env)
+    %{acc | require: requires, nondirectives: nondirectives}
+  end
+
+  # applies the aliases withi `to_as` across the given ast
+  # alias_env is used to expand partial aliases
+  defp apply_aliases(ast, to_as, alias_env) do
     ast
     |> Zipper.zip()
-    |> Zipper.traverse(fn
+    |> Zipper.traverse_while(fn
       {{defx, _, [{:__aliases__, _, _} | _]}, _} = zipper when defx in ~w(defmodule defimpl defprotocol)a ->
-        # move the focus to the body block, zkipping over the alias (and the `for` keyword for `defimpl`)
-        zipper |> Zipper.down() |> Zipper.rightmost() |> Zipper.down() |> Zipper.down() |> Zipper.right()
-
-      {{:alias, _, [{:__aliases__, _, [_, _, _ | _] = aliases}]}, _} = zipper ->
-        # the alias was aliased deeper down. we've lifted that alias to a root, so delete this alias
-        if aliases in to_alias,
-          do: Zipper.remove(zipper),
-          else: zipper
+        # move the focus to the body block, skipping over the alias (and the `for` keyword for `defimpl`)
+        zipper = zipper |> Zipper.down() |> Zipper.rightmost() |> Zipper.down() |> Zipper.down() |> Zipper.right()
+        {:cont, zipper}
+
+      {{:quote, _, _}, _} = zipper ->
+        {:skip, zipper}
+
+      # apply_aliases is only called on nondirectives (+requires), so this alias must exist within a child node like a def
+      # if it's within `to_as` then it's an alias that's already defined further up and is thus a duplicate
+      # either because we lifted and so created a duplicate, or it was just plain ol' user user error.
+      # either way, we can nix it
+      {{:alias, _, [{:__aliases__, _, [_ | _] = modules}]}, _} = zipper ->
+        zipper = if to_as[modules], do: Zipper.remove(zipper), else: zipper
+        {:cont, zipper}
+
+      # We check even modules of 1 length to catch silly situations like
+      # alias A.B.C
+      # alias A.B.C, as: X
+      # That'll then rename all C to X and C will become unused
+      {{:__aliases__, meta, [_ | _] = modules}, _} = zipper ->
+        # if modules |> List.last|> to_string() |> String.ends_with?("Mock") do
+        #   dbg(to_as)
+        #   dbg(alias_env)
+        #   dbg(modules)
+        #   dbg(to_as[modules])
+        #   dbg(AliasEnv.expand(alias_env, modules))
+        #   dbg(to_as[AliasEnv.expand(alias_env, modules)])
+        # end
+        zipper =
+          cond do
+            # There's an alias for this module - replace it with its `as`
+            as = to_as[modules] -> Zipper.replace(zipper, {:__aliases__, meta, [as]})
+            # There's an alias for this modules expansion - replace it with its `as`
+            #
+            # This addresses the following:
+            # given modules=`C.D`, to_as=`A.B.C => C, A.B.C.D => X`,
+            # should yield -> `X` because `C.D -> A.B.C.D -> X`
+            as = to_as[AliasEnv.expand(alias_env, modules)] -> Zipper.replace(zipper, {:__aliases__, meta, [as]})
+            true -> zipper
+          end
 
-      {{:__aliases__, meta, [_, _, _ | _] = aliases}, _} = zipper ->
-        if aliases in to_alias,
-          do: Zipper.replace(zipper, {:__aliases__, meta, [List.last(aliases)]}),
-          else: zipper
+        # minor optimization - don't need to walk the module list!
+        {:skip, zipper}
 
       zipper ->
-        zipper
+        {:cont, zipper}
     end)
     |> Zipper.node()
   end
@@ -490,7 +541,6 @@ defmodule Styler.Style.ModuleDirectives do
     |> Style.reset_newlines()
   end
 
-  # TODO investigate removing this in favor of the Style.post_sort_cleanup(node, comments)
   # "Fixes" the line numbers of nodes who have had their orders changed via sorting or other methods.
   # This "fix" simply ensures that comments don't get wrecked as part of us moving AST nodes willy-nilly.
   #


Test patch:
diff --git a/test/alias_env_test.exs b/test/alias_env_test.exs
new file mode 100644
index 0000000..4021849
--- /dev/null
+++ b/test/alias_env_test.exs
@@ -0,0 +1,75 @@
+# Copyright 2025 Adobe. All rights reserved.
+# This file is licensed to you under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License. You may obtain a copy
+# of the License at http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software distributed under
+# the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
+# OF ANY KIND, either express or implied. See the License for the specific language
+# governing permissions and limitations under the License.
+
+defmodule Styler.AliasEnvTest do
+  use ExUnit.Case, async: true
+
+  import Styler.AliasEnv
+
+  test "define" do
+    {:__block__, [], ast} =
+      quote do
+        alias A.B
+        alias C.D, as: X
+      end
+
+    env = define(ast)
+
+    assert %{B: [:A, :B], X: [:C, :D]} == env
+    assert %{B: [:M, :N, :B], X: [:C, :D]} == define(env, quote(do: alias(M.N.B)))
+  end
+
+  test "expand" do
+    assert expand(%{B: [:A, :B], X: [:C, :D]}, [:B]) == [:A, :B]
+    assert expand(%{B: [:A, :B], X: [:C, :D]}, [:B, :C, :D]) == [:A, :B, :C, :D]
+    assert expand(%{B: [:A, :B], X: [:C, :D]}, [:Not, :Present]) == [:Not, :Present]
+    assert expand(%{}, [:Hi]) == [:Hi]
+  end
+
+  test "expand_ast" do
+    {_, _, aliases} =
+      quote do
+        alias A.B
+        alias A.B.C
+        alias A.B.C.D, as: X
+      end
+
+    ast =
+      quote do
+        A
+        B
+        C
+        X
+      end
+
+    expected =
+      quote do
+        A
+        A.B
+        A.B.C
+        A.B.C.D
+      end
+
+    assert aliases |> define() |> expand_ast(ast) == expected
+  end
+
+  test "invert" do
+    {_, _, aliases} =
+      quote do
+        alias A.B
+        alias A.B.C
+        alias A.B.C, as: X
+        alias A.B.C, as: Y
+      end
+
+    env = define(aliases)
+    assert invert(env) == %{[:A, :B, :C] => :Y, [:A, :B] => :B}
+  end
+end
diff --git a/test/style/module_directives/alias_lifting_test.exs b/test/style/module_directives/alias_lifting_test.exs
index 4557876..2de0724 100644
--- a/test/style/module_directives/alias_lifting_test.exs
+++ b/test/style/module_directives/alias_lifting_test.exs
@@ -202,25 +202,6 @@ defmodule Styler.Style.ModuleDirectives.AliasLiftingTest do
     )
   end
 
-  test "replaces known aliases" do
-    assert_style(
-      """
-      alias A.B.C
-
-      A.B.C.foo()
-      A.B.C.foo()
-      A.B.C.foo()
-      """,
-      """
-      alias A.B.C
-
-      C.foo()
-      C.foo()
-      C.foo()
-      """
-    )
-  end
-
   test "two modules that seem to conflict but don't!" do
     assert_style(
       """
diff --git a/test/style/module_directives_test.exs b/test/style/module_directives_test.exs
index b65518a..d8e618f 100644
--- a/test/style/module_directives_test.exs
+++ b/test/style/module_directives_test.exs
@@ -407,7 +407,7 @@ defmodule Styler.Style.ModuleDirectivesTest do
         alias B.B
         alias D.D
 
-        require A.A
+        require A
         require A.C
 
         @type foo :: :ok
@@ -619,4 +619,97 @@ defmodule Styler.Style.ModuleDirectivesTest do
       )
     end
   end
+
+  describe "apply aliases" do
+    test "replaces known aliases" do
+      assert_style(
+        """
+        alias A.B
+        alias A.B.C
+        alias A.B.C.D, as: X
+
+        A.B.foo()
+        A.B.C.foo()
+        A.B.C.D.woo()
+        C.D.woo()
+        """,
+        """
+        alias A.B
+        alias A.B.C
+        alias A.B.C.D, as: X
+
+        B.foo()
+        C.foo()
+        X.woo()
+        X.woo()
+        """
+      )
+    end
+
+    test "ignores quotes" do
+      assert_style(
+        """
+        alias A.B.C
+
+        A.B.C
+
+        quote do
+          A.B.C
+        end
+        """,
+        """
+        alias A.B.C
+
+        C
+
+        quote do
+          A.B.C
+        end
+        """
+      )
+    end
+
+    test "removes embedded duplicate aliases" do
+      assert_style(
+        """
+        alias A.B
+
+        def foo do
+          alias A.B
+          A.B.bar()
+        end
+        """,
+        """
+        alias A.B
+
+        def foo do
+          B.bar()
+        end
+        """
+      )
+    end
+
+    test "forces a single alias" do
+      assert_style(
+        """
+        alias A.B.C.D.E, as: B
+        alias A.B.C.D.E, as: C
+        alias A.B.C.D.E
+
+        B
+        C
+        E
+        """,
+        """
+        alias A.B.C.D.E
+        alias A.B.C.D.E, as: B
+        alias A.B.C.D.E, as: C
+
+        C
+        C
+        C
+        """
+      )
+    end
+  end
 end