shift comments when making if ast

lib/style.ex

@@ -19,7 +19,7 @@ defmodule Styler.Style do
   alias Styler.Zipper
 
   @type context :: %{
-          comment: [map()],
+          comments: [map()],
           file: :stdin | String.t()
         }
 
@@ -112,13 +112,26 @@ defmodule Styler.Style do
   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
+    shift_comments(comments, [{range, delta}])
+  end
+
+  @doc """
+  Perform a series of shifts in a single pass.
+
+  When shifting comments from block A to block B, naively using two passes of `shift_comments/3` would result
+  in all comments ending up in either region A or region B (because A would move to B, then all B back to A)
+  This function exists to make sure that a comment is only moved once during the swap.
+  """
+  def shift_comments(comments, shifts) do
+    comments
+    |> Enum.map(fn comment ->
+      if delta = Enum.find_value(shifts, fn {range, delta} -> comment.line in range && delta end) do
         %{comment | line: comment.line + delta}
       else
         comment
       end
     end)
+    |> Enum.sort_by(& &1.line)
   end
 
   @doc "Returns true if the ast represents an empty map"

lib/style/blocks.ex

@@ -30,26 +30,25 @@ defmodule Styler.Style.Blocks do
   alias Styler.Style
   alias Styler.Zipper
 
-  # @TODO handle comments https://github.com/adobe/elixir-styler/issues/79
-  def run({node, meta}, ctx), do: {:cont, {style(node), meta}, ctx}
-
   # case statement with exactly 2 `->` cases
   # rewrite to `if` if it's any of 3 trivial cases
-  defp style({:case, _, [head, [{_, [{:->, _, [[lhs_a], a]}, {:->, _, [[lhs_b], b]}]}]]} = trivial_case_ast) do
+  def run({{:case, _, [head, [{_, [{:->, _, [[lhs_a], a]}, {:->, _, [[lhs_b], b]}]}]]}, m} = zipper, ctx) do
     case {lhs_a, lhs_b} do
-      {{_, _, [true]}, {_, _, [false]}} -> if_ast(head, a, b)
-      {{_, _, [true]}, {:_, _, _}} -> if_ast(head, a, b)
-      {{_, _, [false]}, {_, _, [true]}} -> if_ast(head, b, a)
-      _ -> trivial_case_ast
+      {{_, _, [true]}, {_, _, [false]}} -> if_ast(head, a, b, ctx, m)
+      {{_, _, [true]}, {:_, _, _}} -> if_ast(head, a, b, ctx, m)
+      {{_, _, [false]}, {_, _, [true]}} -> if_ast(head, b, a, ctx, m)
+      _ -> {:cont, zipper, ctx}
     end
   end
 
   # `Credo.Check.Refactor.CondStatements`
   # This also detects strings and lists...
-  defp style({:cond, _, [[{_, [{:->, _, [[expr], do_body]}, {:->, _, [[{:__block__, _, [truthy]}], else_body]}]}]]})
-       when is_atom(truthy) and truthy not in [nil, false] do
-    if_ast(expr, do_body, else_body)
-  end
+  def run({{:cond, _, [[{_, [{:->, _, [[head], a]}, {:->, _, [[{:__block__, _, [truthy]}], b]}]}]]}, m}, ctx)
+      when is_atom(truthy) and truthy not in [nil, false],
+      do: if_ast(head, a, b, ctx, m)
+
+  # @TODO handle comments https://github.com/adobe/elixir-styler/issues/79
+  def run({node, meta}, ctx), do: {:cont, {style(node), meta}, ctx}
 
   # Credo.Check.Readability.WithSingleClause
   # rewrite `with success <- single_statement do body else ...elses end`
@@ -130,38 +129,55 @@ defmodule Styler.Style.Blocks do
     Style.update_all_meta(a, fn _ -> nil end) == Style.update_all_meta(b, fn _ -> nil end)
   end
 
-  defp if_ast({_, meta, _} = head, do_body, else_body) do
+  defp if_ast({_, meta, _} = head, do_body, else_body, ctx, zipper_meta) do
+    comments = ctx.comments
     line = meta[:line]
     do_body = Macro.update_meta(do_body, &Keyword.delete(&1, :end_of_expression))
     else_body = Macro.update_meta(else_body, &Keyword.delete(&1, :end_of_expression))
 
     max_do_line = max_line(do_body)
     max_else_line = max_line(else_body)
 
-    #@TODO shift comments with Style.shift_line calls
-    {do_block, else_block, end_line} =
+    # Change ast meta and comment lines to fit the `if` ast
+    #
+    # all +1s in the following code are to accomodate the `else` keyword adding a line to the code.
+    # +2s accomodates both the `else` keyword and the `end` keyword
+    {do_block, else_block, end_line, comments} =
       if max_do_line >= max_else_line do
-        shift = max_else_line - line
-        dbg(shift)
-        do_block = {{:__block__, [line: line], [:do]}, Style.shift_line(do_body, -shift)}
-        # +1 to fit the else keyword
-        else_body = Style.shift_line(else_body, shift + 1)
-        else_block = {{:__block__, [line: max_else_line + 1], [:else]}, else_body}
-        {do_block, else_block, max_do_line + 2}
+        # we're swapping the ordering of two blocks of code
+        # and so must swap the lines of the ast & comments to keep comments where they belong!
+        # the math is: move B up by the length of A, and move A down by the length of B plus one (for the else keyword)
+        else_size = max_else_line - line
+        do_size = max_do_line - max_else_line + 1
+
+        shifts = [
+          # move comments in the `else_body`` down by the size of the `do_body`
+          {line..max_else_line, do_size},
+          # move comments in `do_body` up by the size of the `else_body`
+          {(max_else_line + 1)..max_do_line, -else_size}
+        ]
+
+        comments = Style.shift_comments(comments, shifts)
+
+        do_block = {{:__block__, [line: line], [:do]}, Style.shift_line(do_body, -else_size)}
+        else_block = {{:__block__, [line: max_else_line + 1], [:else]}, Style.shift_line(else_body, else_size + 1)}
+        {do_block, else_block, max_do_line + 2, comments}
       else
+        # much simpler case -- just scootch things in the else down by 1 for the `else` keyword.
         do_block = {{:__block__, [line: line], [:do]}, do_body}
-        # move everything down by 1 to put the `else` keyword on its own line
+        comments = Style.shift_comments(comments, max_do_line..max_else_line, 1)
         else_block = Style.shift_line({{:__block__, [line: max_do_line], [:else]}, else_body}, 1)
-        {do_block, else_block, max_else_line + 2}
+        {do_block, else_block, max_else_line + 2, comments}
       end
-    dbg()
 
     # end line is max_line + 2 to accomodate `else` and `end` both having their own line
-    style({:if, [do: [line: line], end: [line: end_line], line: line], [head, [do_block, else_block]]})
+    if_ast = style({:if, [do: [line: line], end: [line: end_line], line: line], [head, [do_block, else_block]]})
+    {:cont, {if_ast, zipper_meta}, %{ctx | comments: comments}}
   end
 
   defp max_line(ast, start \\ 0) do
-    {_, max_line} =  ast
+    {_, max_line} =
+      ast
       |> Zipper.zip()
       |> Zipper.traverse(start, fn
         {{_, meta, _}, _} = z, max -> {z, max(meta[:line] || max, max)}

test/style/blocks_test.exs

@@ -9,21 +9,25 @@
 # governing permissions and limitations under the License.
 
 defmodule Styler.Style.BlocksTest do
-  use Styler.StyleCase, strict: true, async: true
+  use Styler.StyleCase, async: true
 
   describe "case to if" do
     test "rewrites case true false to if else" do
       assert_style(
         """
         case foo do
+          # a
           true -> :ok
+          # b
           false -> :error
         end
         """,
         """
         if foo do
+          # a
           :ok
         else
+          # b
           :error
         end
         """
@@ -32,14 +36,18 @@ defmodule Styler.Style.BlocksTest do
       assert_style(
         """
         case foo do
+          # a
           true -> :ok
+          # b
           _ -> :error
         end
         """,
         """
         if foo do
+          # a
           :ok
         else
+          # b
           :error
         end
         """
@@ -48,14 +56,18 @@ defmodule Styler.Style.BlocksTest do
       assert_style(
         """
         case foo do
+          # a
           false -> :error
+          # b
           true -> :ok
         end
         """,
         """
         if foo do
+          # b
           :ok
         else
+          # a
           :error
         end
         """
@@ -64,13 +76,17 @@ defmodule Styler.Style.BlocksTest do
       assert_style(
         """
         case foo do
+          # a
           true -> :ok
+          # b
           false -> nil
         end
         """,
         """
         if foo do
+          # a
           :ok
+          # b
         end
         """
       )