crates/ra_assists/src/handlers/replace_unwrap_with_match.rs

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