1 files changed, 44 insertions, 54 deletions

diff --git a/crates/ra_assists/src/handlers/change_return_type_to_result.rs b/crates/ra_assists/src/handlers/change_return_type_to_result.rs
index 5c907097e..c6baa0a57 100644
--- a/crates/ra_assists/src/handlers/change_return_type_to_result.rs
+++ b/crates/ra_assists/src/handlers/change_return_type_to_result.rs
@@ -1,8 +1,6 @@
 use ra_syntax::{
     ast::{self, BlockExpr, Expr, LoopBodyOwner},
-    AstNode,
-    SyntaxKind::{COMMENT, WHITESPACE},
-    SyntaxNode, TextSize,
+    AstNode, SyntaxNode,
 };
 
 use crate::{AssistContext, AssistId, Assists};
@@ -16,39 +14,40 @@ use crate::{AssistContext, AssistId, Assists};
 // ```
 // ->
 // ```
-// fn foo() -> Result<i32, > { Ok(42i32) }
+// fn foo() -> Result<i32, ${0:_}> { Ok(42i32) }
 // ```
 pub(crate) fn change_return_type_to_result(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let fn_def = ctx.find_node_at_offset::<ast::FnDef>();
-    let fn_def = &mut fn_def?;
-    let ret_type = &fn_def.ret_type()?.type_ref()?;
-    if ret_type.syntax().text().to_string().starts_with("Result<") {
+    let ret_type = ctx.find_node_at_offset::<ast::RetType>()?;
+    // FIXME: extend to lambdas as well
+    let fn_def = ret_type.syntax().parent().and_then(ast::FnDef::cast)?;
+
+    let type_ref = &ret_type.type_ref()?;
+    if type_ref.syntax().text().to_string().starts_with("Result<") {
         return None;
     }
 
     let block_expr = &fn_def.body()?;
-    let cursor_in_ret_type =
-        fn_def.ret_type()?.syntax().text_range().contains_range(ctx.frange.range);
-    if !cursor_in_ret_type {
-        return None;
-    }
 
     acc.add(
         AssistId("change_return_type_to_result"),
         "Change return type to Result",
-        ret_type.syntax().text_range(),
-        |edit| {
+        type_ref.syntax().text_range(),
+        |builder| {
             let mut tail_return_expr_collector = TailReturnCollector::new();
             tail_return_expr_collector.collect_jump_exprs(block_expr, false);
             tail_return_expr_collector.collect_tail_exprs(block_expr);
 
             for ret_expr_arg in tail_return_expr_collector.exprs_to_wrap {
-                edit.replace_node_and_indent(&ret_expr_arg, format!("Ok({})", ret_expr_arg));
+                builder.replace_node_and_indent(&ret_expr_arg, format!("Ok({})", ret_expr_arg));
             }
-            edit.replace_node_and_indent(ret_type.syntax(), format!("Result<{}, >", ret_type));
 
-            if let Some(node_start) = result_insertion_offset(&ret_type) {
-                edit.set_cursor(node_start + TextSize::of(&format!("Result<{}, ", ret_type)));
+            match ctx.config.snippet_cap {
+                Some(cap) => {
+                    let snippet = format!("Result<{}, ${{0:_}}>", type_ref);
+                    builder.replace_snippet(cap, type_ref.syntax().text_range(), snippet)
+                }
+                None => builder
+                    .replace(type_ref.syntax().text_range(), format!("Result<{}, _>", type_ref)),
             }
         },
     )
@@ -250,17 +249,8 @@ fn get_tail_expr_from_block(expr: &Expr) -> Option<Vec<NodeType>> {
     }
 }
 
-fn result_insertion_offset(ret_type: &ast::TypeRef) -> Option<TextSize> {
-    let non_ws_child = ret_type
-        .syntax()
-        .children_with_tokens()
-        .find(|it| it.kind() != COMMENT && it.kind() != WHITESPACE)?;
-    Some(non_ws_child.text_range().start())
-}
-
 #[cfg(test)]
 mod tests {
-
     use crate::tests::{check_assist, check_assist_not_applicable};
 
     use super::*;
@@ -273,7 +263,7 @@ mod tests {
                 let test = "test";
                 return 42i32;
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 let test = "test";
                 return Ok(42i32);
             }"#,
@@ -288,7 +278,7 @@ mod tests {
                 let test = "test";
                 return 42i32;
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 let test = "test";
                 return Ok(42i32);
             }"#,
@@ -314,7 +304,7 @@ mod tests {
                 let test = "test";
                 return 42i32;
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 let test = "test";
                 return Ok(42i32);
             }"#,
@@ -329,7 +319,7 @@ mod tests {
                 let test = "test";
                 42i32
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 let test = "test";
                 Ok(42i32)
             }"#,
@@ -343,7 +333,7 @@ mod tests {
             r#"fn foo() -> i32<|> {
                 42i32
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 Ok(42i32)
             }"#,
         );
@@ -359,7 +349,7 @@ mod tests {
                     24i32
                 }
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 if true {
                     Ok(42i32)
                 } else {
@@ -384,7 +374,7 @@ mod tests {
                     24i32
                 }
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 if true {
                     if false {
                         Ok(1)
@@ -413,7 +403,7 @@ mod tests {
                     24i32.await
                 }
             }"#,
-            r#"async fn foo() -> Result<i32, <|>> {
+            r#"async fn foo() -> Result<i32, ${0:_}> {
                 if true {
                     if false {
                         Ok(1.await)
@@ -434,7 +424,7 @@ mod tests {
             r#"fn foo() -> [i32;<|> 3] {
                 [1, 2, 3]
             }"#,
-            r#"fn foo() -> Result<[i32; 3], <|>> {
+            r#"fn foo() -> Result<[i32; 3], ${0:_}> {
                 Ok([1, 2, 3])
             }"#,
         );
@@ -455,7 +445,7 @@ mod tests {
                     24 as i32
                 }
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 if true {
                     if false {
                         Ok(1 as i32)
@@ -480,7 +470,7 @@ mod tests {
                     _ => 24i32,
                 }
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 let my_var = 5;
                 match my_var {
                     5 => Ok(42i32),
@@ -503,7 +493,7 @@ mod tests {
 
                 my_var
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 let my_var = 5;
                 loop {
                     println!("test");
@@ -526,7 +516,7 @@ mod tests {
 
                 my_var
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 let my_var = let x = loop {
                     break 1;
                 };
@@ -549,7 +539,7 @@ mod tests {
 
                 res
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 let my_var = 5;
                 let res = match my_var {
                     5 => 42i32,
@@ -572,7 +562,7 @@ mod tests {
 
                 res
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 let my_var = 5;
                 let res = if my_var == 5 {
                     42i32
@@ -608,7 +598,7 @@ mod tests {
                     },
                 }
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 let my_var = 5;
                 match my_var {
                     5 => {
@@ -641,7 +631,7 @@ mod tests {
                 }
                 53i32
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 let test = "test";
                 if test == "test" {
                     return Ok(24i32);
@@ -672,7 +662,7 @@ mod tests {
 
                 the_field
             }"#,
-            r#"fn foo(the_field: u32) -> Result<u32, <|>> {
+            r#"fn foo(the_field: u32) -> Result<u32, ${0:_}> {
                 let true_closure = || {
                     return true;
                 };
@@ -711,7 +701,7 @@ mod tests {
 
                 t.unwrap_or_else(|| the_field)
             }"#,
-            r#"fn foo(the_field: u32) -> Result<u32, <|>> {
+            r#"fn foo(the_field: u32) -> Result<u32, ${0:_}> {
                 let true_closure = || {
                     return true;
                 };
@@ -749,7 +739,7 @@ mod tests {
                     i += 1;
                 }
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 let test = "test";
                 if test == "test" {
                     return Ok(24i32);
@@ -781,7 +771,7 @@ mod tests {
                     }
                 }
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 let test = "test";
                 if test == "test" {
                     return Ok(24i32);
@@ -819,7 +809,7 @@ mod tests {
                     }
                 }
             }"#,
-            r#"fn foo() -> Result<i32, <|>> {
+            r#"fn foo() -> Result<i32, ${0:_}> {
                 let test = "test";
                 let other = 5;
                 if test == "test" {
@@ -860,7 +850,7 @@ mod tests {
 
                 the_field
             }"#,
-            r#"fn foo(the_field: u32) -> Result<u32, <|>> {
+            r#"fn foo(the_field: u32) -> Result<u32, ${0:_}> {
                 if the_field < 5 {
                     let mut i = 0;
                     loop {
@@ -894,7 +884,7 @@ mod tests {
 
                 the_field
             }"#,
-            r#"fn foo(the_field: u32) -> Result<u32, <|>> {
+            r#"fn foo(the_field: u32) -> Result<u32, ${0:_}> {
                 if the_field < 5 {
                     let mut i = 0;
 
@@ -923,7 +913,7 @@ mod tests {
 
                 the_field
             }"#,
-            r#"fn foo(the_field: u32) -> Result<u32, <|>> {
+            r#"fn foo(the_field: u32) -> Result<u32, ${0:_}> {
                 if the_field < 5 {
                     let mut i = 0;
 
@@ -953,7 +943,7 @@ mod tests {
 
                 the_field
             }"#,
-            r#"fn foo(the_field: u32) -> Result<u32, <|>> {
+            r#"fn foo(the_field: u32) -> Result<u32, ${0:_}> {
                 if the_field < 5 {
                     let mut i = 0;