format properly multiline type annotations

compiler/fmt/src/annotation.rs

@@ -37,10 +37,6 @@ pub enum Newlines {
     No,
 }
 
-pub fn fmt_annotation<'a>(buf: &mut String<'a>, annotation: &'a TypeAnnotation<'a>, indent: u16) {
-    annotation.format(buf, indent);
-}
-
 pub trait Formattable<'a> {
     fn is_multiline(&self) -> bool;
 
@@ -313,8 +309,18 @@ impl<'a> Formattable<'a> for TypeAnnotation<'a> {
                 rhs.value.format(buf, indent);
             }
 
-            SpaceBefore(ann, _spaces) | SpaceAfter(ann, _spaces) => {
-                ann.format_with_options(buf, parens, newlines, indent)
+            SpaceBefore(ann, spaces) => {
+                newline(buf, indent + INDENT);
+                fmt_comments_only(buf, spaces.iter(), NewlineAt::Bottom, indent + INDENT);
+                ann.format_with_options(buf, parens, newlines, indent + INDENT)
+            }
+            SpaceAfter(ann, spaces) => {
+                ann.format_with_options(buf, parens, newlines, indent);
+                fmt_comments_only(buf, spaces.iter(), NewlineAt::Bottom, indent);
+                // seems like a lot of spaceAfter are not constructible
+                // since most of the time we use "SpaceBefore".
+                // Is this specific case really unreachable?
+                unreachable!("cannot have space after type annotation");
             }
 
             Malformed(raw) => buf.push_str(raw),

compiler/fmt/src/def.rs

@@ -36,7 +36,11 @@ impl<'a> Formattable<'a> for Def<'a> {
         match self {
             Annotation(loc_pattern, loc_annotation) => {
                 loc_pattern.format(buf, indent);
+                if loc_annotation.is_multiline() {
                     buf.push_str(" :");
+                } else {
+                    buf.push_str(" : ");
+                }
                 loc_annotation.format(buf, indent);
             }
             Alias { name, vars, ann } => {

compiler/fmt/tests/test_fmt.rs

@@ -782,8 +782,30 @@ mod test_fmt {
     }
 
     #[test]
-    fn comments_in_record_annotation() {
-        expr_formats_same(
+    fn multiline_type_definition() {
+        expr_formats_same(indoc!(
+            r#"
+                f :
+                    Int
+
+                f"#
+        ));
+    }
+
+    #[test]
+    fn multiline_empty_record_type_definition() {
+        expr_formats_same(indoc!(
+            r#"
+                f :
+                    {}
+
+                f"#
+        ));
+    }
+
+    #[test]
+    fn type_definition_comment_after_colon() {
+        expr_formats_to(
             indoc!(
                 r#"
                 f : # comment
@@ -791,6 +813,14 @@ mod test_fmt {
 
                 f"#
             ),
+            indoc!(
+                r#"
+                f :
+                    # comment
+                    {}
+
+                f"#
+            ),
         );
     }