From 99d9e943fb8b03e81b231086aad33de7675ad11e Mon Sep 17 00:00:00 2001
From: Matt Enlow <enlow@adobe.com>
Date: Thu, 7 Dec 2023 14:50:43 -0700
Subject: [PATCH] maintain meta & comment lines when rewriting case/cond to if
 statements (#97)

Closes #79
---
 lib/style.ex               |  19 ++++-
 lib/style/blocks.ex        |  97 +++++++++++++++++--------
 test/style/blocks_test.exs | 143 +++++++++++++++++++++++++++++++++++--
 3 files changed, 219 insertions(+), 40 deletions(-)

diff --git a/lib/style.ex b/lib/style.ex
index c6d6d90b..09e81ef7 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 7485d1d5..02c145e5 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
diff --git a/test/style/blocks_test.exs b/test/style/blocks_test.exs
index d38b0c01..e34c693a 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
         """
       )