fix handling of comments on binary operators with assoc rewriting (#264)

WhatsApp · Mar 29, 2021 · 309babd · 309babd
1 parent 1d56a5e
commit 309babd
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 10 deletions.
diff --git a/src/erlfmt_format.erl b/src/erlfmt_format.erl
@@ -312,13 +312,17 @@ binary_op_to_algebra(Op, Meta0, Left0, Right0) ->
     {op, _Meta, Op, Left, Right} = rewrite_assoc(Op, Meta, Left0, Right0),
     binary_op_to_algebra(Op, Meta, Left, Right, ?INDENT).
 
-rewrite_assoc(Op, MetaA, A, {op, MetaBC, Op, B, C} = BC) ->
-    case has_no_comments_or_parens(MetaBC) of
-        false ->
-            {op, MetaA, Op, A, BC};
+rewrite_assoc(Op, MetaABC0, A, {op, MetaBC0, Op, B0, C} = BC) ->
+    case erlfmt_scan:get_anno(parens, MetaBC0, false) of
+        true ->
+            {op, MetaABC0, Op, A, BC};
         _ ->
+            {Pre, Post} = comments(MetaBC0),
+            MetaBC = erlfmt_scan:delete_annos([pre_comments, post_comments], MetaBC0),
+            B = erlfmt_recomment:put_pre_comments(B0, Pre),
             AB = rewrite_assoc(Op, update_meta_location(MetaBC, A, B), A, B),
-            rewrite_assoc(Op, update_meta_location(MetaA, AB, C), AB, C)
+            MetaABC = erlfmt_recomment:put_post_comments(MetaABC0, Post),
+            rewrite_assoc(Op, update_meta_location(MetaABC, AB, C), AB, C)
     end;
 rewrite_assoc(Op, Meta, Left, Right) ->
     {op, Meta, Op, Left, Right}.

diff --git a/test/erlfmt_format_SUITE.erl b/test/erlfmt_format_SUITE.erl
@@ -754,7 +754,14 @@ binary_operator_more(Config) when is_list(Config) ->
         "Foo orelse\n"
         "    %% orelse comment\n"
         "    Bar orelse\n"
-        "        Baz\n"
+        "    Baz\n"
+    ),
+    ?assertSame(
+        "Foo orelse\n"
+        "    %% orelse comment\n"
+        "    Bar orelse\n"
+        "    %% orelse comment\n"
+        "    Baz\n"
     ),
     ?assertFormat(
         "equivalent(L1, R1) andalso %% checking L1 and R1\n"
@@ -765,10 +772,10 @@ binary_operator_more(Config) when is_list(Config) ->
         "equivalent(L1, R1) andalso\n"
         "    %% checking L2 and R2\n"
         "    equivalent(L2, R2) andalso\n"
-        "        %% checking L3 and R3\n"
-        "        equivalent(L3, R3) andalso\n"
-        "            %% checking L4 and R4\n"
-        "            equivalent(L4, R4)\n"
+        "    %% checking L3 and R3\n"
+        "    equivalent(L3, R3) andalso\n"
+        "    %% checking L4 and R4\n"
+        "    equivalent(L4, R4)\n"
     ).
 
 tuple(Config) when is_list(Config) ->
@@ -3372,6 +3379,11 @@ comment(Config) when is_list(Config) ->
         "%% post comment\n"
         "b.\n"
     ),
+    ?assertSame(
+        "a\n"
+        "%% post comment\n"
+        ".\n"
+    ),
     ?assertFormat(
         "\"a,\n"
         "b\" % c\n"