fixup whitespace when adding missing impl items

2 files changed, 89 insertions, 7 deletions

diff --git a/crates/assists/src/handlers/add_missing_impl_members.rs b/crates/assists/src/handlers/add_missing_impl_members.rs
index 81b61ebf8..83a2ada9a 100644
--- a/crates/assists/src/handlers/add_missing_impl_members.rs
+++ b/crates/assists/src/handlers/add_missing_impl_members.rs
@@ -48,7 +48,6 @@ enum AddMissingImplMembersMode {
 //     fn foo(&self) -> u32 {
 //         ${0:todo!()}
 //     }
-//
 // }
 // ```
 pub(crate) fn add_missing_impl_members(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
@@ -89,8 +88,8 @@ pub(crate) fn add_missing_impl_members(acc: &mut Assists, ctx: &AssistContext) -
 // impl Trait for () {
 //     Type X = ();
 //     fn foo(&self) {}
-//     $0fn bar(&self) {}
 //
+//     $0fn bar(&self) {}
 // }
 // ```
 pub(crate) fn add_missing_default_members(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
@@ -240,15 +239,18 @@ struct S;
 
 impl Foo for S {
     fn bar(&self) {}
+
     $0type Output;
+
     const CONST: usize = 42;
+
     fn foo(&self) {
         todo!()
     }
+
     fn baz(&self) {
         todo!()
     }
-
 }"#,
         );
     }
@@ -281,10 +283,10 @@ struct S;
 
 impl Foo for S {
     fn bar(&self) {}
+
     fn foo(&self) {
         ${0:todo!()}
     }
-
 }"#,
         );
     }
@@ -599,6 +601,7 @@ trait Foo {
 struct S;
 impl Foo for S {
     $0type Output;
+
     fn foo(&self) {
         todo!()
     }
@@ -708,4 +711,56 @@ impl Tr for () {
 }"#,
         )
     }
+
+    #[test]
+    fn test_whitespace_fixup_preserves_bad_tokens() {
+        check_assist(
+            add_missing_impl_members,
+            r#"
+trait Tr {
+    fn foo();
+}
+
+impl Tr for ()<|> {
+    +++
+}"#,
+            r#"
+trait Tr {
+    fn foo();
+}
+
+impl Tr for () {
+    fn foo() {
+        ${0:todo!()}
+    }
+    +++
+}"#,
+        )
+    }
+
+    #[test]
+    fn test_whitespace_fixup_preserves_comments() {
+        check_assist(
+            add_missing_impl_members,
+            r#"
+trait Tr {
+    fn foo();
+}
+
+impl Tr for ()<|> {
+    // very important
+}"#,
+            r#"
+trait Tr {
+    fn foo();
+}
+
+impl Tr for () {
+    fn foo() {
+        ${0:todo!()}
+    }
+    // very important
+}"#,
+        )
+    }
 }
diff --git a/crates/syntax/src/ast/edit.rs b/crates/syntax/src/ast/edit.rs
index 190746e09..b295b5bc6 100644
--- a/crates/syntax/src/ast/edit.rs
+++ b/crates/syntax/src/ast/edit.rs
@@ -91,29 +91,56 @@ impl ast::AssocItemList {
             res = make_multiline(res);
         }
         items.into_iter().for_each(|it| res = res.append_item(it));
-        res
+        res.fixup_trailing_whitespace().unwrap_or(res)
     }
 
     #[must_use]
     pub fn append_item(&self, item: ast::AssocItem) -> ast::AssocItemList {
-        let (indent, position) = match self.assoc_items().last() {
+        let (indent, position, whitespace) = match self.assoc_items().last() {
             Some(it) => (
                 leading_indent(it.syntax()).unwrap_or_default().to_string(),
                 InsertPosition::After(it.syntax().clone().into()),
+                "\n\n",
             ),
             None => match self.l_curly_token() {
                 Some(it) => (
                     "    ".to_string() + &leading_indent(self.syntax()).unwrap_or_default(),
                     InsertPosition::After(it.into()),
+                    "\n",
                 ),
                 None => return self.clone(),
             },
         };
-        let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
+        let ws = tokens::WsBuilder::new(&format!("{}{}", whitespace, indent));
         let to_insert: ArrayVec<[SyntaxElement; 2]> =
             [ws.ws().into(), item.syntax().clone().into()].into();
         self.insert_children(position, to_insert)
     }
+
+    /// Remove extra whitespace between last item and closing curly brace.
+    fn fixup_trailing_whitespace(&self) -> Option<ast::AssocItemList> {
+        let first_token_after_items = self
+            .assoc_items()
+            .last()?
+            .syntax()
+            .next_sibling_or_token()?;
+        let last_token_before_curly = self
+            .r_curly_token()?
+            .prev_sibling_or_token()?;
+        if last_token_before_curly != first_token_after_items {
+            // there is something more between last item and
+            // right curly than just whitespace - bail out
+            return None;
+        }
+        let whitespace = last_token_before_curly
+            .clone()
+            .into_token()
+            .and_then(ast::Whitespace::cast)?;
+        let text = whitespace.syntax().text();
+        let newline = text.rfind("\n")?;
+        let keep = tokens::WsBuilder::new(&text[newline..]);
+        Some(self.replace_children(first_token_after_items..=last_token_before_curly, std::iter::once(keep.ws().into())))
+    }
 }
 
 impl ast::RecordExprFieldList {