Preserve comments in receive after after (#271)

* add TODO for broken receive test that shows how comment is moved * getting there * always break when comment is present * an alternative * better * added concat test * comments do not indent * remove superfluous pattern match * maybe this is better * add a test * special case comments * preserve breaks * reformat code
WhatsApp · Apr 19, 2021 · e532852 · e532852
1 parent c909c32
commit e532852
Show file tree

Hide file tree

Showing 6 changed files with 154 additions and 34 deletions.
diff --git a/src/erlfmt_format.erl b/src/erlfmt_format.erl
@@ -189,13 +189,15 @@ do_expr_to_algebra({'case', _Meta, Expr, Clauses}) ->
     Prefix = surround(<<"case">>, <<" ">>, expr_to_algebra(Expr), <<" ">>, <<"of">>),
     surround_block(Prefix, clauses_to_algebra(Clauses), <<"end">>);
 do_expr_to_algebra({'receive', _Meta, Clauses}) ->
-    surround_block(<<"receive">>, clauses_to_algebra(Clauses), <<"end">>);
-do_expr_to_algebra({'receive', _Meta, [], AfterExpr, AfterBody}) ->
-    AfterD = receive_after_to_algebra(AfterExpr, AfterBody),
+    surround_block(<<"receive">>, expr_to_algebra(Clauses), <<"end">>);
+do_expr_to_algebra({'receive', _Meta, empty, After}) ->
+    AfterD = expr_to_algebra(After),
     concat(force_breaks(), (line(<<"receive">>, line(AfterD, <<"end">>))));
-do_expr_to_algebra({'receive', _Meta, Clauses, AfterExpr, AfterBody}) ->
-    AfterD = receive_after_to_algebra(AfterExpr, AfterBody),
-    surround_block(<<"receive">>, clauses_to_algebra(Clauses), line(AfterD, <<"end">>));
+do_expr_to_algebra({'receive', _Meta, Clauses, After}) ->
+    AfterD = expr_to_algebra(After),
+    surround_block(<<"receive">>, expr_to_algebra(Clauses), line(AfterD, <<"end">>));
+do_expr_to_algebra({after_clause, _Meta, Expr, Body}) ->
+    receive_after_to_algebra(Expr, Body);
 do_expr_to_algebra({'try', _Meta, Exprs, OfClauses, CatchClauses, After}) ->
     try_to_algebra(Exprs, OfClauses, CatchClauses, After);
 do_expr_to_algebra({'catch', _Meta, Exprs}) ->
@@ -776,12 +778,17 @@ single_clause_spec_to_algebra(Name, {spec_clause, CMeta, Head, Body, Guards}) ->
     {args, AMeta, Args} = Head,
     clauses_to_algebra([{spec_clause, CMeta, {call, AMeta, Name, Args}, Body, Guards}]).
 
-receive_after_to_algebra(Expr, [HBody | _] = Body) ->
-    ExprD = do_expr_to_algebra(Expr),
+receive_after_to_algebra(Expr, Body) ->
+    ExprD = expr_to_algebra(Expr),
     BodyD = block_to_algebra(Body),
-    HeadD = concat([<<"after ">>, ExprD, <<" ->">>]),
-    Doc = concat([HeadD, break(), maybe_force_breaks(has_break_between(Expr, HBody)), BodyD]),
-    combine_comments(element(2, Expr), group(nest(Doc, ?INDENT))).
+    Nest = fun(List) -> group(nest(concat(break(), concat(List)), ?INDENT, break)) end,
+    case has_comments(Expr) of
+        false ->
+            MaybeForceBreaks = maybe_force_breaks(has_break_between(Expr, Body)),
+            concat([<<"after">>, <<" ">>, ExprD, <<" ->">>, Nest([MaybeForceBreaks, BodyD])]);
+        true ->
+            concat([<<"after">>, Nest([ExprD, <<" ->">>, Nest([BodyD])])])
+    end.
 
 try_to_algebra(Body, OfClauses, CatchClauses, After) ->
     Clauses =
@@ -900,3 +907,9 @@ comments_with_pre_dot(Meta) ->
 comments(Meta) ->
     [] = erlfmt_scan:get_anno(pre_dot_comments, Meta, []),
     {erlfmt_scan:get_anno(pre_comments, Meta, []), erlfmt_scan:get_anno(post_comments, Meta, [])}.
+
+has_comments(Expr) ->
+    case comments(Expr) of
+        {[], []} -> false;
+        _ -> true
+    end.
diff --git a/src/erlfmt_parse.yrl b/src/erlfmt_parse.yrl
@@ -397,11 +397,15 @@ cr_clause -> expr clause_guard clause_body :
     {clause, ?range_anno('$1', '$3'), '$1', '$2', ?val('$3')}.
 
 receive_expr -> 'receive' cr_clauses 'end' :
-        {'receive',?range_anno('$1', '$3'),'$2'}.
-receive_expr -> 'receive' 'after' expr clause_body 'end' :
-        {'receive',?range_anno('$1', '$5'),[],'$3',?val('$4')}.
-receive_expr -> 'receive' cr_clauses 'after' expr clause_body 'end' :
-        {'receive',?range_anno('$1', '$6'),'$2','$4',?val('$5')}.
+        Clauses = {clauses, ?range_anno('$1', '$3'), '$2'},
+        {'receive',?range_anno('$1', '$3'),Clauses}.
+receive_expr -> 'receive' 'after' expr clause_body 'end':
+        After =  {after_clause, ?range_anno('$2', '$5'), '$3', ?val('$4')},
+        {'receive',?range_anno('$1', '$5'),empty,After}.
+receive_expr -> 'receive' cr_clauses 'after' expr clause_body 'end':
+        After =  {after_clause, ?range_anno('$3', '$6'), '$4', ?val('$5')},
+        Clauses = {clauses, ?range_upto_anno('$1', '$3'), '$2'},
+        {'receive',?range_anno('$1', '$6'),Clauses,After}.
 
 fun_expr -> 'fun' atom_or_var_or_macro '/' integer_or_var_or_macro :
     Anno = ?range_anno('$1', '$4'),

diff --git a/src/erlfmt_recomment.erl b/src/erlfmt_recomment.erl
@@ -217,13 +217,19 @@ insert_nested({'case', Meta, Expr0, Clauses0}, Comments0) ->
     Clauses = insert_expr_container(Clauses0, Comments1),
     {{'case', Meta, Expr, Clauses}, []};
 insert_nested({'receive', Meta, Clauses0}, Comments0) ->
-    Clauses = insert_expr_container(Clauses0, Comments0),
+    {Clauses, []} = insert_expr(Clauses0, Comments0),
     {{'receive', Meta, Clauses}, []};
-insert_nested({'receive', Meta, Clauses0, AfterExpr0, AfterBody0}, Comments0) ->
-    {Clauses, Comments1} = insert_expr_list(Clauses0, Comments0),
-    {AfterExpr, Comments2} = insert_expr(AfterExpr0, Comments1),
-    AfterBody = insert_expr_container(AfterBody0, Comments2),
-    {{'receive', Meta, Clauses, AfterExpr, AfterBody}, []};
+insert_nested({'receive', Meta, empty, After0}, Comments0) ->
+    {After, []} = insert_expr(After0, Comments0),
+    {{'receive', Meta, empty, After}, []};
+insert_nested({'receive', Meta, Clauses0, After0}, Comments0) ->
+    {Clauses, Comments1} = insert_expr(Clauses0, Comments0),
+    {After, []} = insert_expr(After0, Comments1),
+    {{'receive', Meta, Clauses, After}, []};
+insert_nested({after_clause, Meta, Expr0, Body0}, Comments0) ->
+    {Expr, Comments1} = insert_expr(Expr0, Comments0),
+    Body = insert_expr_container(Body0, Comments1),
+    {{after_clause, Meta, Expr, Body}, []};
 insert_nested({'if', Meta, Clauses0}, Comments0) ->
     Clauses = insert_expr_container(Clauses0, Comments0),
     {{'if', Meta, Clauses}, []};

diff --git a/test/erlfmt_SUITE.erl b/test/erlfmt_SUITE.erl
@@ -730,7 +730,7 @@ clauses(Config) when is_list(Config) ->
         parse_expr("case X of true -> ok end")
     ),
     ?assertMatch(
-        {'receive', _, [{clause, _, {var, _, '_'}, empty, [{atom, _, true}]}]},
+        {'receive', _, {clauses, _, [{clause, _, {var, _, '_'}, empty, [{atom, _, true}]}]}},
         parse_expr("receive _ -> true end")
     ),
     ?assertMatch(

diff --git a/test/erlfmt_SUITE_data/simple_comments.erl b/test/erlfmt_SUITE_data/simple_comments.erl
@@ -160,7 +160,7 @@ call() ->
     end,
     receive
         ok -> ok
-    %% comment 5
+        %% comment 5
     after 1 ->
         %% comment 6
         ok

diff --git a/test/erlfmt_format_SUITE.erl b/test/erlfmt_format_SUITE.erl
@@ -2480,7 +2480,9 @@ receive_expression(Config) when is_list(Config) ->
     ?assertSame(
         "receive\n"
         "after 0 ->\n"
-        "    some:long(Expression)\n"
+        "    some:long(\n"
+        "        Expression\n"
+        "    )\n"
         "end\n",
         25
     ),
@@ -2574,13 +2576,46 @@ receive_expression(Config) when is_list(Config) ->
         "    2 -> two\n"
         "end\n"
     ),
-    % ?assertSame(
-    %     "receive\n"
-    %     "after\n"
-    %     "    % foo\n"
-    %     "    0 -> ok\n"
-    %     "end\n"
-    % ),
+    ?assertFormat(
+        "receive\n"
+        "after\n"
+        "    0 -> ok\n"
+        "    % receive post comment\n"
+        "end\n",
+        "receive\n"
+        "after 0 ->\n"
+        "    ok\n"
+        "    % receive post comment\n"
+        "end\n"
+    ),
+    ?assertFormat(
+        "receive\n"
+        "after\n"
+        "    % before zero\n"
+        "    0 -> ok\n"
+        "end\n",
+        "receive\n"
+        "after\n"
+        "    % before zero\n"
+        "    0 -> ok\n"
+        "end\n"
+    ),
+    ?assertFormat(
+        "receive\n"
+        "after\n"
+        "    %% foo\n"
+        "    0 ->\n"
+        "        \"abc\"\n"
+        "        \"def\"\n"
+        "end.\n",
+        "receive\n"
+        "after\n"
+        "    %% foo\n"
+        "    0 ->\n"
+        "        \"abc\"\n"
+        "        \"def\"\n"
+        "end.\n"
+    ),
     ?assertFormat(
         "receive\n"
         "    1 -> ok\n"
@@ -2591,12 +2626,74 @@ receive_expression(Config) when is_list(Config) ->
         "end\n",
         "receive\n"
         "    1 -> ok\n"
-        "%% after receive\n"
-        "\n"
+        "    %% after receive\n"
         "after 0 ->\n"
         "    ok\n"
         "    %% after after for receive\n"
         "end\n"
+    ),
+    ?assertFormat(
+        "receive\n"
+        "    1 -> ok\n"
+        "    %% after receive\n"
+        "after\n"
+        "    %% before zero\n"
+        "    0 -> ok\n"
+        "    %% after after for receive\n"
+        "end\n",
+        "receive\n"
+        "    1 -> ok\n"
+        "    %% after receive\n"
+        "after\n"
+        "    %% before zero\n"
+        "    0 ->\n"
+        "        ok\n"
+        "        %% after after for receive\n"
+        "end\n"
+    ),
+    ?assertFormat(
+        "receive\n"
+        "    1 -> ok\n"
+        "after 0 -> % after arrow moves one line above\n"
+        "    a(\n"
+        "        b,\n"
+        "        c\n"
+        "    )\n"
+        "end\n",
+        "receive\n"
+        "    1 -> ok\n"
+        "    % after arrow moves one line above\n"
+        "after 0 ->\n"
+        "    a(\n"
+        "        b,\n"
+        "        c\n"
+        "    )\n"
+        "end\n"
+    ),
+    ?assertFormat(
+        "receive\n"
+        "    1 -> ok\n"
+        "    %% after receive\n"
+        "after\n"
+        "    %% before zero\n"
+        "    0 -> a(\n"
+        "        b,\n"
+        "        c\n"
+        "    )\n"
+        "    %% after after for receive\n"
+        "end\n",
+        "receive\n"
+        "    1 -> ok\n"
+        "    %% after receive\n"
+        "after\n"
+        "    %% before zero\n"
+        "    0 ->\n"
+        "        a(\n"
+        "            b,\n"
+        "            c\n"
+        "        )\n"
+        "        %% after after for receive\n"
+        "end\n"
     ).
 
 try_expression(Config) when is_list(Config) ->