crates/ra_assists/src/handlers/replace_unwrap_with_match.rs

   1 use std::iter;
   2
   3 use ra_syntax::{
   4     ast::{self, edit::IndentLevel, make},
   5     AstNode,
   6 };
   7
   8 use crate::{utils::TryEnum, AssistContext, AssistId, Assists};
   9
  10 // Assist: replace_unwrap_with_match
  11 //
  12 // Replaces `unwrap` a `match` expression. Works for Result and Option.
  13 //
  14 // ```
  15 // enum Result<T, E> { Ok(T), Err(E) }
  16 // fn main() {
  17 //     let x: Result<i32, i32> = Result::Ok(92);
  18 //     let y = x.<|>unwrap();
  19 // }
  20 // ```
  21 // ->
  22 // ```
  23 // enum Result<T, E> { Ok(T), Err(E) }
  24 // fn main() {
  25 //     let x: Result<i32, i32> = Result::Ok(92);
  26 //     let y = match x {
  27 //         Ok(a) => a,
  28 //         _ => unreachable!(),
  29 //     };
  30 // }
  31 // ```
  32 pub(crate) fn replace_unwrap_with_match(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  33     let method_call: ast::MethodCallExpr = ctx.find_node_at_offset()?;
  34     let name = method_call.name_ref()?;
  35     if name.text() != "unwrap" {
  36         return None;
  37     }
  38     let caller = method_call.expr()?;
  39     let ty = ctx.sema.type_of_expr(&caller)?;
  40     let happy_variant = TryEnum::from_ty(&ctx.sema, &ty)?.happy_case();
  41     let target = method_call.syntax().text_range();
  42     acc.add(AssistId("replace_unwrap_with_match"), "Replace unwrap with match", target, |edit| {
  43         let ok_path = make::path_unqualified(make::path_segment(make::name_ref(happy_variant)));
  44         let it = make::bind_pat(make::name("a")).into();
  45         let ok_tuple = make::tuple_struct_pat(ok_path, iter::once(it)).into();
  46
  47         let bind_path = make::path_unqualified(make::path_segment(make::name_ref("a")));
  48         let ok_arm = make::match_arm(iter::once(ok_tuple), make::expr_path(bind_path));
  49
  50         let unreachable_call = make::unreachable_macro_call().into();
  51         let err_arm = make::match_arm(iter::once(make::placeholder_pat().into()), unreachable_call);
  52
  53         let match_arm_list = make::match_arm_list(vec![ok_arm, err_arm]);
  54         let match_expr = make::expr_match(caller.clone(), match_arm_list);
  55         let match_expr = IndentLevel::from_node(method_call.syntax()).increase_indent(match_expr);
  56
  57         edit.set_cursor(caller.syntax().text_range().start());
  58         edit.replace_ast::<ast::Expr>(method_call.into(), match_expr);
  59     })
  60 }
  61
  62 #[cfg(test)]
  63 mod tests {
  64     use super::*;
  65     use crate::tests::{check_assist, check_assist_target};
  66
  67     #[test]
  68     fn test_replace_result_unwrap_with_match() {
  69         check_assist(
  70             replace_unwrap_with_match,
  71             r"
  72 enum Result<T, E> { Ok(T), Err(E) }
  73 fn i<T>(a: T) -> T { a }
  74 fn main() {
  75     let x: Result<i32, i32> = Result::Ok(92);
  76     let y = i(x).<|>unwrap();
  77 }
  78             ",
  79             r"
  80 enum Result<T, E> { Ok(T), Err(E) }
  81 fn i<T>(a: T) -> T { a }
  82 fn main() {
  83     let x: Result<i32, i32> = Result::Ok(92);
  84     let y = <|>match i(x) {
  85         Ok(a) => a,
  86         _ => unreachable!(),
  87     };
  88 }
  89             ",
  90         )
  91     }
  92
  93     #[test]
  94     fn test_replace_option_unwrap_with_match() {
  95         check_assist(
  96             replace_unwrap_with_match,
  97             r"
  98 enum Option<T> { Some(T), None }
  99 fn i<T>(a: T) -> T { a }
 100 fn main() {
 101     let x = Option::Some(92);
 102     let y = i(x).<|>unwrap();
 103 }
 104             ",
 105             r"
 106 enum Option<T> { Some(T), None }
 107 fn i<T>(a: T) -> T { a }
 108 fn main() {
 109     let x = Option::Some(92);
 110     let y = <|>match i(x) {
 111         Some(a) => a,
 112         _ => unreachable!(),
 113     };
 114 }
 115             ",
 116         );
 117     }
 118
 119     #[test]
 120     fn test_replace_result_unwrap_with_match_chaining() {
 121         check_assist(
 122             replace_unwrap_with_match,
 123             r"
 124 enum Result<T, E> { Ok(T), Err(E) }
 125 fn i<T>(a: T) -> T { a }
 126 fn main() {
 127     let x: Result<i32, i32> = Result::Ok(92);
 128     let y = i(x).<|>unwrap().count_zeroes();
 129 }
 130             ",
 131             r"
 132 enum Result<T, E> { Ok(T), Err(E) }
 133 fn i<T>(a: T) -> T { a }
 134 fn main() {
 135     let x: Result<i32, i32> = Result::Ok(92);
 136     let y = <|>match i(x) {
 137         Ok(a) => a,
 138         _ => unreachable!(),
 139     }.count_zeroes();
 140 }
 141             ",
 142         )
 143     }
 144
 145     #[test]
 146     fn replace_unwrap_with_match_target() {
 147         check_assist_target(
 148             replace_unwrap_with_match,
 149             r"
 150 enum Option<T> { Some(T), None }
 151 fn i<T>(a: T) -> T { a }
 152 fn main() {
 153     let x = Option::Some(92);
 154     let y = i(x).<|>unwrap();
 155 }
 156             ",
 157             r"i(x).unwrap()",
 158         );
 159     }
 160 }