4 files changed, 74 insertions, 44 deletions

diff --git a/crates/assists/src/assist_context.rs b/crates/assists/src/assist_context.rs
index 79574b9ac..11c171fc2 100644
--- a/crates/assists/src/assist_context.rs
+++ b/crates/assists/src/assist_context.rs
@@ -6,6 +6,7 @@ use algo::find_covering_element;
 use base_db::{FileId, FileRange};
 use hir::Semantics;
 use ide_db::{
+    label::Label,
     source_change::{SourceChange, SourceFileEdit},
     RootDatabase,
 };
@@ -157,8 +158,9 @@ impl Assists {
         if !self.is_allowed(&id) {
             return None;
         }
-        let label = Assist::new(id, label.into(), None, target);
-        self.add_impl(label, f)
+        let label = Label::new(label.into());
+        let assist = Assist { id, label, group: None, target };
+        self.add_impl(assist, f)
     }
 
     pub(crate) fn add_group(
@@ -172,12 +174,12 @@ impl Assists {
         if !self.is_allowed(&id) {
             return None;
         }
-
-        let label = Assist::new(id, label.into(), Some(group.clone()), target);
-        self.add_impl(label, f)
+        let label = Label::new(label.into());
+        let assist = Assist { id, label, group: Some(group.clone()), target };
+        self.add_impl(assist, f)
     }
 
-    fn add_impl(&mut self, label: Assist, f: impl FnOnce(&mut AssistBuilder)) -> Option<()> {
+    fn add_impl(&mut self, assist: Assist, f: impl FnOnce(&mut AssistBuilder)) -> Option<()> {
         let source_change = if self.resolve {
             let mut builder = AssistBuilder::new(self.file);
             f(&mut builder);
@@ -186,7 +188,7 @@ impl Assists {
             None
         };
 
-        self.buf.push((label, source_change));
+        self.buf.push((assist, source_change));
         Some(())
     }
 
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..14834480a 100644
--- a/crates/assists/src/lib.rs
+++ b/crates/assists/src/lib.rs
@@ -19,7 +19,7 @@ pub mod ast_transform;
 
 use base_db::FileRange;
 use hir::Semantics;
-use ide_db::{source_change::SourceChange, RootDatabase};
+use ide_db::{label::Label, source_change::SourceChange, RootDatabase};
 use syntax::TextRange;
 
 pub(crate) use crate::assist_context::{AssistContext, Assists};
@@ -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 label: Label,
+    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)]
@@ -113,32 +113,6 @@ 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
-    }
 }
 
 mod handlers {
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!()}
     }
-
 }
 "#####,
     )