3 files changed, 70 insertions, 33 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/assists/src/lib.rs b/crates/assists/src/lib.rs
index ae90d68a3..c589b08dc 100644
--- a/crates/assists/src/lib.rs
+++ b/crates/assists/src/lib.rs
@@ -66,13 +66,13 @@ pub struct GroupLabel(pub String);
 
 #[derive(Debug, Clone)]
 pub struct Assist {
-    id: AssistId,
+    pub id: AssistId,
     /// Short description of the assist, as shown in the UI.
     label: String,
-    group: Option<GroupLabel>,
+    pub group: Option<GroupLabel>,
     /// Target ranges are used to sort assists: the smaller the target range,
     /// the more specific assist is, and so it should be sorted first.
-    target: TextRange,
+    pub target: TextRange,
 }
 
 #[derive(Debug, Clone)]
@@ -82,6 +82,11 @@ pub struct ResolvedAssist {
 }
 
 impl Assist {
+    fn new(id: AssistId, label: String, group: Option<GroupLabel>, target: TextRange) -> Assist {
+        assert!(label.starts_with(char::is_uppercase));
+        Assist { id, label, group, target }
+    }
+
     /// Return all the assists applicable at the given position.
     ///
     /// Assists are returned in the "unresolved" state, that is only labels are
@@ -114,30 +119,8 @@ impl Assist {
         acc.finish_resolved()
     }
 
-    pub(crate) fn new(
-        id: AssistId,
-        label: String,
-        group: Option<GroupLabel>,
-        target: TextRange,
-    ) -> Assist {
-        assert!(label.starts_with(|c: char| c.is_uppercase()));
-        Assist { id, label, group, target }
-    }
-
-    pub fn id(&self) -> AssistId {
-        self.id
-    }
-
-    pub fn label(&self) -> String {
-        self.label.clone()
-    }
-
-    pub fn group(&self) -> Option<GroupLabel> {
-        self.group.clone()
-    }
-
-    pub fn target(&self) -> TextRange {
-        self.target
+    pub fn label(&self) -> &str {
+        self.label.as_str()
     }
 }
 
diff --git a/crates/assists/src/tests/generated.rs b/crates/assists/src/tests/generated.rs
index d16e6fb0a..173567003 100644
--- a/crates/assists/src/tests/generated.rs
+++ b/crates/assists/src/tests/generated.rs
@@ -82,8 +82,8 @@ trait Trait {
 impl Trait for () {
     Type X = ();
     fn foo(&self) {}
-    $0fn bar(&self) {}
 
+    $0fn bar(&self) {}
 }
 "#####,
     )
@@ -115,7 +115,6 @@ impl Trait<u32> for () {
     fn foo(&self) -> u32 {
         ${0:todo!()}
     }
-
 }
 "#####,
     )