Repository: adobe/elixir-styler (elixir)

Issue:
block manipulation is cruel to comments
the `if/unless/with` statement manipulation was all implemented under the SingleNode module, which ignores comments (.. because it's suppose to only manipulate single nodes). given that if and unless both swap blocks around without updating lines/comments, comments will end up in the wrong places


here's a tdd:

```ex
    test "moves comments" do
      assert_style(
        """
        if not :ok do
          # do
          do_body
        else
          # else
          else_body
        end
        """,
        """
        if :ok do
          # else
          else_body
        else
          # do
          do_body
        end
        """
      )
    end
```
fails with

```
  1) test if/else moves comments (Styler.Style.SingleNodeTest)
     test/style/single_node_test.exs:671
     Assertion with == failed
     code:  assert styled == expected
     left:  "if :ok do\n  # do\n  # else\n  else_body\nelse\n  do_body\nend"
     right: "if :ok do\n  # else\n  else_body\nelse\n  # do\n  do_body\nend"
     stacktrace:
       test/style/single_node_test.exs:672: (test)
```

a fix would move this multi-node manipulation to its own style (`Blocks`?) and update the line meta of all nodes in the do/else blocks, and likewise move comments within those ranges to their new homes. examples of manipulating line numbers and comment line numbers can be found in the `Defs` style

Pull request description:
maintain meta & comment lines when rewriting case/cond to if statements
Fixes #79 for `case` and `cond`. still need to address if/unless/with (with could be pretty tricky!)

Patch:
diff --git a/lib/style.ex b/lib/style.ex
index c6d6d90..09e81ef 100644
--- a/lib/style.ex
+++ b/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"
diff --git a/lib/style/blocks.ex b/lib/style/blocks.ex
index 7485d1d..02c145e 100644
--- a/lib/style/blocks.ex
+++ b/lib/style/blocks.ex
@@ -28,36 +28,27 @@ defmodule Styler.Style.Blocks do
   @behaviour Styler.Style
 
   alias Styler.Style
-
-  # @TODO handle comments https://github.com/adobe/elixir-styler/issues/79
-  def run({node, meta}, ctx), do: {:cont, {style(node), meta}, ctx}
-
-  defmacrop trivial_case(head, a, a_body, b, b_body) do
-    quote do
-      {:case, _,
-       [
-         unquote(head),
-         [
-           {_,
-            [
-              {:->, _, [[unquote(a)], unquote(a_body)]},
-              {:->, _, [[unquote(b)], unquote(b_body)]}
-            ]}
-         ]
-       ]}
+  alias Styler.Zipper
+
+  # case statement with exactly 2 `->` cases
+  # rewrite to `if` if it's any of 3 trivial cases
+  def run({{:case, _, [head, [{_, [{:->, _, [[lhs_a], a]}, {:->, _, [[lhs_b], b]}]}]]}, zm} = zipper, ctx) do
+    case {lhs_a, lhs_b} do
+      {{_, _, [true]}, {_, _, [false]}} -> if_ast(head, a, b, ctx, zm)
+      {{_, _, [true]}, {:_, _, _}} -> if_ast(head, a, b, ctx, zm)
+      {{_, _, [false]}, {_, _, [true]}} -> if_ast(head, b, a, ctx, zm)
+      _ -> {:cont, zipper, ctx}
     end
   end
 
-  defp style(trivial_case(head, {_, _, [true]}, do_, {_, _, [false]}, else_)), do: styled_if(head, do_, else_)
-  defp style(trivial_case(head, {_, _, [false]}, else_, {_, _, [true]}, do_)), do: styled_if(head, do_, else_)
-  defp style(trivial_case(head, {_, _, [true]}, do_, {:_, _, _}, else_)), do: styled_if(head, do_, else_)
-
   # `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
-    styled_if(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`
@@ -138,11 +129,57 @@ defmodule Styler.Style.Blocks do
     Style.update_all_meta(a, fn _ -> nil end) == Style.update_all_meta(b, fn _ -> nil end)
   end
 
-  defp styled_if(head, do_body, else_body) do
-    {_, meta, _} = head
+  defp if_ast({_, meta, _} = head, do_body, else_body, ctx, zipper_meta) do
+    comments = ctx.comments
     line = meta[:line]
-    # @TODO figure out appropriate line meta for `else` and `if->end->line`
-    children = [head, [{{:__block__, [line: line], [:do]}, do_body}, {{:__block__, [], [:else]}, else_body}]]
-    style({:if, [line: line, do: [line: line], end: []], children})
+    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)
+    end_line = max(max_do_line, max_else_line)
+
+    # Change ast meta and comment lines to fit the `if` ast
+    {do_block, else_block, comments} =
+      if max_do_line >= max_else_line do
+        # 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
+
+        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}
+        ]
+
+        do_block = {{:__block__, [line: line], [:do]}, Style.shift_line(do_body, -else_size)}
+        else_block = {{:__block__, [line: max_else_line], [:else]}, Style.shift_line(else_body, else_size + 1)}
+        {do_block, else_block, Style.shift_comments(comments, shifts)}
+      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}
+        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, comments}
+      end
+
+    # end line is max_line + 2 to accomodate `else` and `end` both having their own line
+    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) do
+    {_, max_line} =
+      ast
+      |> Zipper.zip()
+      |> Zipper.traverse(0, fn
+        {{_, meta, _}, _} = z, max -> {z, max(meta[:line] || max, max)}
+        z, max -> {z, max}
+      end)
+
+    max_line
   end
 end


Test patch:
diff --git a/test/style/blocks_test.exs b/test/style/blocks_test.exs
index d38b0c0..e34c693 100644
--- a/test/style/blocks_test.exs
+++ b/test/style/blocks_test.exs
@@ -16,14 +16,18 @@ defmodule Styler.Style.BlocksTest 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,16 +56,19 @@ defmodule Styler.Style.BlocksTest do
       assert_style(
         """
         case foo do
-          false -> :error
+          # a
           true -> :ok
+          # b
+          false -> nil
         end
         """,
         """
         if foo do
+          # a
           :ok
-        else
-          :error
         end
+
+        # b
         """
       )
 
@@ -65,7 +76,7 @@ defmodule Styler.Style.BlocksTest do
         """
         case foo do
           true -> :ok
-          false -> nil
+          _ -> nil
         end
         """,
         """
@@ -79,31 +90,149 @@ defmodule Styler.Style.BlocksTest do
         """
         case foo do
           true -> :ok
-          _ -> nil
+          false ->
+            Logger.warning("it's false")
+            nil
         end
         """,
         """
         if foo do
           :ok
+        else
+          Logger.warning("it's false")
+          nil
         end
         """
       )
+    end
 
+    test "block swapping comments" do
       assert_style(
         """
         case foo do
+          false ->
+            # a
+            :error
+          true ->
+            # b
+            :ok
+        end
+        """,
+        """
+        if foo do
+          # b
+          :ok
+        else
+          # a
+          :error
+        end
+        """
+      )
+
+      assert_style(
+        """
+        case foo do
+          # a
+          false ->
+            :error
+          # b
+          true ->
+            :ok
+        end
+        """,
+        """
+        if foo do
+          # b
+          :ok
+        else
+          # a
+          :error
+        end
+        """
+      )
+
+      assert_style(
+        """
+        case foo do
+          # a
+          false -> :error
+          # b
           true -> :ok
+        end
+        """,
+        """
+        if foo do
+          # b
+          :ok
+        else
+          # a
+          :error
+        end
+        """
+      )
+    end
+
+    test "complex comments" do
+      assert_style(
+        """
+        case foo do
           false ->
+            #a
+            actual(code)
+
+            #b
+            if foo do
+              #c
+              doing_stuff()
+              #d
+            end
+
+            #e
+            :ok
+          true ->
+            #f
             Logger.warning("it's false")
+
+            if 1 do
+              # g
+              :yay
+            else
+              # h
+              :ohno
+            end
+
+            # i
             nil
         end
         """,
         """
         if foo do
-          :ok
-        else
+          # f
           Logger.warning("it's false")
+
+          if 1 do
+            # g
+            :yay
+          else
+            # h
+            :ohno
+          end
+
+          # i
           nil
+        else
+          # a
+          actual(code)
+
+          # b
+          if foo do
+            # c
+            doing_stuff()
+            # d
+          end
+
+          # e
+          :ok
         end
         """
       )