Allow comments after the -> in a when

crates/compiler/fmt/src/expr.rs

@@ -2,7 +2,10 @@ use crate::annotation::{Formattable, Newlines, Parens};
 use crate::collection::{fmt_collection, Braces};
 use crate::def::fmt_defs;
 use crate::pattern::fmt_pattern;
-use crate::spaces::{count_leading_newlines, fmt_comments_only, fmt_spaces, NewlineAt, INDENT};
+use crate::spaces::{
+    count_leading_newlines, fmt_comments_only, fmt_spaces, fmt_spaces_no_blank_lines, NewlineAt,
+    INDENT,
+};
 use crate::Buf;
 use roc_module::called_via::{self, BinOp};
 use roc_parse::ast::{
@@ -756,15 +759,16 @@ fn fmt_when<'a, 'buf>(
 
         buf.push_str(" ->");
 
-        if is_multiline_expr {
-            buf.newline();
-        } else {
-            buf.spaces(1);
-        }
-
         match expr.value {
             Expr::SpaceBefore(nested, spaces) => {
-                fmt_comments_only(buf, spaces.iter(), NewlineAt::Bottom, indent + (INDENT * 2));
+                fmt_spaces_no_blank_lines(buf, spaces.iter(), indent + (INDENT * 2));
+
+                if is_multiline_expr {
+                    buf.ensure_ends_in_newline();
+                } else {
+                    buf.spaces(1);
+                }
+
                 nested.format_with_options(
                     buf,
                     Parens::NotNeeded,
@@ -773,6 +777,12 @@ fn fmt_when<'a, 'buf>(
                 );
             }
             _ => {
+                if is_multiline_expr {
+                    buf.ensure_ends_in_newline();
+                } else {
+                    buf.spaces(1);
+                }
+
                 expr.format_with_options(
                     buf,
                     Parens::NotNeeded,

crates/compiler/fmt/src/spaces.rs

@@ -32,26 +32,44 @@ pub fn fmt_default_spaces<'a, 'buf>(
     }
 }
 
+/// Like fmt_spaces, but disallows two consecutive newlines.
+pub fn fmt_spaces_no_blank_lines<'a, 'buf, I>(buf: &mut Buf<'buf>, spaces: I, indent: u16)
+where
+    I: Iterator<Item = &'a CommentOrNewline<'a>>,
+{
+    fmt_spaces_max_consecutive_newlines(buf, spaces, 1, indent)
+}
+
 pub fn fmt_spaces<'a, 'buf, I>(buf: &mut Buf<'buf>, spaces: I, indent: u16)
 where
     I: Iterator<Item = &'a CommentOrNewline<'a>>,
+{
+    fmt_spaces_max_consecutive_newlines(buf, spaces, 2, indent)
+}
+
+fn fmt_spaces_max_consecutive_newlines<'a, 'buf, I>(
+    buf: &mut Buf<'buf>,
+    spaces: I,
+    max_consecutive_newlines: usize,
+    indent: u16,
+) where
+    I: Iterator<Item = &'a CommentOrNewline<'a>>,
 {
     use self::CommentOrNewline::*;
 
     // Only ever print two newlines back to back.
     // (Two newlines renders as one blank line.)
     let mut consecutive_newlines = 0;
-
     let mut encountered_comment = false;
 
     for space in spaces {
         match space {
             Newline => {
-                if !encountered_comment && (consecutive_newlines < 2) {
+                if !encountered_comment && (consecutive_newlines < max_consecutive_newlines) {
                     buf.newline();
 
                     // Don't bother incrementing it if we're already over the limit.
-                    // There's no upside, and it might eventually overflow,
+                    // There's no upside, and it might eventually overflow.
                     consecutive_newlines += 1;
                 }
             }

crates/compiler/fmt/tests/test_fmt.rs

@@ -4044,6 +4044,13 @@ mod test_fmt {
                     4
                         / 5 # comment 5
                         < 1 # comment 6
+                46 # first pattern comment
+                    | 95 # alternative comment 1
+                    | 126 # alternative comment 2
+                    | 150 -> # This comment goes after the ->
+                    # This comment is for the expr
+                    Str.appendScalar output (Num.toU32 byte)
+                        |> Result.withDefault "" # this will never fail
                 _ ->
                     42
             "#