Merge pull request #3472 from rtfeldman/fix-comments

Allow single-line comments in multiline annotations

crates/compiler/fmt/src/collection.rs

@@ -1,8 +1,8 @@
-use roc_parse::ast::{Collection, ExtractSpaces};
+use roc_parse::ast::{Collection, CommentOrNewline, ExtractSpaces};
 
 use crate::{
     annotation::{Formattable, Newlines},
-    spaces::{count_leading_newlines, fmt_comments_only, NewlineAt, INDENT},
+    spaces::{fmt_comments_only, NewlineAt, INDENT},
     Buf,
 };
 
@@ -41,24 +41,43 @@ pub fn fmt_collection<'a, 'buf, T: ExtractSpaces<'a> + Formattable>(
         buf.push(start);
 
         for (index, item) in items.iter().enumerate() {
-            let item = item.extract_spaces();
             let is_first_item = index == 0;
-
+            let item = item.extract_spaces();
-            buf.newline();
-
-            if !item.before.is_empty() {
             let is_only_newlines = item.before.iter().all(|s| s.is_newline());
 
-                if !is_first_item
+            if item.before.is_empty() || is_only_newlines {
-                    && !is_only_newlines
+                buf.ensure_ends_in_newline();
-                    && count_leading_newlines(item.before.iter()) > 1
+            } else {
+                if is_first_item {
+                    // The first item in a multiline collection always begins with exactly
+                    // one newline (so the delimiter is at the end of its own line),
+                    // and that newline appears before the first comment (if there is one).
+                    buf.ensure_ends_in_newline();
+                } else {
+                    if item.before.starts_with(&[CommentOrNewline::Newline]) {
+                        buf.ensure_ends_in_newline();
+                    }
+
+                    if item
+                        .before
+                        .starts_with(&[CommentOrNewline::Newline, CommentOrNewline::Newline])
                     {
+                        // If there's a comment, and it's not on the first item,
+                        // and it's preceded by at least one blank line, maintain 1 blank line.
+                        // (We already ensured that it ends in a newline, so this will turn that
+                        // into a blank line.)
+
+                        buf.newline();
+                    }
+                }
+
+                fmt_comments_only(buf, item.before.iter(), NewlineAt::None, item_indent);
+
+                if !is_only_newlines {
+                    if item.before.ends_with(&[CommentOrNewline::Newline]) {
                         buf.newline();
                     }
 
-                fmt_comments_only(buf, item.before.iter(), NewlineAt::Bottom, item_indent);
-
-                if !is_only_newlines && count_leading_newlines(item.before.iter().rev()) > 0 {
                     buf.newline();
                 }
             }
@@ -68,21 +87,33 @@ pub fn fmt_collection<'a, 'buf, T: ExtractSpaces<'a> + Formattable>(
             buf.push(',');
 
             if !item.after.is_empty() {
-                fmt_comments_only(buf, item.after.iter(), NewlineAt::Top, item_indent);
+                if item.after.iter().any(|s| s.is_newline()) {
+                    buf.newline();
+                }
+
+                fmt_comments_only(buf, item.after.iter(), NewlineAt::None, item_indent);
             }
         }
 
-        if count_leading_newlines(items.final_comments().iter()) > 1 {
+        if items.final_comments().iter().any(|s| s.is_newline()) {
+            buf.newline();
+        }
+
+        if items
+            .final_comments()
+            .starts_with(&[CommentOrNewline::Newline, CommentOrNewline::Newline])
+        {
             buf.newline();
         }
 
         fmt_comments_only(
             buf,
             items.final_comments().iter(),
-            NewlineAt::Top,
+            NewlineAt::None,
             item_indent,
         );
-        buf.newline();
+
+        buf.ensure_ends_in_newline();
         buf.indent(braces_indent);
     } else {
         // is_multiline == false

crates/compiler/fmt/tests/test_fmt.rs

@@ -2347,8 +2347,7 @@ mod test_fmt {
             indoc!(
                 r#"
                 f : {
-                    x : Int *,
+                    x : Int *, # comment 1
-                    # comment 1
                     # comment 2
                 }
 
@@ -4588,7 +4587,7 @@ mod test_fmt {
     }
 
     #[test]
-    fn multiline_tag_union_annotation() {
+    fn multiline_tag_union_annotation_no_comments() {
         expr_formats_same(indoc!(
             r#"
             b : [
@@ -4694,8 +4693,7 @@ mod test_fmt {
                 b : [
                     True,
                     # comment 1
-                    False,
+                    False, # comment 2
-                    # comment 2
                     # comment 3
                 ]
 
@@ -5082,6 +5080,38 @@ mod test_fmt {
         );
     }
 
+    #[test]
+    fn comments_in_multiline_tag_union_annotation() {
+        expr_formats_to(
+            indoc!(
+                r#"
+                    UnionAnn : [
+                        Foo, # comment 1
+                        Bar, # comment 2
+                        Baz, # comment 3
+                             # comment 4 line 1
+                             # comment 4 line 2
+                    ]
+
+                    0
+                "#
+            ),
+            indoc!(
+                r#"
+                    UnionAnn : [
+                        Foo, # comment 1
+                        Bar, # comment 2
+                        Baz, # comment 3
+                        # comment 4 line 1
+                        # comment 4 line 2
+                    ]
+
+                    0
+                "#
+            ),
+        );
+    }
+
     #[test]
     /// Test that everything under examples/ is formatted correctly
     /// If this test fails on your diff, it probably means you need to re-format the examples.