crates/ide_assists/src/handlers/replace_let_with_if_let.rs

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