Assist: replace unwrap with match

Author: unrealhoang

crates/ra_assists/src/doc_tests/generated.rs

@@ -622,6 +622,30 @@ fn process(map: HashMap<String, String>) {}
     )
 }
 
+#[test]
+fn doctest_replace_unwrap_with_match() {
+    check(
+        "replace_unwrap_with_match",
+        r#####"
+enum Result<T, E> { Ok(T), Err(E) }
+fn main() {
+    let x: Result<i32, i32> = Result::Ok(92);
+    let y = x.<|>unwrap();
+}
+"#####,
+        r#####"
+enum Result<T, E> { Ok(T), Err(E) }
+fn main() {
+    let x: Result<i32, i32> = Result::Ok(92);
+    let y = match x {
+        Ok(a) => a,
+        _ => unreachable!(),
+    };
+}
+"#####,
+    )
+}
+
 #[test]
 fn doctest_split_import() {
     check(

crates/ra_assists/src/handlers/replace_unwrap_with_match.rs

@@ -0,0 +1,188 @@
+use std::iter;
+
+use ra_syntax::{
+    ast::{self, make},
+    AstNode,
+};
+use hir::name;
+
+use crate::{Assist, AssistCtx, AssistId};
+use ast::edit::IndentLevel;
+
+// Assist: replace_unwrap_with_match
+//
+// Replaces `unwrap` a `match` expression. Works for Result and Option.
+//
+// ```
+// enum Result<T, E> { Ok(T), Err(E) }
+// fn main() {
+//     let x: Result<i32, i32> = Result::Ok(92);
+//     let y = x.<|>unwrap();
+// }
+// ```
+// ->
+// ```
+// enum Result<T, E> { Ok(T), Err(E) }
+// fn main() {
+//     let x: Result<i32, i32> = Result::Ok(92);
+//     let y = match x {
+//         Ok(a) => a,
+//         _ => unreachable!(),
+//     };
+// }
+// ```
+pub(crate) fn replace_unwrap_with_match(ctx: AssistCtx) -> Option<Assist> {
+    let method_call: ast::MethodCallExpr = ctx.find_node_at_offset()?;
+    let name = method_call.name_ref()?;
+    if name.text() != "unwrap" {
+        return None;
+    }
+    let caller = method_call.expr()?;
+    let ty = ctx.sema.type_of_expr(&caller)?;
+
+    let type_name = ty.as_adt()?.name(ctx.sema.db);
+    if type_name == name![Result] {
+        let ok_path = make::path_unqualified(make::path_segment(make::name_ref("Ok")));
+        let bind_pat = make::bind_pat(make::name("a")).into();
+        let ok_tuple = make::tuple_struct_pat(ok_path, iter::once(bind_pat)).into();
+        let bind_path = make::path_unqualified(make::path_segment(make::name_ref("a")));
+
+        let ok_arm = make::match_arm(iter::once(ok_tuple), make::expr_path(bind_path));
+
+        let unreachable_call = make::unreachable_macro_call().into();
+        let err_arm = make::match_arm(iter::once(make::placeholder_pat().into()), unreachable_call);
+        let match_arm_list = make::match_arm_list(vec![ok_arm, err_arm]);
+        let match_expr = make::expr_match(method_call.expr()?, match_arm_list);
+        let match_expr = IndentLevel::from_node(method_call.syntax()).increase_indent(match_expr);
+
+        return ctx.add_assist(AssistId("replace_unwrap_with_match"), "Replace unwrap with match", |edit| {
+            edit.target(method_call.syntax().text_range());
+            edit.set_cursor(caller.syntax().text_range().start());
+            edit.replace_ast::<ast::Expr>(method_call.into(), match_expr);
+        });
+    }
+
+    if type_name == name![Option] {
+        let some_path = make::path_unqualified(make::path_segment(make::name_ref("Some")));
+        let bind_pat = make::bind_pat(make::name("a")).into();
+        let some_tuple = make::tuple_struct_pat(some_path, iter::once(bind_pat)).into();
+        let bind_path = make::path_unqualified(make::path_segment(make::name_ref("a")));
+        let some_arm = make::match_arm(iter::once(some_tuple), make::expr_path(bind_path));
+
+        let unreachable_call = make::unreachable_macro_call().into();
+        let none_arm = make::match_arm(iter::once(make::placeholder_pat().into()), unreachable_call);
+        let match_arm_list = make::match_arm_list(vec![some_arm, none_arm]);
+        let match_expr = make::expr_match(method_call.expr()?, match_arm_list);
+        let match_expr = IndentLevel::from_node(method_call.syntax()).increase_indent(match_expr);
+
+        return ctx.add_assist(AssistId("replace_unwrap_with_match"), "Replace unwrap with match", |edit| {
+            edit.target(method_call.syntax().text_range());
+            edit.set_cursor(caller.syntax().text_range().start());
+            edit.replace_ast::<ast::Expr>(method_call.into(), match_expr);
+        });
+    }
+
+    None
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::helpers::{check_assist, check_assist_target};
+
+    #[test]
+    fn test_replace_result_unwrap_with_match() {
+        check_assist(
+            replace_unwrap_with_match,
+            "
+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 = i(x).<|>unwrap();
+}
+            ",
+            "
+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) {
+        Ok(a) => a,
+        _ => unreachable!(),
+    };
+}
+            ",
+        )
+    }
+
+    #[test]
+    fn test_replace_option_unwrap_with_match() {
+        check_assist(
+            replace_unwrap_with_match,
+            "
+enum Option<T> { Some(T), None }
+fn i<T>(a: T) -> T { a }
+fn main() {
+    let x = Option::Some(92);
+    let y = i(x).<|>unwrap();
+}
+            ",
+            "
+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) {
+        Some(a) => a,
+        _ => unreachable!(),
+    };
+}
+            ",
+        );
+    }
+
+    #[test]
+    fn test_replace_result_unwrap_with_match_chaining() {
+        check_assist(
+            replace_unwrap_with_match,
+            "
+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 = i(x).<|>unwrap().count_zeroes();
+}
+            ",
+            "
+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) {
+        Ok(a) => a,
+        _ => unreachable!(),
+    }.count_zeroes();
+}
+            ",
+        )
+    }
+
+
+    #[test]
+    fn replace_unwrap_with_match_target() {
+        check_assist_target(
+            replace_unwrap_with_match,
+            "
+enum Option<T> { Some(T), None }
+fn i<T>(a: T) -> T { a }
+fn main() {
+    let x = Option::Some(92);
+    let y = i(x).<|>unwrap();
+}
+            ",
+            "i(x).unwrap()"
+        );
+    }
+
+}

crates/ra_assists/src/lib.rs

@@ -119,6 +119,7 @@ mod handlers {
     mod remove_mut;
     mod replace_if_let_with_match;
     mod replace_qualified_name_with_use;
+    mod replace_unwrap_with_match;
     mod split_import;
 
     pub(crate) fn all() -> &'static [AssistHandler] {
@@ -154,6 +155,7 @@ mod handlers {
             remove_mut::remove_mut,
             replace_if_let_with_match::replace_if_let_with_match,
             replace_qualified_name_with_use::replace_qualified_name_with_use,
+            replace_unwrap_with_match::replace_unwrap_with_match,
             split_import::split_import,
         ]
     }

crates/ra_hir/src/lib.rs

@@ -70,6 +70,6 @@ pub use hir_def::{
     type_ref::Mutability,
 };
 pub use hir_expand::{
-    name::Name, HirFileId, InFile, MacroCallId, MacroCallLoc, MacroDefId, MacroFile, Origin,
+    name::Name, name::name, HirFileId, InFile, MacroCallId, MacroCallLoc, MacroDefId, MacroFile, Origin,
 };
 pub use hir_ty::{display::HirDisplay, CallableDef};

crates/ra_hir_expand/src/name.rs

@@ -155,6 +155,7 @@ pub mod known {
         Ok,
         Future,
         Result,
+        Option,
         Output,
         Target,
         Box,

crates/ra_syntax/src/ast/make.rs

@@ -250,6 +250,10 @@ pub fn token(kind: SyntaxKind) -> SyntaxToken {
         .unwrap_or_else(|| panic!("unhandled token: {:?}", kind))
 }
 
+pub fn unreachable_macro_call() -> ast::MacroCall {
+    ast_from_text(&format!("unreachable!()"))
+}
+
 fn ast_from_text<N: AstNode>(text: &str) -> N {
     let parse = SourceFile::parse(text);
     let node = parse.tree().syntax().descendants().find_map(N::cast).unwrap();

docs/user/assists.md

@@ -597,6 +597,29 @@ use std::collections::HashMap;
 fn process(map: HashMap<String, String>) {}
 ```
 
+## `replace_unwrap_with_match`
+
+Replaces `unwrap` a `match` expression. Works for Result and Option.
+
+```rust
+// BEFORE
+enum Result<T, E> { Ok(T), Err(E) }
+fn main() {
+    let x: Result<i32, i32> = Result::Ok(92);
+    let y = x.┃unwrap();
+}
+
+// AFTER
+enum Result<T, E> { Ok(T), Err(E) }
+fn main() {
+    let x: Result<i32, i32> = Result::Ok(92);
+    let y = match x {
+        Ok(a) => a,
+        _ => unreachable!(),
+    };
+}
+```
+
 ## `split_import`
 
 Wraps the tail of import into braces.