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, Assist, AssistCtx, AssistId};
   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(ctx: AssistCtx) -> Option<Assist> {
  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     ctx.add_assist(
  43         AssistId("replace_unwrap_with_match"),
  44         "Replace unwrap with match",
  45         target,
  46         |edit| {
  47             let ok_path = make::path_unqualified(make::path_segment(make::name_ref(happy_variant)));
  48             let it = make::bind_pat(make::name("a")).into();
  49             let ok_tuple = make::tuple_struct_pat(ok_path, iter::once(it)).into();
  50
  51             let bind_path = make::path_unqualified(make::path_segment(make::name_ref("a")));
  52             let ok_arm = make::match_arm(iter::once(ok_tuple), make::expr_path(bind_path));
  53
  54             let unreachable_call = make::unreachable_macro_call().into();
  55             let err_arm =
  56                 make::match_arm(iter::once(make::placeholder_pat().into()), unreachable_call);
  57
  58             let match_arm_list = make::match_arm_list(vec![ok_arm, err_arm]);
  59             let match_expr = make::expr_match(caller.clone(), match_arm_list);
  60             let match_expr =
  61                 IndentLevel::from_node(method_call.syntax()).increase_indent(match_expr);
  62
  63             edit.set_cursor(caller.syntax().text_range().start());
  64             edit.replace_ast::<ast::Expr>(method_call.into(), match_expr);
  65         },
  66     )
  67 }
  68
  69 #[cfg(test)]
  70 mod tests {
  71     use super::*;
  72     use crate::tests::{check_assist, check_assist_target};
  73
  74     #[test]
  75     fn test_replace_result_unwrap_with_match() {
  76         check_assist(
  77             replace_unwrap_with_match,
  78             r"
  79 enum Result<T, E> { Ok(T), Err(E) }
  80 fn i<T>(a: T) -> T { a }
  81 fn main() {
  82     let x: Result<i32, i32> = Result::Ok(92);
  83     let y = i(x).<|>unwrap();
  84 }
  85             ",
  86             r"
  87 enum Result<T, E> { Ok(T), Err(E) }
  88 fn i<T>(a: T) -> T { a }
  89 fn main() {
  90     let x: Result<i32, i32> = Result::Ok(92);
  91     let y = <|>match i(x) {
  92         Ok(a) => a,
  93         _ => unreachable!(),
  94     };
  95 }
  96             ",
  97         )
  98     }
  99
 100     #[test]
 101     fn test_replace_option_unwrap_with_match() {
 102         check_assist(
 103             replace_unwrap_with_match,
 104             r"
 105 enum Option<T> { Some(T), None }
 106 fn i<T>(a: T) -> T { a }
 107 fn main() {
 108     let x = Option::Some(92);
 109     let y = i(x).<|>unwrap();
 110 }
 111             ",
 112             r"
 113 enum Option<T> { Some(T), None }
 114 fn i<T>(a: T) -> T { a }
 115 fn main() {
 116     let x = Option::Some(92);
 117     let y = <|>match i(x) {
 118         Some(a) => a,
 119         _ => unreachable!(),
 120     };
 121 }
 122             ",
 123         );
 124     }
 125
 126     #[test]
 127     fn test_replace_result_unwrap_with_match_chaining() {
 128         check_assist(
 129             replace_unwrap_with_match,
 130             r"
 131 enum Result<T, E> { Ok(T), Err(E) }
 132 fn i<T>(a: T) -> T { a }
 133 fn main() {
 134     let x: Result<i32, i32> = Result::Ok(92);
 135     let y = i(x).<|>unwrap().count_zeroes();
 136 }
 137             ",
 138             r"
 139 enum Result<T, E> { Ok(T), Err(E) }
 140 fn i<T>(a: T) -> T { a }
 141 fn main() {
 142     let x: Result<i32, i32> = Result::Ok(92);
 143     let y = <|>match i(x) {
 144         Ok(a) => a,
 145         _ => unreachable!(),
 146     }.count_zeroes();
 147 }
 148             ",
 149         )
 150     }
 151
 152     #[test]
 153     fn replace_unwrap_with_match_target() {
 154         check_assist_target(
 155             replace_unwrap_with_match,
 156             r"
 157 enum Option<T> { Some(T), None }
 158 fn i<T>(a: T) -> T { a }
 159 fn main() {
 160     let x = Option::Some(92);
 161     let y = i(x).<|>unwrap();
 162 }
 163             ",
 164             r"i(x).unwrap()",
 165         );
 166     }
 167 }