Support map comprehensions

WhatsApp · Sep 5, 2023 · 0e0e907 · 0e0e907
1 parent 7f1bed6
commit 0e0e907
Show file tree

Hide file tree

Showing 5 changed files with 94 additions and 2 deletions.
diff --git a/src/erlfmt_format.erl b/src/erlfmt_format.erl
@@ -169,6 +169,8 @@ do_expr_to_algebra({record_name, Meta, Name}) ->
     record_name_to_algebra(Meta, Name);
 do_expr_to_algebra({lc, _Meta, Expr, LcExprs}) ->
     comprehension_to_algebra(Expr, LcExprs, <<"[">>, <<"]">>);
+do_expr_to_algebra({mc, _Meta, Expr, LcExprs}) ->
+    comprehension_to_algebra(Expr, LcExprs, <<"#{">>, <<"}">>);
 do_expr_to_algebra({bc, _Meta, Expr, LcExprs}) ->
     comprehension_to_algebra(Expr, LcExprs, <<"<<">>, <<">>">>);
 do_expr_to_algebra({generate, _Meta, Left, Right}) ->

diff --git a/src/erlfmt_parse.yrl b/src/erlfmt_parse.yrl
@@ -26,6 +26,7 @@ pat_expr pat_expr_max map_pat_expr record_pat_expr
 pat_argument_list pat_exprs
 list list_exprs
 list_comprehension lc_expr lc_exprs
+map_comprehension
 binary_comprehension
 tuple
 record_expr record_name record_field_name record_tuple record_field record_fields
@@ -238,6 +239,7 @@ expr_max -> atomic : '$1'.
 expr_max -> list : '$1'.
 expr_max -> binary : '$1'.
 expr_max -> list_comprehension : '$1'.
+expr_max -> map_comprehension : '$1'.
 expr_max -> binary_comprehension : '$1'.
 expr_max -> tuple : '$1'.
 expr_max -> '(' expr ')' : set_parens('$2').
@@ -326,6 +328,8 @@ dotted_seq -> float '.' dotted_seq : ['$1' | '$3'].
 
 list_comprehension -> '[' expr '||' lc_exprs ']' :
     {lc, ?range_anno('$1', '$5'), '$2', '$4'}.
+map_comprehension -> '#' '{' expr '||' lc_exprs '}' :
+    {mc, ?range_anno('$1', '$6'), '$3', '$5'}.
 binary_comprehension -> '<<' expr_max '||' lc_exprs '>>' :
     {bc, ?range_anno('$1', '$5'), '$2', '$4'}.
 
@@ -702,6 +706,7 @@ Erlang code.
     | af_remote_call()
     | af_args(abstract_expr())
     | af_list_comprehension()
+    | af_map_comprehension()
     | af_binary_comprehension()
     | af_block()
     | af_if()
@@ -734,6 +739,9 @@ Erlang code.
 -type af_list_comprehension() ::
     {'lc', anno(), af_template(), af_qualifier_seq()}.
 
+-type af_map_comprehension() ::
+    {'mc', anno(), af_assoc(abstract_expr()), af_qualifier_seq()}.
+
 -type af_binary_comprehension() ::
     {'bc', anno(), af_template(), af_qualifier_seq()}.
 

diff --git a/src/erlfmt_recomment.erl b/src/erlfmt_recomment.erl
@@ -173,7 +173,7 @@ insert_nested({record, Meta, Expr0, Name, Values0}, Comments0) ->
     Values = insert_expr_container(Values0, Comments1),
     {{record, Meta, Expr, Name, Values}, []};
 insert_nested({Comprehension, Meta, Expr0, LcExprs0}, Comments0) when
-    Comprehension =:= lc; Comprehension =:= bc
+    Comprehension =:= lc; Comprehension =:= bc; Comprehension =:= mc
 ->
     {Expr, Comments1} = insert_expr(Expr0, Comments0),
     LcExprs = insert_expr_container(LcExprs0, Comments1),

diff --git a/test/erlfmt_SUITE.erl b/test/erlfmt_SUITE.erl
@@ -50,6 +50,7 @@
     annos/1,
     shebang/1,
     dotted/1,
+    map_comprehension/1,
     snapshot_simple_comments/1,
     snapshot_big_binary/1,
     snapshot_attributes/1,
@@ -141,7 +142,8 @@ groups() ->
             shebang,
             do_not_crash_on_bad_record,
             raw_string_anno,
-            dotted
+            dotted,
+            map_comprehension
         ]},
         {snapshot_tests, [parallel], [
             snapshot_simple_comments,
@@ -981,6 +983,14 @@ dotted(Config) when is_list(Config) ->
         parse_expr("#Port<0.1>")
     ).
 
+map_comprehension(Config) when is_list(Config) ->
+    ?assertMatch(
+        {mc, _, {map_field_assoc, _, {var, _, 'A'}, {var, _, 'B'}}, [
+            {generate, _, {map_field_exact, _, {var, _, 'A'}, {var, _, 'B'}}, {var, _, 'M'}}
+        ]},
+        parse_expr("#{A => B || A := B <- M}")
+    ).
+
 parse_expr(String) ->
     {function, _, [{clause, _, _, empty, [Expr]}]} = parse_form("f() -> " ++ String ++ "."),
     Expr.

diff --git a/test/erlfmt_format_SUITE.erl b/test/erlfmt_format_SUITE.erl
@@ -46,6 +46,7 @@
     record_index/1,
     record_field/1,
     list_comprehension/1,
+    map_comprehension/1,
     binary_comprehension/1,
     call/1,
     block/1,
@@ -146,6 +147,7 @@ groups() ->
         ]},
         {comprehensions, [parallel], [
             list_comprehension,
+            map_comprehension,
             binary_comprehension
         ]}
     ].
@@ -1994,6 +1996,7 @@ update_edgecase(Config) when is_list(Config) ->
 list_comprehension(Config) when is_list(Config) ->
     ?assertFormat("[X||X<-List]", "[X || X <- List]\n"),
     ?assertSame("[X || {X, Y} <- Results, X >= Y]\n"),
+    ?assertSame("[X || X := Y <- ResultMap, X >= Y]\n"),
     ?assertSame("[X || <<X, Y>> <= Results, X >= Y]\n"),
     ?assertFormat(
         "[[Very, Long, Expression] || X <- Y, X < 10]",
@@ -2080,9 +2083,78 @@ list_comprehension(Config) when is_list(Config) ->
         "]).\n"
     ).
 
+map_comprehension(Config) when is_list(Config) ->
+    ?assertFormat("#{X=>X||X<-List}", "#{X => X || X <- List}\n"),
+    ?assertFormat("#{X=>Y||X:=Y<-Map}", "#{X => Y || X := Y <- Map}\n"),
+    ?assertSame("#{X => X || {X, Y} <- Results, X >= Y}\n"),
+    ?assertSame("#{X => X || <<X, Y>> <= Results, X >= Y}\n"),
+    ?assertFormat(
+        "#{[Very, Long, Expression] => X || X <- Y, X < 10}",
+        "#{\n"
+        "    [\n"
+        "        Very,\n"
+        "        Long,\n"
+        "        Expression\n"
+        "    ] => X\n"
+        " || X <- Y, X < 10\n"
+        "}\n",
+        25
+    ),
+    ?assertFormat(
+        "#{X => X || X <- Y, X < 10\n"
+        " % trailing comment\n"
+        "}",
+        "#{\n"
+        "    X => X\n"
+        " || X <- Y,\n"
+        "    X < 10\n"
+        "    % trailing comment\n"
+        "}\n"
+    ),
+    ?assertFormat(
+        "#{{Very, Long, Expression} => X || X <- Y, X < 10}",
+        "#{\n"
+        "    {Very, Long,\n"
+        "        Expression} => X\n"
+        " || X <- Y, X < 10\n"
+        "}\n",
+        25
+    ),
+    ?assertFormat(
+        "#{X => X || X <- LongExpr123, X < 10}",
+        "#{\n"
+        "    X => X\n"
+        " || X <- LongExpr123,\n"
+        "    X < 10\n"
+        "}\n",
+        25
+    ),
+    ?assertFormat(
+        "#{X => X || X <- VeryLongExpression, X < 10}",
+        "#{\n"
+        "    X => X\n"
+        " || X <-\n"
+        "        VeryLongExpression,\n"
+        "    X < 10\n"
+        "}\n",
+        25
+    ),
+    ?assertFormat(
+        "#{X => X || X := VeryLongExpression <- VeryLongExpression, X < 10}",
+        "#{\n"
+        "    X => X\n"
+        " || X :=\n"
+        "        VeryLongExpression <-\n"
+        "        VeryLongExpression,\n"
+        "    X < 10\n"
+        "}\n",
+        25
+    ).
+
 binary_comprehension(Config) when is_list(Config) ->
     ?assertFormat("<<X||X<-List>>", "<<X || X <- List>>\n"),
     ?assertSame("<<X || <<X, Y>> <= Results, X >= Y>>\n"),
+    ?assertSame("<<X || X := Y <- ResultMap, X >= Y>>\n"),
     ?assertSame(
         "<<\n"
         "    X\n"