crates/ra_assists/src/handlers/replace_let_with_if_let.rs

   1 use std::iter::once;
   2
   3 use ra_syntax::{
   4     ast::{
   5         self,
   6         edit::{AstNodeEdit, IndentLevel},
   7         make,
   8     },
   9     AstNode, T,
  10 };
  11
  12 use crate::{
  13     assist_ctx::{Assist, AssistCtx},
  14     utils::TryEnum,
  15     AssistId,
  16 };
  17
  18 // Assist: replace_let_with_if_let
  19 //
  20 // Replaces `let` with an `if-let`.
  21 //
  22 // ```
  23 // # enum Option<T> { Some(T), None }
  24 //
  25 // fn main(action: Action) {
  26 //     <|>let x = compute();
  27 // }
  28 //
  29 // fn compute() -> Option<i32> { None }
  30 // ```
  31 // ->
  32 // ```
  33 // # enum Option<T> { Some(T), None }
  34 //
  35 // fn main(action: Action) {
  36 //     if let Some(x) = compute() {
  37 //     }
  38 // }
  39 //
  40 // fn compute() -> Option<i32> { None }
  41 // ```
  42 pub(crate) fn replace_let_with_if_let(ctx: AssistCtx) -> Option<Assist> {
  43     let let_kw = ctx.find_token_at_offset(T![let])?;
  44     let let_stmt = let_kw.ancestors().find_map(ast::LetStmt::cast)?;
  45     let init = let_stmt.initializer()?;
  46     let original_pat = let_stmt.pat()?;
  47     let ty = ctx.sema.type_of_expr(&init)?;
  48     let happy_variant = TryEnum::from_ty(ctx.sema, &ty).map(|it| it.happy_case());
  49
  50     ctx.add_assist(AssistId("replace_let_with_if_let"), "Replace with if-let", |edit| {
  51         let with_placeholder: ast::Pat = match happy_variant {
  52             None => make::placeholder_pat().into(),
  53             Some(var_name) => make::tuple_struct_pat(
  54                 make::path_unqualified(make::path_segment(make::name_ref(var_name))),
  55                 once(make::placeholder_pat().into()),
  56             )
  57             .into(),
  58         };
  59         let block =
  60             IndentLevel::from_node(let_stmt.syntax()).increase_indent(make::block_expr(None, None));
  61         let if_ = make::expr_if(make::condition(init, Some(with_placeholder)), block);
  62         let stmt = make::expr_stmt(if_);
  63
  64         let placeholder = stmt.syntax().descendants().find_map(ast::PlaceholderPat::cast).unwrap();
  65         let target_offset =
  66             let_stmt.syntax().text_range().start() + placeholder.syntax().text_range().start();
  67         let stmt = stmt.replace_descendant(placeholder.into(), original_pat);
  68
  69         edit.replace_ast(ast::Stmt::from(let_stmt), ast::Stmt::from(stmt));
  70         edit.target(let_kw.text_range());
  71         edit.set_cursor(target_offset);
  72     })
  73 }
  74
  75 #[cfg(test)]
  76 mod tests {
  77     use crate::tests::check_assist;
  78
  79     use super::*;
  80
  81     #[test]
  82     fn replace_let_unknown_enum() {
  83         check_assist(
  84             replace_let_with_if_let,
  85             r"
  86 enum E<T> { X(T), Y(T) }
  87
  88 fn main() {
  89     <|>let x = E::X(92);
  90 }
  91             ",
  92             r"
  93 enum E<T> { X(T), Y(T) }
  94
  95 fn main() {
  96     if let <|>x = E::X(92) {
  97     }
  98 }
  99             ",
 100         )
 101     }
 102 }