Snippetify unwrap -> match

1 files changed, 29 insertions, 12 deletions

diff --git a/crates/ra_assists/src/handlers/replace_unwrap_with_match.rs b/crates/ra_assists/src/handlers/replace_unwrap_with_match.rs
index b379b55a8..cff7dfb81 100644
--- a/crates/ra_assists/src/handlers/replace_unwrap_with_match.rs
+++ b/crates/ra_assists/src/handlers/replace_unwrap_with_match.rs
@@ -9,7 +9,10 @@ use ra_syntax::{
     AstNode,
 };
 
-use crate::{utils::TryEnum, AssistContext, AssistId, Assists};
+use crate::{
+    utils::{render_snippet, Cursor, TryEnum},
+    AssistContext, AssistId, Assists,
+};
 
 // Assist: replace_unwrap_with_match
 //
@@ -29,7 +32,7 @@ use crate::{utils::TryEnum, AssistContext, AssistId, Assists};
 //     let x: Result<i32, i32> = Result::Ok(92);
 //     let y = match x {
 //         Ok(a) => a,
-//         _ => unreachable!(),
+//         $0_ => unreachable!(),
 //     };
 // }
 // ```
@@ -43,7 +46,7 @@ pub(crate) fn replace_unwrap_with_match(acc: &mut Assists, ctx: &AssistContext)
     let ty = ctx.sema.type_of_expr(&caller)?;
     let happy_variant = TryEnum::from_ty(&ctx.sema, &ty)?.happy_case();
     let target = method_call.syntax().text_range();
-    acc.add(AssistId("replace_unwrap_with_match"), "Replace unwrap with match", target, |edit| {
+    acc.add(AssistId("replace_unwrap_with_match"), "Replace unwrap with match", target, |builder| {
         let ok_path = make::path_unqualified(make::path_segment(make::name_ref(happy_variant)));
         let it = make::bind_pat(make::name("a")).into();
         let ok_tuple = make::tuple_struct_pat(ok_path, iter::once(it)).into();
@@ -58,16 +61,30 @@ pub(crate) fn replace_unwrap_with_match(acc: &mut Assists, ctx: &AssistContext)
         let match_expr = make::expr_match(caller.clone(), match_arm_list)
             .indent(IndentLevel::from_node(method_call.syntax()));
 
-        edit.set_cursor(caller.syntax().text_range().start());
-        edit.replace_ast::<ast::Expr>(method_call.into(), match_expr);
+        let range = method_call.syntax().text_range();
+        match ctx.config.snippet_cap {
+            Some(cap) => {
+                let err_arm = match_expr
+                    .syntax()
+                    .descendants()
+                    .filter_map(ast::MatchArm::cast)
+                    .last()
+                    .unwrap();
+                let snippet =
+                    render_snippet(cap, match_expr.syntax(), Cursor::Before(err_arm.syntax()));
+                builder.replace_snippet(cap, range, snippet)
+            }
+            None => builder.replace(range, match_expr.to_string()),
+        }
     })
 }
 
 #[cfg(test)]
 mod tests {
-    use super::*;
     use crate::tests::{check_assist, check_assist_target};
 
+    use super::*;
+
     #[test]
     fn test_replace_result_unwrap_with_match() {
         check_assist(
@@ -85,9 +102,9 @@ enum Result<T, E> { Ok(T), Err(E) }
 fn i<T>(a: T) -> T { a }
 fn main() {
     let x: Result<i32, i32> = Result::Ok(92);
-    let y = <|>match i(x) {
+    let y = match i(x) {
         Ok(a) => a,
-        _ => unreachable!(),
+        $0_ => unreachable!(),
     };
 }
             ",
@@ -111,9 +128,9 @@ enum Option<T> { Some(T), None }
 fn i<T>(a: T) -> T { a }
 fn main() {
     let x = Option::Some(92);
-    let y = <|>match i(x) {
+    let y = match i(x) {
         Some(a) => a,
-        _ => unreachable!(),
+        $0_ => unreachable!(),
     };
 }
             ",
@@ -137,9 +154,9 @@ enum Result<T, E> { Ok(T), Err(E) }
 fn i<T>(a: T) -> T { a }
 fn main() {
     let x: Result<i32, i32> = Result::Ok(92);
-    let y = <|>match i(x) {
+    let y = match i(x) {
         Ok(a) => a,
-        _ => unreachable!(),
+        $0_ => unreachable!(),
     }.count_zeroes();
 }
             ",