Repository: adobe/elixir-styler
Pull request #8 (2023-04-06 00:42:23)

Optimize module directives, remove concept of directive groups
Changes:

- directives: will remove newlines between groups: no more empty space between your aliases etc
- directives: orders different directives in all contexts, not just module bodies (functions, quote do, etc, will now be organized)
- rename `Zipper.make_node/2` -> `Zipper.replace_children/2`, which made it a little clearer to me what problems it can solve! 
- `@moduledoc` `@behaviour` and `@shortdoc` will not be noticed outside of a defmodule if they are alone in the context with no `use/alias/import/require`. you can see this in the implementation of run: it only catches `defmodule, use, alias, import, require`. i feel like this is.. too much of an edgecase to deal with for the moment, as it'll only occur in super weird macro situations? ala `quote do @moduledoc ...; @behaviour ...; end`

## Looking for feedback on

- whatever your heart desires!
- lmk if you have a better name for the zipper api change (`Zipper.set_children`?)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0b14858..3dbc089 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -1,16 +1,23 @@
 # Changelog
 
+## main
+
+### Improvements
+
+* Enabled `Defs` style and overhauled it to properly handles comments
+
 ## v0.2.0
 
 ### Improvements
 
-* Added `Styler.Style.ModuleDirectives`
+* Added `ModuleDirectives` style
     * Note that this is potentially destructive in some rare cases. See moduledoc for more.
+    * This supersedes the `Aliases` style, which has been removed.
 * `mix style -` reads and writes to stdin/stdout
 
 ### Fixes
 
-* `Styler.Style.Pipes` is now aware of `unless` blocks
+* `Pipes` style is now aware of `unless` blocks
 
 ## v0.1.1
 
diff --git a/README.md b/README.md
index 2694330..cac89dd 100644
--- a/README.md
+++ b/README.md
@@ -12,7 +12,7 @@ Add `:styler` as a dependency to your project's `mix.exs`:
 ```elixir
 def deps do
   [
-    {:styler, "~> 0.2", only: [:dev, :test], runtime: false},
+    {:styler, "~> 0.3", only: [:dev, :test], runtime: false},
   ]
 end
 ```
@@ -64,13 +64,6 @@ You can find the currently-enabled styles in the `Mix.Tasks.Style` module, insid
 
 If you're using Credo and Styler, we recommend disabling these rules in Credo to save on unnecessary checks in CI.
 
-### Styler and Comments...
-
-Styler is currently unaware of comments, so you may find that they end up in really odd spots after a rewrite.
-
-If you find that a comment was put somewhere weird after using Styler, you'll just have to manually put it back where you want it after.
-Feel free to grumble about it in an Issue so that we can properly prioritize making this work better in the future.
-
 ## Thanks & Inspiration
 
 ### [Sourceror](https://github.com/doorgan/sourceror/)
diff --git a/lib/mix/tasks/style.ex b/lib/mix/tasks/style.ex
index ea80869..c7f2a01 100644
--- a/lib/mix/tasks/style.ex
+++ b/lib/mix/tasks/style.ex
@@ -36,7 +36,8 @@ defmodule Mix.Tasks.Style do
   @styles [
     Styler.Style.ModuleDirectives,
     Styler.Style.Pipes,
-    Styler.Style.Simple
+    Styler.Style.Simple,
+    Styler.Style.Defs
   ]
 
   @impl Mix.Task
diff --git a/lib/style.ex b/lib/style.ex
index 1329b04..294a484 100644
--- a/lib/style.ex
+++ b/lib/style.ex
@@ -30,4 +30,31 @@ defmodule Styler.Style do
   by calling `Zipper.skip(zipper)` to skip an entire subtree you know is of no interest to your Style.
   """
   @callback run(Zipper.zipper(), context()) :: {Zipper.command(), Zipper.zipper(), context()}
+
+  @doc """
+  Set the line of all comments with `line` in `range_start..range_end` to instead have line `range_start`
+  """
+  def displace_comments(comments, range) do
+    Enum.map(comments, fn comment ->
+      if comment.line in range do
+        %{comment | line: range.first}
+      else
+        comment
+      end
+    end)
+  end
+
+  @doc """
+  Change the `line` of all comments with `line` in `range` by adding `delta` to it.
+  A positive delta will move the lines further down a file, while a negative delta will move them up.
+  """
+  def shift_comments(comments, range, delta) do
+    Enum.map(comments, fn comment ->
+      if comment.line in range do
+        %{comment | line: comment.line + delta}
+      else
+        comment
+      end
+    end)
+  end
 end
diff --git a/lib/style/defs.ex b/lib/style/defs.ex
index 0cf762e..fc67a16 100644
--- a/lib/style/defs.ex
+++ b/lib/style/defs.ex
@@ -10,12 +10,6 @@
 
 defmodule Styler.Style.Defs do
   @moduledoc """
-  NOT ENABLED
-  Currently has a bug where it puts comments into bad places no matter what, since it's
-  always rewriting every head. It's been run on our codebase once though...
-
-  --------------------------------------
-
   Styles function heads so that they're as small as possible.
 
   The goal is that a function head fits on a single line.
@@ -41,6 +35,7 @@ defmodule Styler.Style.Defs do
 
   @behaviour Styler.Style
 
+  alias Styler.Style
   alias Styler.Zipper
 
   # a def with no body like
@@ -48,53 +43,99 @@ defmodule Styler.Style.Defs do
   #  def example(foo, bar \\ nil)
   #
   def run({{def, meta, [head]}, _} = zipper, ctx) when def in [:def, :defp] do
-    # There won't be any defs deeper in here, so lets skip ahead if we can
-    {:skip, Zipper.replace(zipper, {def, meta, [flatten_head(head, meta[:line])]}), ctx}
+    {_fn_name, head_meta, _children} = head
+    first_line = meta[:line]
+    last_line = head_meta[:closing][:line]
+
+    if first_line == last_line do
+      # Already collapsed
+      {:skip, zipper, ctx}
+    else
+      comments = Style.displace_comments(ctx.comments, first_line..last_line)
+      node = {def, meta, [flatten_head(head, meta[:line])]}
+      {:skip, Zipper.replace(zipper, node), %{ctx | comments: comments}}
+    end
   end
 
   # all the other kinds of defs!
   def run({{def, def_meta, [head, body]}, _} = zipper, ctx) when def in [:def, :defp] do
-    def_start_line = def_meta[:line]
-    # order matters here! the end of the def is where the `do`s line is - but only if there's a do end block.
-    # otherwise it's just where the end of the (def) expression is.
-    def_end_line = (def_meta[:do] || def_meta[:end_of_expression])[:line]
+    {def_line, do_line, end_line} =
+      if def_meta[:do] do
+        # This is a def with a do block, like
+        #
+        #  def example(foo, bar \\ nil) do
+        #    :ok
+        #  end
+        #
+        def_line = def_meta[:line]
+        do_line = def_meta[:do][:line]
+        end_line = def_meta[:end][:line]
+        {def_line, do_line, end_line}
+      else
+        # This is a def with a keyword do, like
+        #
+        #  def example(foo, bar \\ nil), do: :ok
+        #
+        [{{:__block__, do_meta, [:do]}, {_, body_meta, _}}] = body
+        def_line = def_meta[:line]
+        do_line = do_meta[:line]
+        end_line = body_meta[:closing][:line] || do_meta[:line]
+        {def_line, do_line, end_line}
+      end
 
-    if def_start_line == def_end_line do
-      {:skip, zipper, ctx}
-    else
-      head = flatten_head(head, def_start_line)
-
-      {def_meta, body_meta_rewriter} =
-        if def_meta[:do] do
-          # we're in a `def do ... end`
-          delta = def_end_line - def_start_line
-          up_by_delta = &(&1 - delta)
-
-          # this is what does the shrinking of the `def ... do` stanza
-          def_meta =
-            def_meta
-            |> Keyword.replace_lazy(:do, &Keyword.put(&1, :line, def_start_line))
-            |> Keyword.replace_lazy(:end, &Keyword.update!(&1, :line, up_by_delta))
-
-          # move all body line #s up by the amount we squished the head by
-          {def_meta, collapse_lines(up_by_delta)}
-        else
-          # we're in a `def, do:`
-          to_same_line = fn _ -> def_start_line end
-          {def_meta, collapse_lines(to_same_line)}
-        end
-
-      body = update_all_meta(body, body_meta_rewriter)
-
-      # There won't be any defs deeper in here, so lets skip ahead if we can
-      # @TODO this skips checking the body, which can be incorrect if therey's a `quote do def do ...` inside of it
-      {:skip, Zipper.replace(zipper, {def, def_meta, [head, body]}), ctx}
+    delta = def_line - do_line
+    move_up = &(&1 + delta)
+    set_to_def_line = fn _ -> def_line end
+
+    cond do
+      def_line == end_line ->
+        # Already collapsed
+        {:skip, zipper, ctx}
+
+      def_meta[:do] ->
+        # We're working on a def do ... end
+        def_meta =
+          def_meta
+          |> Keyword.replace_lazy(:do, &Keyword.update!(&1, :line, set_to_def_line))
+          |> Keyword.replace_lazy(:end, &Keyword.update!(&1, :line, move_up))
+
+        head = flatten_head(head, def_line)
+        body = update_all_meta(body, shift_lines(move_up))
+        node = {def, def_meta, [head, body]}
+
+        comments =
+          ctx.comments
+          |> Style.displace_comments(def_line..do_line)
+          |> Style.shift_comments(do_line..end_line, delta)
+
+        # @TODO this skips checking the body, which can be incorrect if therey's a `quote do def do ...` inside of it
+        {:skip, Zipper.replace(zipper, node), %{ctx | comments: comments}}
+
+      true ->
+        # We're working on a Keyword def do:
+        head = flatten_head(head, def_line)
+        body = update_all_meta(body, collapse_lines(set_to_def_line))
+        node = {def, def_meta, [head, body]}
+
+        comments = Style.displace_comments(ctx.comments, def_line..end_line)
+
+        # @TODO this skips checking the body, which can be incorrect if therey's a `quote do def do ...` inside of it
+        {:skip, Zipper.replace(zipper, node), %{ctx | comments: comments}}
     end
   end
 
   def run(zipper, ctx), do: {:cont, zipper, ctx}
 
   defp collapse_lines(line_mover) do
+    fn meta ->
+      meta
+      |> Keyword.replace_lazy(:line, line_mover)
+      |> Keyword.replace_lazy(:closing, &Keyword.replace_lazy(&1, :line, line_mover))
+      |> Keyword.delete(:newlines)
+    end
+  end
+
+  defp shift_lines(line_mover) do
     fn meta ->
       meta
       |> Keyword.replace_lazy(:line, line_mover)
diff --git a/lib/style/module_directives.ex b/lib/style/module_directives.ex
index 1fdbd62..2fa624b 100644
--- a/lib/style/module_directives.ex
+++ b/lib/style/module_directives.ex
@@ -65,18 +65,59 @@ defmodule Styler.Style.ModuleDirectives do
   """
   @behaviour Styler.Style
 
-  @directives ~w(alias import require)a
+  alias Styler.Zipper
+
+  @directives ~w(alias import require use)a
   @attr_directives ~w(moduledoc shortdoc behaviour)a
 
   # module names ending with these suffixes will not have a default moduledoc appended
   @dont_moduledoc ~w(Test Mixfile MixProject Controller Endpoint Repo Router Socket View HTML JSON)
   @moduledoc_false {:@, [], [{:moduledoc, [], [{:__block__, [], [false]}]}]}
 
-  def run({{:defmodule, def_meta, [name, [{mod_do, {:__block__, children_meta, children}}]]}, zipper_meta}, ctx) do
-    {directives, other} =
-      Enum.split_with(children, fn
+  def run({{:defmodule, _, children}, _} = zipper, ctx) do
+    [{:__aliases__, _, aliases}, [{{:__block__, do_meta, [:do]}, _module_body}]] = children
+    # Move the zipper's focus to the module's body
+    name = aliases |> List.last() |> to_string()
+    add_moduledoc? = do_meta[:format] != :keyword and not String.ends_with?(name, @dont_moduledoc)
+    body_zipper = zipper |> Zipper.down() |> Zipper.right() |> Zipper.down() |> Zipper.down() |> Zipper.right()
+
+    case Zipper.node(body_zipper) do
+      {:__block__, _, _} ->
+        {:skip, organize_directives(body_zipper, add_moduledoc?), ctx}
+
+      {:@, _, [{:moduledoc, _, _}]} ->
+        # a module whose only child is a moduledoc. nothing to do here!
+        # seems weird at first blush but lots of projects/libraries do this with their root namespace module
+        {:skip, zipper, ctx}
+
+      only_child ->
+        # There's only one child, and it's not a moduledoc. Conditionally add a moduledoc, then style the only_child
+        if add_moduledoc? do
+          body_zipper
+          |> Zipper.replace({:__block__, [], [@moduledoc_false, only_child]})
+          |> Zipper.down()
+          |> Zipper.right()
+          |> run(ctx)
+        else
+          run(body_zipper, ctx)
+        end
+    end
+  end
+
+  def run({{d, _, _}, _} = zipper, ctx) when d in @directives do
+    {:skip, zipper |> Zipper.up() |> organize_directives(), ctx}
+  end
+
+  def run(zipper, ctx), do: {:cont, zipper, ctx}
+
+  defp organize_directives(parent, add_moduledoc? \\ false) do
+    {directives, nondirectives} =
+      parent
+      |> Zipper.node()
+      |> Zipper.children()
+      |> Enum.split_with(fn
         {:@, _, [{attr, _, _}]} -> attr in @attr_directives
-        {directive, _, _} -> directive in [:use | @directives]
+        {directive, _, _} -> directive in @directives
         _ -> false
       end)
 
@@ -86,21 +127,17 @@ defmodule Styler.Style.ModuleDirectives do
         {directive, _, _} -> directive
       end)
 
-    # TODO: (optimization)
-    # now that we have use/import/alias/require, we might as well run
-    # them through the sort/expand/dedupe functionality and skip them in the traversal
     shortdocs = directives[:"@shortdoc"] || []
-    moduledocs = directives[:"@moduledoc"] || if needs_moduledoc?(name), do: [@moduledoc_false], else: []
-    # TODO sort behaviours?
-    behaviours = directives[:"@behaviour"] || []
-    behaviours = List.update_at(behaviours, -1, &set_newlines(&1, 2))
+    moduledocs = directives[:"@moduledoc"] || if add_moduledoc?, do: [@moduledoc_false], else: []
+    behaviours = expand_and_sort(directives[:"@behaviour"] || [])
 
-    uses = directives[:use] || []
-    imports = directives[:import] || []
-    aliases = directives[:alias] || []
-    requires = directives[:require] || []
+    uses = (directives[:use] || []) |> Enum.flat_map(&expand_directive/1) |> reset_newlines()
 
-    children =
+    imports = expand_and_sort(directives[:import] || [])
+    aliases = expand_and_sort(directives[:alias] || [])
+    requires = expand_and_sort(directives[:require] || [])
+
+    directives =
       Enum.concat([
         shortdocs,
         moduledocs,
@@ -108,88 +145,44 @@ defmodule Styler.Style.ModuleDirectives do
         uses,
         imports,
         aliases,
-        requires,
-        other
+        requires
       ])
 
-    zipper = {{:defmodule, def_meta, [name, [{mod_do, {:__block__, children_meta, children}}]]}, zipper_meta}
-    {:cont, zipper, ctx}
-  end
-
-  # a module whose only child is a moduledoc. pass it on through
-  def run({{:defmodule, _, [_, [{_, {:@, _, [{:moduledoc, _, _}]}}]]}, _} = zipper, ctx), do: {:cont, zipper, ctx}
+    parent = Zipper.update(parent, &Zipper.replace_children(&1, directives))
 
-  def run({{:defmodule, def_meta, [name, [{mod_do, mod_children}]]}, zipper_meta} = zipper, ctx) do
-    # a module with a single child. lets add moduledoc false
-    # ... unless it's a `defmodule Foo, do: ...`, that is
-    if needs_moduledoc?(name, mod_do) do
-      mod_children = {:__block__, [], [@moduledoc_false, mod_children]}
-      {:cont, {{:defmodule, def_meta, [name, [{mod_do, mod_children}]]}, zipper_meta}, ctx}
+    if Enum.empty?(nondirectives) do
+      parent
     else
-      {:cont, zipper, ctx}
+      {last_directive, meta} = parent |> Zipper.down() |> Zipper.rightmost()
+      {last_directive, %{meta | r: nondirectives}}
     end
   end
 
-  def run({{:use, _, _} = directive, meta}, ctx) do
-    [last | rest] = directive |> expand_directive() |> Enum.reverse()
-    meta = %{meta | l: rest ++ meta.l}
-
-    zipper =
-      case meta.r do
-        [{:use, _, _} | _] -> {last, meta}
-        _ -> {set_newlines(last, 2), meta}
-      end
-
-    {:skip, zipper, ctx}
-  end
-
-  def run({{d, _, _} = directive, %{l: left, r: right} = meta}, ctx) when d in @directives do
-    {right, directives} = consume_directive_group(d, [directive | right], [])
-
-    [last | rest] =
-      directives
-      # Credo does case-agnostic sorting, so we have to match that here
-      |> Enum.map(&{&1, &1 |> Macro.to_string() |> String.downcase()})
-      # a splash of deduping for happiness
-      |> Enum.uniq_by(&elem(&1, 1))
-      |> List.keysort(1, :desc)
-      |> Enum.map(&(&1 |> elem(0) |> set_newlines(1)))
-
-    zipper = {set_newlines(last, 2), %{meta | r: right, l: rest ++ left}}
-
-    {:skip, zipper, ctx}
+  defp expand_and_sort(directives) do
+    # sorting is done with `downcase` to match Credo
+    directives
+    |> Enum.flat_map(&expand_directive/1)
+    |> Enum.map(&{&1, &1 |> Macro.to_string() |> String.downcase()})
+    |> Enum.uniq_by(&elem(&1, 1))
+    |> List.keysort(1)
+    |> Enum.map(&elem(&1, 0))
+    |> reset_newlines()
   end
 
-  def run(zipper, ctx), do: {:cont, zipper, ctx}
-
-  def needs_moduledoc?({_, _, aliases}) do
-    name = aliases |> List.last() |> to_string()
-    not String.ends_with?(name, @dont_moduledoc)
-  end
-
-  def needs_moduledoc?(name, {_, do_meta, _}) do
-    needs_moduledoc?(name) and do_meta[:format] != :keyword
-  end
-
-  defp consume_directive_group(d, [{d, meta, _} = directive | siblings], directives) do
-    directives = expand_directive(directive) ++ directives
-
-    if meta[:end_of_expression][:newlines] == 1,
-      do: consume_directive_group(d, siblings, directives),
-      else: {siblings, directives}
-  end
-
-  defp consume_directive_group(_, siblings, directives), do: {siblings, directives}
-
   # alias Foo.{Bar, Baz}
   # =>
   # alias Foo.Bar
   # alias Foo.Baz
-  defp expand_directive({directive, _, [{{:., _, [{_, _, module}, :{}]}, _, right}]}) do
-    Enum.map(right, fn {_, meta, segments} -> {directive, meta, [{:__aliases__, [], module ++ segments}]} end)
-  end
+  defp expand_directive({directive, _, [{{:., _, [{_, _, module}, :{}]}, _, right}]}),
+    do: Enum.map(right, fn {_, meta, segments} -> {directive, meta, [{:__aliases__, [], module ++ segments}]} end)
+
+  defp expand_directive(other), do: [other]
+
+  defp reset_newlines([]), do: []
+  defp reset_newlines(directives), do: reset_newlines(directives, [])
 
-  defp expand_directive(alias), do: [alias]
+  defp reset_newlines([directive], acc), do: Enum.reverse([set_newlines(directive, 2) | acc])
+  defp reset_newlines([directive | rest], acc), do: reset_newlines(rest, [set_newlines(directive, 1) | acc])
 
   defp set_newlines({directive, meta, children}, newline) do
     updated_meta = Keyword.update(meta, :end_of_expression, [newlines: newline], &Keyword.put(&1, :newlines, newline))
diff --git a/lib/style/pipes.ex b/lib/style/pipes.ex
index d80f311..b816496 100644
--- a/lib/style/pipes.ex
+++ b/lib/style/pipes.ex
@@ -28,7 +28,7 @@ defmodule Styler.Style.Pipes do
   # we're in a multi-pipe, so only need to fix pipe_start
   def run({{:|>, _, [{:|>, _, _} | _]}, _} = zipper, ctx), do: {:cont, zipper |> check_start() |> Zipper.next(), ctx}
   # this is a single pipe, since valid pipelines are consumed by the previous head
-  def run({{:|>, meta, [lhs, {fun, _, args}]}, _} = zipper, ctx) do
+  def run({{:|>, _, [lhs, {fun, meta, args}]}, _} = zipper, ctx) do
     if valid_pipe_start?(lhs) do
       # `a |> f(b, c)` => `f(a, b, c)`
       {:cont, Zipper.replace(zipper, {fun, meta, [lhs | args]}), ctx}
diff --git a/lib/zipper.ex b/lib/zipper.ex
index 5e77d5d..a140e4d 100644
--- a/lib/zipper.ex
+++ b/lib/zipper.ex
@@ -58,12 +58,12 @@ defmodule Styler.Zipper do
   @doc """
   Returns a new branch node, given an existing node and new children.
   """
-  @spec make_node(tree, [tree]) :: tree
-  def make_node({form, meta, _}, args) when is_atom(form), do: {form, meta, args}
-  def make_node({_form, meta, args}, [first | rest]) when is_list(args), do: {first, meta, rest}
-  def make_node({_, _}, [left, right]), do: {left, right}
-  def make_node({_, _}, args), do: {:{}, [], args}
-  def make_node(list, children) when is_list(list), do: children
+  @spec replace_children(tree, [tree]) :: tree
+  def replace_children({form, meta, _}, children) when is_atom(form), do: {form, meta, children}
+  def replace_children({_form, meta, args}, [first | rest]) when is_list(args), do: {first, meta, rest}
+  def replace_children({_, _}, [left, right]), do: {left, right}
+  def replace_children({_, _}, children), do: {:{}, [], children}
+  def replace_children(list, children) when is_list(list), do: children
 
   @doc """
   Creates a zipper from a tree node.
@@ -112,14 +112,14 @@ defmodule Styler.Zipper do
   def up({tree, meta}) do
     children = Enum.reverse(meta.l, [tree | meta.r])
     {parent, parent_meta} = meta.ptree
-    {make_node(parent, children), parent_meta}
+    {replace_children(parent, children), parent_meta}
   end
 
   @doc """
   Returns the zipper of the left sibling of the node at this zipper, or nil.
   """
   @spec left(zipper) :: zipper | nil
-  def left({tree, %{l: [ltree | l], r: r} = meta}), do: {ltree, %{meta | l: l, r: [tree | r || []]}}
+  def left({tree, %{l: [ltree | l], r: r} = meta}), do: {ltree, %{meta | l: l, r: [tree | r]}}
   def left(_), do: nil
 
   @doc """
@@ -171,7 +171,7 @@ defmodule Styler.Zipper do
   @spec remove(zipper) :: zipper
   def remove({_, nil}), do: raise(ArgumentError, message: "Cannot remove the top level node.")
   def remove({_, %{l: [left | rest]} = meta}), do: prev_down({left, %{meta | l: rest}})
-  def remove({_, %{ptree: {parent, parent_meta}, r: children}}), do: {make_node(parent, children), parent_meta}
+  def remove({_, %{ptree: {parent, parent_meta}, r: children}}), do: {replace_children(parent, children), parent_meta}
 
   @doc """
   Inserts the item as the left sibling of the node at this zipper, without
diff --git a/mix.exs b/mix.exs
index f710b3a..3adb507 100644
--- a/mix.exs
+++ b/mix.exs
@@ -12,7 +12,7 @@ defmodule Styler.MixProject do
   use Mix.Project
 
   # Don't forget to bump the README when doing non-patch version changes
-  @version "0.2.0"
+  @version "0.3.0"
   @url "https://github.com/adobe/elixir-styler"
 
   def project do
@@ -41,7 +41,7 @@ defmodule Styler.MixProject do
 
   defp deps do
     [
-      {:ex_doc, "~> 0.29.4", runtime: false, only: :dev}
+      {:ex_doc, "~> 0.29", runtime: false, only: :dev}
     ]
   end
 


Test patch:
diff --git a/test/style/defs_test.exs b/test/style/defs_test.exs
index 54f7a66..cb88dfd 100644
--- a/test/style/defs_test.exs
+++ b/test/style/defs_test.exs
@@ -15,13 +15,19 @@ defmodule Styler.Style.DefsTest do
     test "function with do keyword" do
       assert_style(
         """
+        # Top comment
         def save(
+               # Socket comment
                %Socket{assigns: %{user: user, live_action: :new}} = initial_socket,
+               # Params comment
                params
              ),
              do: :ok
         """,
         """
+        # Top comment
+        # Socket comment
+        # Params comment
         def save(%Socket{assigns: %{user: user, live_action: :new}} = initial_socket, params), do: :ok
         """
       )
@@ -36,9 +42,11 @@ defmodule Styler.Style.DefsTest do
 
     test "block function body doesn't get newlined" do
       assert_style("""
+      # Here's a comment
       def some_function(%{id: id, type: type, processed_at: processed_at} = file, params, _)
           when type == :file and is_nil(processed_at) do
         with {:ok, results} <- FileProcessor.process(file) do
+          # This comment could make sense
           {:ok, post_process_the_results_somehow(results)}
         end
       end
@@ -57,12 +65,13 @@ defmodule Styler.Style.DefsTest do
         """
         def save(
                %Socket{assigns: %{user: user, live_action: :new}} = initial_socket,
-               params
+               params # Comments in the darndest places
              ) do
           :ok
         end
         """,
         """
+        # Comments in the darndest places
         def save(%Socket{assigns: %{user: user, live_action: :new}} = initial_socket, params) do
           :ok
         end
@@ -73,16 +82,22 @@ defmodule Styler.Style.DefsTest do
     test "no body" do
       assert_style(
         """
+        # Top comment
         def no_body(
-          foo,
-          bar
+          foo, # This is a foo
+          bar  # This is a bar
         )
 
+        # Another comment for this head
         def no_body(nil, _), do: nil
         """,
         """
+        # Top comment
+        # This is a foo
+        # This is a bar
         def no_body(foo, bar)
 
+        # Another comment for this head
         def no_body(nil, _), do: nil
         """
       )
@@ -94,18 +109,38 @@ defmodule Styler.Style.DefsTest do
         def foo(%{
           bar: baz
           })
+          # Self-documenting code!
           when baz in [
-            :a,
-            :b
+            :a, # Obviously, this is a
+            :b  # ... and this is b
           ],
           do: :never_write_code_like_this
         """,
         """
+        # Self-documenting code!
+        # Obviously, this is a
+        # ... and this is b
         def foo(%{bar: baz}) when baz in [:a, :b], do: :never_write_code_like_this
         """
       )
     end
 
+    test "keyword do with a list" do
+      assert_style(
+        """
+        def foo,
+          do: [
+            # Weirdo comment
+            :never_write_code_like_this
+          ]
+        """,
+        """
+        # Weirdo comment
+        def foo, do: [:never_write_code_like_this]
+        """
+      )
+    end
+
     test "rewrites subsequent definitions" do
       assert_style(
         """
@@ -113,12 +148,14 @@ defmodule Styler.Style.DefsTest do
 
         def foo(
           too,
+          # Long long is too long
           long
         ), do: :ok
         """,
         """
         def foo(), do: :ok
 
+        # Long long is too long
         def foo(too, long), do: :ok
         """
       )
@@ -127,23 +164,55 @@ defmodule Styler.Style.DefsTest do
     test "when clause with block do" do
       assert_style(
         """
+        # Foo takes a bar
         def foo(%{
           bar: baz
           })
+          # Baz should be either :a or :b
           when baz in [
             :a,
             :b
           ]
-          do
+          do # Weird place for a comment
+          # Above the body
           :never_write_code_like_this
+          # Below the body
         end
         """,
         """
+        # Foo takes a bar
+        # Baz should be either :a or :b
+        # Weird place for a comment
         def foo(%{bar: baz}) when baz in [:a, :b] do
+          # Above the body
           :never_write_code_like_this
+          # Below the body
         end
         """
       )
     end
+
+    test "Doesn't move stuff around if it would make the line too long" do
+      assert_style("""
+      @doc "this is a doc"
+      # And also a comment
+      def wow_this_function_name_is_super_long(it_also, has_a, ton_of, arguments),
+        do: "this is going to end up making the line too long if we inline it"
+
+      @doc "this is another function"
+      # And it also has a comment
+      def this_one_fits_on_one_line, do: :ok
+      """)
+    end
+
+    test "Doesn't collapse pipe chains in a def do ... end" do
+      assert_style("""
+      def foo(some_list) do
+        some_list
+        |> Enum.reject(&is_nil/1)
+        |> Enum.map(&transform/1)
+      end
+      """)
+    end
   end
 end
diff --git a/test/style/module_directives_test.exs b/test/style/module_directives_test.exs
index fce7f25..4c69bb8 100644
--- a/test/style/module_directives_test.exs
+++ b/test/style/module_directives_test.exs
@@ -19,6 +19,11 @@ defmodule Styler.Style.ModuleDirectivesTest do
         defmodule A do
         end
 
+        defmodule B do
+          defmodule C do
+          end
+        end
+
         defmodule Bar do
           alias Bop
 
@@ -34,12 +39,23 @@ defmodule Styler.Style.ModuleDirectivesTest do
         end
 
         def Foo, do: :ok
+
+        defmodule Foo do
+          alias Foo.{Bar, Baz}
+        end
         """,
         """
         defmodule A do
           @moduledoc false
         end
 
+        defmodule B do
+          @moduledoc false
+          defmodule C do
+            @moduledoc false
+          end
+        end
+
         defmodule Bar do
           @moduledoc false
           alias Bop
@@ -57,6 +73,12 @@ defmodule Styler.Style.ModuleDirectivesTest do
         end
 
         def Foo, do: :ok
+
+        defmodule Foo do
+          @moduledoc false
+          alias Foo.Bar
+          alias Foo.Baz
+        end
         """
       )
     end
@@ -75,21 +97,38 @@ defmodule Styler.Style.ModuleDirectivesTest do
       assert_style(
         """
         defmodule Foo do
+          @behaviour Lawful
           require A
           alias A
 
+          use B
+
           def c(x), do: y
+
+          import C
           @behaviour Chaotic
+          @doc "d doc"
           def d do
             alias X
             alias H
+
+            alias Z
             import Ecto.Query
             X.foo()
           end
           @shortdoc "it's pretty short"
           import A
+          alias C
+          alias D
+
+          require C
           require B
+
           use A
+
+          alias C
+          alias A
+
           @moduledoc "README.md"
                      |> File.read!()
                      |> String.split("<!-- MDOC !-->")
@@ -104,23 +143,31 @@ defmodule Styler.Style.ModuleDirectivesTest do
                      |> String.split("<!-- MDOC !-->")
                      |> Enum.fetch!(1)
           @behaviour Chaotic
+          @behaviour Lawful
 
+          use B
           use A
 
           import A
+          import C
 
           alias A
+          alias C
+          alias D
 
           require A
           require B
+          require C
 
           def c(x), do: y
 
+          @doc "d doc"
           def d do
+            import Ecto.Query
+
             alias H
             alias X
-
-            import Ecto.Query
+            alias Z
 
             X.foo()
           end
@@ -148,13 +195,12 @@ defmodule Styler.Style.ModuleDirectivesTest do
           #{d} A
           """,
           """
+          #{d} A
           #{d} A.A
           #{d} A.B
           #{d} A.C
-          #{d} D
-
-          #{d} A
           #{d} B
+          #{d} D
           """
         )
       end
@@ -163,8 +209,8 @@ defmodule Styler.Style.ModuleDirectivesTest do
     test "expands use but does not sort it" do
       assert_style(
         """
-        use A
         use D
+        use A
         use A.{
           C,
           B
@@ -172,8 +218,8 @@ defmodule Styler.Style.ModuleDirectivesTest do
         import F
         """,
         """
-        use A
         use D
+        use A
         use A.C
         use A.B
 
@@ -185,6 +231,7 @@ defmodule Styler.Style.ModuleDirectivesTest do
     test "interwoven directives w/o the context of a module" do
       assert_style(
         """
+        @type foo :: :ok
         alias D
         alias A.{B}
         require A.{
@@ -195,14 +242,15 @@ defmodule Styler.Style.ModuleDirectivesTest do
         alias A
         """,
         """
+        alias A
         alias A.B
+        alias B
         alias D
 
         require A.A
         require A.C
 
-        alias A
-        alias B
+        @type foo :: :ok
         """
       )
     end
diff --git a/test/style/pipes_test.exs b/test/style/pipes_test.exs
index 5c00d54..0ab3652 100644
--- a/test/style/pipes_test.exs
+++ b/test/style/pipes_test.exs
@@ -175,6 +175,18 @@ defmodule Styler.Style.PipesTest do
       )
     end
 
+    test "keeps invocation on a single line" do
+      assert_style(
+        """
+        foo
+        |> bar(baz, bop, boom)
+        """,
+        """
+        bar(foo, baz, bop, boom)
+        """
+      )
+    end
+
     test "extracts blocks successfully" do
       assert_style(
         """
diff --git a/test/style_test.exs b/test/style_test.exs
new file mode 100644
index 0000000..6d93669
--- /dev/null
+++ b/test/style_test.exs
@@ -0,0 +1,118 @@
+defmodule Styler.StyleTest do
+  use ExUnit.Case, async: true
+
+  import Styler.Style, only: [displace_comments: 2, shift_comments: 3]
+
+  @code """
+  # Above module
+  defmodule Foo do
+    # Top of module
+
+    @moduledoc "This is a moduledoc"
+
+
+    # Above spec
+    @spec some_fun(
+      atom()
+    )
+    :: atom()
+
+    @doc "This is a function"
+    # Above def
+    def some_fun(
+      # Before arg
+      arg
+      # After arg
+    ) do
+      # In function body
+      :ok
+    end
+
+    # After def
+  end
+  # After module
+  """
+
+  @comments @code |> Styler.string_to_quoted_with_comments() |> elem(1)
+
+  describe "displace_comments/2" do
+    test "Doesn't lose any comments" do
+      new_comments = displace_comments(@comments, 0..0)
+
+      for comment <- @comments do
+        assert Enum.any?(new_comments, &(&1.text == comment.text))
+      end
+    end
+
+    test "Moves comments within the range to the start of the range" do
+      # Sanity-check line numbers in test fixture
+      before = [
+        {"# Above def", 15},
+        {"# Before arg", 17},
+        {"# After arg", 19}
+      ]
+
+      for {text, line} <- before do
+        assert line == Enum.find(@comments, &(&1.text == text)).line
+      end
+
+      # Simulate collapsing the `def` on lines 16-20
+      new_comments = displace_comments(@comments, 16..20)
+
+      expected = [
+        {"# Above def", 15},
+        {"# Before arg", 16},
+        {"# After arg", 16}
+      ]
+
+      for {text, line} <- expected do
+        assert line == Enum.find(new_comments, &(&1.text == text)).line
+      end
+    end
+  end
+
+  describe "shift_comments/3" do
+    test "Doesn't lose any comments" do
+      new_comments = shift_comments(@comments, 1..30, 1)
+
+      for comment <- @comments do
+        assert Enum.any?(new_comments, &(&1.text == comment.text))
+      end
+    end
+
+    test "Moves comments after the specified line by the specified delta" do
+      # Sanity-check line numbers in test fixture
+      before = [
+        {"# Above module", 1},
+        {"# Top of module", 3},
+        {"# Above def", 15},
+        {"# In function body", 21},
+        {"# After def", 25},
+        {"# After module", 27}
+      ]
+
+      for {text, line} <- before do
+        assert line == Enum.find(@comments, &(&1.text == text)).line
+      end
+
+      # Simulate collapsing the `def` on lines 16-20, shifting everything afterword up by 4
+      new_comments = shift_comments(@comments, 21..23, -4)
+
+      expected = [
+        # Before 21 doesn't get moved
+        {"# Above module", 1},
+        {"# Top of module", 3},
+        {"# Above def", 15},
+        # 21 does get moved
+        {"# In function body", 17},
+        # After 23 doesn't get moved
+        {"# After def", 25},
+        {"# After module", 27}
+      ]
+
+      for {text, line} <- expected do
+        assert line == Enum.find(new_comments, &(&1.text == text)).line
+      end
+    end
+  end
+end
diff --git a/test/support/style_case.ex b/test/support/style_case.ex
index 22af6b4..6aa7017 100644
--- a/test/support/style_case.ex
+++ b/test/support/style_case.ex
@@ -21,9 +21,9 @@ defmodule Styler.StyleCase do
     unless style, do: raise(ArgumentError, "missing required `:style` option")
 
     quote do
-      @style unquote(style)
       import unquote(__MODULE__), only: [assert_style: 1, assert_style: 2]
 
+      @style unquote(style)
       def style(code), do: unquote(__MODULE__).style(code, @style)
     end
   end
@@ -33,20 +33,23 @@ defmodule Styler.StyleCase do
 
     quote bind_quoted: [before: before, expected: expected] do
       expected = String.trim(expected)
-      {styled_ast, styled} = style(before)
+      {styled_ast, styled, styled_comments} = style(before)
 
       if styled != expected and ExUnit.configuration()[:trace] do
         IO.puts("\n======Given=============\n")
         IO.puts(before)
-        {before_ast, _} = Styler.string_to_quoted_with_comments(before)
+        {before_ast, before_comments} = Styler.string_to_quoted_with_comments(before)
         dbg(before_ast)
+        dbg(before_comments)
         IO.puts("======Expected==========\n")
         IO.puts(expected)
-        {expected_ast, _} = Styler.string_to_quoted_with_comments(expected)
+        {expected_ast, expected_comments} = Styler.string_to_quoted_with_comments(expected)
         dbg(expected_ast)
+        dbg(expected_comments)
         IO.puts("======Got===============\n")
         IO.puts(styled)
         dbg(styled_ast)
+        dbg(styled_comments)
         IO.puts("========================\n")
       end
 
@@ -64,6 +67,6 @@ defmodule Styler.StyleCase do
 
     styled_ast = Zipper.root(zipper)
     styled_code = Styler.quoted_to_string(styled_ast, comments)
-    {styled_ast, styled_code}
+    {styled_ast, styled_code, comments}
   end
 end
diff --git a/test/zipper_test.exs b/test/zipper_test.exs
index 84e4afa..aed5d32 100644
--- a/test/zipper_test.exs
+++ b/test/zipper_test.exs
@@ -45,26 +45,26 @@ defmodule StylerTest.ZipperTest do
     end
   end
 
-  describe "make_node/2" do
+  describe "replace_children/2" do
     test "2-tuples" do
-      assert Zipper.make_node({1, 2}, [3, 4]) == {3, 4}
+      assert Zipper.replace_children({1, 2}, [3, 4]) == {3, 4}
     end
 
     test "changing to 2-tuples arity" do
-      assert Zipper.make_node({1, 2}, [3, 4, 5]) == {:{}, [], [3, 4, 5]}
-      assert Zipper.make_node({1, 2}, [3]) == {:{}, [], [3]}
+      assert Zipper.replace_children({1, 2}, [3, 4, 5]) == {:{}, [], [3, 4, 5]}
+      assert Zipper.replace_children({1, 2}, [3]) == {:{}, [], [3]}
     end
 
     test "lists" do
-      assert Zipper.make_node([1, 2, 3], [:a, :b, :c]) == [:a, :b, :c]
+      assert Zipper.replace_children([1, 2, 3], [:a, :b, :c]) == [:a, :b, :c]
     end
 
     test "unqualified calls" do
-      assert Zipper.make_node({:foo, [], [1, 2]}, [:a, :b]) == {:foo, [], [:a, :b]}
+      assert Zipper.replace_children({:foo, [], [1, 2]}, [:a, :b]) == {:foo, [], [:a, :b]}
     end
 
     test "qualified calls" do
-      assert Zipper.make_node({{:., [], [1, 2]}, [], [3, 4]}, [:a, :b, :c]) == {:a, [], [:b, :c]}
+      assert Zipper.replace_children({{:., [], [1, 2]}, [], [3, 4]}, [:a, :b, :c]) == {:a, [], [:b, :c]}
     end
   end