5 files changed, 82 insertions, 43 deletions

diff --git a/crates/ra_assists/src/assist_context.rs b/crates/ra_assists/src/assist_context.rs
index 0dcd9df61..005c17776 100644
--- a/crates/ra_assists/src/assist_context.rs
+++ b/crates/ra_assists/src/assist_context.rs
@@ -194,20 +194,30 @@ impl AssistBuilder {
     pub(crate) fn insert(&mut self, offset: TextSize, text: impl Into<String>) {
         self.edit.insert(offset, text.into())
     }
-    /// Append specified `text` at the given `offset`
+    /// Append specified `snippet` at the given `offset`
     pub(crate) fn insert_snippet(
         &mut self,
         _cap: SnippetCap,
         offset: TextSize,
-        text: impl Into<String>,
+        snippet: impl Into<String>,
     ) {
         self.is_snippet = true;
-        self.edit.insert(offset, text.into())
+        self.insert(offset, snippet);
     }
     /// Replaces specified `range` of text with a given string.
     pub(crate) fn replace(&mut self, range: TextRange, replace_with: impl Into<String>) {
         self.edit.replace(range, replace_with.into())
     }
+    /// Replaces specified `range` of text with a given `snippet`.
+    pub(crate) fn replace_snippet(
+        &mut self,
+        _cap: SnippetCap,
+        range: TextRange,
+        snippet: impl Into<String>,
+    ) {
+        self.is_snippet = true;
+        self.replace(range, snippet);
+    }
     pub(crate) fn replace_ast<N: AstNode>(&mut self, old: N, new: N) {
         algo::diff(old.syntax(), new.syntax()).into_text_edit(&mut self.edit)
     }
diff --git a/crates/ra_assists/src/handlers/add_function.rs b/crates/ra_assists/src/handlers/add_function.rs
index a0709630d..24f931a85 100644
--- a/crates/ra_assists/src/handlers/add_function.rs
+++ b/crates/ra_assists/src/handlers/add_function.rs
@@ -10,7 +10,11 @@ use ra_syntax::{
 };
 use rustc_hash::{FxHashMap, FxHashSet};
 
-use crate::{assist_config::SnippetCap, utils::render_snippet, AssistContext, AssistId, Assists};
+use crate::{
+    assist_config::SnippetCap,
+    utils::{render_snippet, Cursor},
+    AssistContext, AssistId, Assists,
+};
 
 // Assist: add_function
 //
@@ -81,7 +85,11 @@ struct FunctionTemplate {
 impl FunctionTemplate {
     fn to_string(&self, cap: Option<SnippetCap>) -> String {
         let f = match cap {
-            Some(cap) => render_snippet(cap, self.fn_def.syntax(), self.placeholder_expr.syntax()),
+            Some(cap) => render_snippet(
+                cap,
+                self.fn_def.syntax(),
+                Cursor::Replace(self.placeholder_expr.syntax()),
+            ),
             None => self.fn_def.to_string(),
         };
         format!("{}{}{}", self.leading_ws, f, self.trailing_ws)
diff --git a/crates/ra_assists/src/handlers/add_missing_impl_members.rs b/crates/ra_assists/src/handlers/add_missing_impl_members.rs
index 22e1156d2..d7aa06947 100644
--- a/crates/ra_assists/src/handlers/add_missing_impl_members.rs
+++ b/crates/ra_assists/src/handlers/add_missing_impl_members.rs
@@ -11,7 +11,7 @@ use ra_syntax::{
 use crate::{
     assist_context::{AssistContext, Assists},
     ast_transform::{self, AstTransform, QualifyPaths, SubstituteTypeParams},
-    utils::{get_missing_assoc_items, resolve_target_trait},
+    utils::{get_missing_assoc_items, render_snippet, resolve_target_trait, Cursor},
     AssistId,
 };
 
@@ -45,7 +45,7 @@ enum AddMissingImplMembersMode {
 // }
 //
 // impl Trait<u32> for () {
-//     fn foo(&self) -> u32 {
+//     $0fn foo(&self) -> u32 {
 //         todo!()
 //     }
 //
@@ -89,7 +89,7 @@ pub(crate) fn add_missing_impl_members(acc: &mut Assists, ctx: &AssistContext) -
 // impl Trait for () {
 //     Type X = ();
 //     fn foo(&self) {}
-//     fn bar(&self) {}
+//     $0fn bar(&self) {}
 //
 // }
 // ```
@@ -147,7 +147,7 @@ fn add_missing_impl_members_inner(
     }
 
     let target = impl_def.syntax().text_range();
-    acc.add(AssistId(assist_id), label, target, |edit| {
+    acc.add(AssistId(assist_id), label, target, |builder| {
         let n_existing_items = impl_item_list.assoc_items().count();
         let source_scope = ctx.sema.scope_for_def(trait_);
         let target_scope = ctx.sema.scope(impl_item_list.syntax());
@@ -162,13 +162,21 @@ fn add_missing_impl_members_inner(
             })
             .map(|it| edit::remove_attrs_and_docs(&it));
         let new_impl_item_list = impl_item_list.append_items(items);
-        let cursor_position = {
-            let first_new_item = new_impl_item_list.assoc_items().nth(n_existing_items).unwrap();
-            first_new_item.syntax().text_range().start()
+        let first_new_item = new_impl_item_list.assoc_items().nth(n_existing_items).unwrap();
+
+        let original_range = impl_item_list.syntax().text_range();
+        match ctx.config.snippet_cap {
+            None => builder.replace(original_range, new_impl_item_list.to_string()),
+            Some(cap) => builder.replace_snippet(
+                cap,
+                original_range,
+                render_snippet(
+                    cap,
+                    new_impl_item_list.syntax(),
+                    Cursor::Before(first_new_item.syntax()),
+                ),
+            ),
         };
-
-        edit.replace_ast(impl_item_list, new_impl_item_list);
-        edit.set_cursor(cursor_position);
     })
 }
 
@@ -222,7 +230,7 @@ struct S;
 
 impl Foo for S {
     fn bar(&self) {}
-    <|>type Output;
+    $0type Output;
     const CONST: usize = 42;
     fn foo(&self) {
         todo!()
@@ -263,7 +271,7 @@ struct S;
 
 impl Foo for S {
     fn bar(&self) {}
-    <|>fn foo(&self) {
+    $0fn foo(&self) {
         todo!()
     }
 
@@ -283,7 +291,7 @@ impl Foo for S { <|> }"#,
 trait Foo { fn foo(&self); }
 struct S;
 impl Foo for S {
-    <|>fn foo(&self) {
+    $0fn foo(&self) {
         todo!()
     }
 }"#,
@@ -302,7 +310,7 @@ impl Foo<u32> for S { <|> }"#,
 trait Foo<T> { fn foo(&self, t: T) -> &T; }
 struct S;
 impl Foo<u32> for S {
-    <|>fn foo(&self, t: u32) -> &u32 {
+    $0fn foo(&self, t: u32) -> &u32 {
         todo!()
     }
 }"#,
@@ -321,7 +329,7 @@ impl<U> Foo<U> for S { <|> }"#,
 trait Foo<T> { fn foo(&self, t: T) -> &T; }
 struct S;
 impl<U> Foo<U> for S {
-    <|>fn foo(&self, t: U) -> &U {
+    $0fn foo(&self, t: U) -> &U {
         todo!()
     }
 }"#,
@@ -340,7 +348,7 @@ impl Foo for S {}<|>"#,
 trait Foo { fn foo(&self); }
 struct S;
 impl Foo for S {
-    <|>fn foo(&self) {
+    $0fn foo(&self) {
         todo!()
     }
 }"#,
@@ -365,7 +373,7 @@ mod foo {
 }
 struct S;
 impl foo::Foo for S {
-    <|>fn foo(&self, bar: foo::Bar) {
+    $0fn foo(&self, bar: foo::Bar) {
         todo!()
     }
 }"#,
@@ -390,7 +398,7 @@ mod foo {
 }
 struct S;
 impl foo::Foo for S {
-    <|>fn foo(&self, bar: foo::Bar<u32>) {
+    $0fn foo(&self, bar: foo::Bar<u32>) {
         todo!()
     }
 }"#,
@@ -415,7 +423,7 @@ mod foo {
 }
 struct S;
 impl foo::Foo<u32> for S {
-    <|>fn foo(&self, bar: foo::Bar<u32>) {
+    $0fn foo(&self, bar: foo::Bar<u32>) {
         todo!()
     }
 }"#,
@@ -443,7 +451,7 @@ mod foo {
 struct Param;
 struct S;
 impl foo::Foo<Param> for S {
-    <|>fn foo(&self, bar: Param) {
+    $0fn foo(&self, bar: Param) {
         todo!()
     }
 }"#,
@@ -470,7 +478,7 @@ mod foo {
 }
 struct S;
 impl foo::Foo for S {
-    <|>fn foo(&self, bar: foo::Bar<u32>::Assoc) {
+    $0fn foo(&self, bar: foo::Bar<u32>::Assoc) {
         todo!()
     }
 }"#,
@@ -497,7 +505,7 @@ mod foo {
 }
 struct S;
 impl foo::Foo for S {
-    <|>fn foo(&self, bar: foo::Bar<foo::Baz>) {
+    $0fn foo(&self, bar: foo::Bar<foo::Baz>) {
         todo!()
     }
 }"#,
@@ -522,7 +530,7 @@ mod foo {
 }
 struct S;
 impl foo::Foo for S {
-    <|>fn foo(&self, bar: dyn Fn(u32) -> i32) {
+    $0fn foo(&self, bar: dyn Fn(u32) -> i32) {
         todo!()
     }
 }"#,
@@ -580,7 +588,7 @@ trait Foo {
 }
 struct S;
 impl Foo for S {
-    <|>type Output;
+    $0type Output;
     fn foo(&self) {
         todo!()
     }
@@ -614,7 +622,7 @@ trait Foo {
 }
 struct S;
 impl Foo for S {
-    <|>fn valid(some: u32) -> bool { false }
+    $0fn valid(some: u32) -> bool { false }
 }"#,
         )
     }
@@ -637,7 +645,7 @@ trait Foo<T = Self> {
 
 struct S;
 impl Foo for S {
-    <|>fn bar(&self, other: &Self) {
+    $0fn bar(&self, other: &Self) {
         todo!()
     }
 }"#,
@@ -662,7 +670,7 @@ trait Foo<T1, T2 = Self> {
 
 struct S<T>;
 impl Foo<T> for S<T> {
-    <|>fn bar(&self, this: &T, that: &Self) {
+    $0fn bar(&self, this: &T, that: &Self) {
         todo!()
     }
 }"#,
diff --git a/crates/ra_assists/src/tests/generated.rs b/crates/ra_assists/src/tests/generated.rs
index 1d82c245d..2522ec5df 100644
--- a/crates/ra_assists/src/tests/generated.rs
+++ b/crates/ra_assists/src/tests/generated.rs
@@ -150,7 +150,7 @@ trait Trait {
 impl Trait for () {
     Type X = ();
     fn foo(&self) {}
-    fn bar(&self) {}
+    $0fn bar(&self) {}
 
 }
 "#####,
@@ -180,7 +180,7 @@ trait Trait<T> {
 }
 
 impl Trait<u32> for () {
-    fn foo(&self) -> u32 {
+    $0fn foo(&self) -> u32 {
         todo!()
     }
 
diff --git a/crates/ra_assists/src/utils.rs b/crates/ra_assists/src/utils.rs
index 8a26a6808..9af27180b 100644
--- a/crates/ra_assists/src/utils.rs
+++ b/crates/ra_assists/src/utils.rs
@@ -15,18 +15,31 @@ use crate::assist_config::SnippetCap;
 
 pub(crate) use insert_use::insert_use_statement;
 
-pub(crate) fn render_snippet(
-    _cap: SnippetCap,
-    node: &SyntaxNode,
-    placeholder: &SyntaxNode,
-) -> String {
-    assert!(placeholder.ancestors().any(|it| it == *node));
-    let range = placeholder.text_range() - node.text_range().start();
+#[derive(Clone, Copy, Debug)]
+pub(crate) enum Cursor<'a> {
+    Replace(&'a SyntaxNode),
+    Before(&'a SyntaxNode),
+}
+
+impl<'a> Cursor<'a> {
+    fn node(self) -> &'a SyntaxNode {
+        match self {
+            Cursor::Replace(node) | Cursor::Before(node) => node,
+        }
+    }
+}
+
+pub(crate) fn render_snippet(_cap: SnippetCap, node: &SyntaxNode, cursor: Cursor) -> String {
+    assert!(cursor.node().ancestors().any(|it| it == *node));
+    let range = cursor.node().text_range() - node.text_range().start();
     let range: ops::Range<usize> = range.into();
 
-    let mut placeholder = placeholder.to_string();
+    let mut placeholder = cursor.node().to_string();
     escape(&mut placeholder);
-    let tab_stop = format!("${{0:{}}}", placeholder);
+    let tab_stop = match cursor {
+        Cursor::Replace(placeholder) => format!("${{0:{}}}", placeholder),
+        Cursor::Before(placeholder) => format!("$0{}", placeholder),
+    };
 
     let mut buf = node.to_string();
     buf.replace_range(range, &tab_stop);