Feedback: dry out term, make fmt_spaces_with_newline_mode more readable, explicitly implement all the variants in ann_lift_spaces

crates/compiler/fmt/src/annotation.rs

@@ -1,6 +1,7 @@
 use crate::{
     collection::{fmt_collection, Braces},
     expr::merge_spaces_conservative,
+    pattern::pattern_lift_spaces_after,
     spaces::{fmt_comments_only, fmt_spaces, NewlineAt, INDENT},
     Buf,
 };
@@ -963,11 +964,64 @@ pub fn ann_lift_spaces<'a, 'b: 'a>(
                 }
             }
         }
-        _ => Spaces {
+        TypeAnnotation::BoundVariable(_)
+        | TypeAnnotation::Inferred
+        | TypeAnnotation::Wildcard
+        | TypeAnnotation::Malformed(_) => Spaces {
             before: &[],
             item: *ann,
             after: &[],
         },
+        TypeAnnotation::Where(inner, clauses) => {
+            let new_inner = ann_lift_spaces_before(arena, &inner.value);
+            let new_clauses = arena.alloc_slice_copy(clauses);
+            let after = if let Some(last) = new_clauses.last_mut() {
+                let lifted = implements_clause_lift_spaces_after(arena, &last.value);
+                last.value = lifted.item;
+                lifted.after
+            } else {
+                &[]
+            };
+            Spaces {
+                before: new_inner.before,
+                item: TypeAnnotation::Where(arena.alloc(Loc::at_zero(new_inner.item)), new_clauses),
+                after,
+            }
+        }
+        TypeAnnotation::As(ann, comments, type_header) => {
+            let new_ann = ann_lift_spaces_before(arena, &ann.value);
+            let new_header = type_head_lift_spaces_after(arena, type_header);
+            Spaces {
+                before: new_ann.before,
+                item: TypeAnnotation::As(
+                    arena.alloc(Loc::at_zero(new_ann.item)),
+                    comments,
+                    new_header.item,
+                ),
+                after: new_header.after,
+            }
+        }
+    }
+}
+
+fn implements_clause_lift_spaces_after<'a, 'b: 'a>(
+    arena: &'a Bump,
+    value: &ImplementsClause<'b>,
+) -> SpacesAfter<'a, ImplementsClause<'a>> {
+    let new_abilities = arena.alloc_slice_copy(value.abilities);
+    let after = if let Some(last) = new_abilities.last_mut() {
+        let lifted = ann_lift_spaces_after(arena, &last.value);
+        last.value = lifted.item;
+        lifted.after
+    } else {
+        &[]
+    };
+    SpacesAfter {
+        item: ImplementsClause {
+            var: value.var,
+            abilities: new_abilities,
+        },
+        after,
     }
 }
 
@@ -993,6 +1047,27 @@ pub fn ann_lift_spaces_after<'a, 'b: 'a>(
     }
 }
 
+pub fn type_head_lift_spaces_after<'a, 'b: 'a>(
+    arena: &'a Bump,
+    header: &TypeHeader<'b>,
+) -> SpacesAfter<'a, TypeHeader<'a>> {
+    let new_vars = arena.alloc_slice_copy(header.vars);
+    let after = if let Some(last) = new_vars.last_mut() {
+        let lifted = pattern_lift_spaces_after(arena, &last.value);
+        last.value = lifted.item;
+        lifted.after
+    } else {
+        &[]
+    };
+    SpacesAfter {
+        item: TypeHeader {
+            name: header.name,
+            vars: new_vars,
+        },
+        after,
+    }
+}
+
 type Sp<'a> = &'a [CommentOrNewline<'a>];
 
 #[derive(Copy, Clone, Debug)]

crates/compiler/fmt/src/spaces.rs

@@ -37,17 +37,22 @@ pub fn fmt_spaces_with_newline_mode(
         mode,
         SpacesNewlineMode::SkipNewlinesAtStart | SpacesNewlineMode::SkipNewlinesAtBoth
     ) {
-        while let Some(CommentOrNewline::Newline) = spaces.first() {
-            spaces = &spaces[1..];
-        }
+        let skip_count = spaces
+            .iter()
+            .take_while(|s| *s == &CommentOrNewline::Newline)
+            .count();
+        spaces = &spaces[skip_count..];
     }
     if matches!(
         mode,
         SpacesNewlineMode::SkipNewlinesAtEnd | SpacesNewlineMode::SkipNewlinesAtBoth
     ) {
-        while let Some(CommentOrNewline::Newline) = spaces.last() {
-            spaces = &spaces[..spaces.len() - 1];
-        }
+        let skip_count = spaces
+            .iter()
+            .rev()
+            .take_while(|s| *s == &CommentOrNewline::Newline)
+            .count();
+        spaces = &spaces[..spaces.len() - skip_count];
     }
     fmt_spaces(buf, spaces.iter(), indent);
 }

crates/compiler/parse/src/type_annotation.rs

@@ -111,7 +111,9 @@ fn parse_type_alias_after_as<'a>() -> impl Parser<'a, TypeHeader<'a>, EType<'a>>
     )
 }
 
-fn term<'a>(stop_at_surface_has: bool) -> impl Parser<'a, Loc<TypeAnnotation<'a>>, EType<'a>> {
+fn term_fragment<'a>(
+    stop_at_surface_has: bool,
+) -> impl Parser<'a, Loc<TypeAnnotation<'a>>, EType<'a>> {
     one_of!(
         loc_wildcard(),
         loc_inferred(),
@@ -124,8 +126,14 @@ fn term<'a>(stop_at_surface_has: bool) -> impl Parser<'a, Loc<TypeAnnotation<'a>
             EType::TTagUnion,
             tag_union_type(stop_at_surface_has)
         )),
-        loc(specialize_err(EType::TApply, concrete_type())),
         loc(parse_type_variable(stop_at_surface_has)),
+    )
+}
+
+fn term<'a>(stop_at_surface_has: bool) -> impl Parser<'a, Loc<TypeAnnotation<'a>>, EType<'a>> {
+    one_of!(
+        term_fragment(stop_at_surface_has),
+        loc(specialize_err(EType::TApply, concrete_type())),
         fail(EType::TStart),
     )
     .trace("type_annotation:term")
@@ -137,19 +145,8 @@ fn term_or_apply_with_as<'a>(
     map_with_arena(
         and(
             one_of!(
-                loc_wildcard(),
-                loc_inferred(),
-                specialize_err(EType::TInParens, loc_type_in_parens(stop_at_surface_has)),
-                loc(specialize_err(
-                    EType::TRecord,
-                    record_type(stop_at_surface_has)
-                )),
-                loc(specialize_err(
-                    EType::TTagUnion,
-                    tag_union_type(stop_at_surface_has)
-                )),
+                term_fragment(stop_at_surface_has),
                 loc(applied_type(stop_at_surface_has)),
-                loc(parse_type_variable(stop_at_surface_has)),
                 fail(EType::TStart),
             ),
             // Inline alias notation, e.g. [Nil, Cons a (List a)] as List a