Preserve formatting for blocks

lib/sourceror/comments.ex

@@ -23,8 +23,31 @@ defmodule Sourceror.Comments do
         quoted
 
       _ ->
-        line = Sourceror.get_line(quoted)
-        {:__block__, [trailing_comments: leftovers, leading_comments: [], line: line], [quoted]}
+        if match?({:__block__, _, _}, quoted) and not Sourceror.Identifier.do_block?(quoted) do
+          {last, args} = Sourceror.get_args(quoted) |> List.pop_at(-1)
+          line = Sourceror.get_line(last)
+
+          last =
+            {:__block__,
+             [
+               __sourceror__: %{trailing_block: true},
+               trailing_comments: leftovers,
+               leading_comments: [],
+               line: line
+             ], [last]}
+
+          {:__block__, Sourceror.get_meta(quoted), args ++ [last]}
+        else
+          line = Sourceror.get_line(quoted)
+
+          {:__block__,
+           [
+             __sourceror__: %{trailing_block: true},
+             trailing_comments: leftovers,
+             leading_comments: [],
+             line: line
+           ], [quoted]}
+        end
     end
   end
 
@@ -149,6 +172,14 @@ defmodule Sourceror.Comments do
       Enum.concat([acc, leading_comments, trailing_comments])
       |> Enum.sort_by(& &1.line)
 
+    quoted =
+      if meta[:__sourceror__][:trailing_block] do
+        {_, _, [quoted]} = quoted
+        quoted
+      else
+        quoted
+      end
+
     quoted =
       Macro.update_meta(quoted, fn meta ->
         meta
@@ -161,26 +192,29 @@ defmodule Sourceror.Comments do
 
   defp collapse_trailing_comments(quoted, trailing_comments) do
     meta = Sourceror.get_meta(quoted)
+    trailing_block? = meta[:__sourceror__][:trailing_block]
 
     comments =
       Enum.map(trailing_comments, fn comment ->
         line = meta[:end_of_expression][:line] || meta[:line]
 
-        %{comment | line: line - 2, previous_eol_count: 1}
+        %{comment | line: line - 1}
       end)
 
     comments =
       case comments do
         [first | rest] ->
-          [%{first | previous_eol_count: 0} | rest]
+          prev_eol_count = if trailing_block?, do: first.previous_eol_count, else: 0
+
+          [%{first | previous_eol_count: prev_eol_count} | rest]
 
         _ ->
           comments
       end
 
     case List.pop_at(comments, -1) do
       {last, rest} when is_map(last) ->
-        rest ++ [%{last | next_eol_count: 2}]
+        rest ++ [%{last | next_eol_count: 0}]
 
       _ ->
         comments

lib/sourceror/identifier.ex

@@ -115,4 +115,11 @@ defmodule Sourceror.Identifier do
   defp trim_leading_while_valid_identifier(other) do
     other
   end
+
+  @doc false
+  def do_block?({:__block__, _, args}) do
+    not is_nil(args[:do])
+  end
+
+  def do_block?(_), do: false
 end

lib/sourceror/lines_corrector.ex

@@ -75,7 +75,11 @@ defmodule Sourceror.LinesCorrector do
 
     last_line =
       if has_trailing_comments?(quoted) do
-        last_line + length(meta[:trailing_comments] || []) + 2
+        if Sourceror.Identifier.do_block?(quoted) do
+          last_line + length(meta[:trailing_comments] || []) + 2
+        else
+          last_line + length(meta[:trailing_comments] || []) + 1
+        end
       else
         last_line
       end

test/comments_test.exs

@@ -142,7 +142,7 @@ defmodule SourcerorTest.CommentsTest do
 
       assert [
                %{line: 5, text: "# A"},
-               %{line: 8, text: "# B"}
+               %{line: 7, text: "# B"}
              ] = comments
 
       quoted =
@@ -158,7 +158,7 @@ defmodule SourcerorTest.CommentsTest do
 
       assert [
                %{line: 5, text: "# A"},
-               %{line: 8, text: "# B"}
+               %{line: 7, text: "# B"}
              ] = comments
 
       assert Sourceror.to_string(quoted, collapse_comments: true, correct_lines: true) ==

test/sourceror_test.exs

@@ -14,6 +14,13 @@ defmodule SourcerorTest do
       assert_same(~S"""
       foo()
 
+      baz()
+
+      # Bar
+      """)
+
+      assert_same(~S"""
+      foo()
       # Bar
       """)
     end