crates/ide_assists/src/handlers/convert_while_to_loop.rs

   1 use std::iter::once;
   2
   3 use syntax::{
   4     ast::{
   5         self,
   6         edit::{AstNodeEdit, IndentLevel},
   7         make, HasLoopBody,
   8     },
   9     AstNode, T,
  10 };
  11
  12 use crate::{
  13     assist_context::{AssistContext, Assists},
  14     utils::invert_boolean_expression,
  15     AssistId, AssistKind,
  16 };
  17
  18 // Assist: convert_while_to_loop
  19 //
  20 // Replace a while with a loop.
  21 //
  22 // ```
  23 // fn main() {
  24 //     $0while cond {
  25 //         foo();
  26 //     }
  27 // }
  28 // ```
  29 // ->
  30 // ```
  31 // fn main() {
  32 //     loop {
  33 //         if !cond {
  34 //             break;
  35 //         }
  36 //         foo();
  37 //     }
  38 // }
  39 // ```
  40 pub(crate) fn convert_while_to_loop(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
  41     let while_kw = ctx.find_token_syntax_at_offset(T![while])?;
  42     let while_expr: ast::WhileExpr = while_kw.parent().and_then(ast::WhileExpr::cast)?;
  43     let while_body = while_expr.loop_body()?;
  44     let cond = while_expr.condition()?;
  45
  46     // Don't handle while let
  47     if cond.pat().is_some() {
  48         return None;
  49     };
  50
  51     let cond_expr = cond.expr()?;
  52
  53     let target = while_expr.syntax().text_range();
  54     acc.add(
  55         AssistId("convert_while_to_loop", AssistKind::RefactorRewrite),
  56         "Convert while to loop",
  57         target,
  58         |edit| {
  59             let while_indent_level = IndentLevel::from_node(while_expr.syntax());
  60
  61             let replacement = {
  62                 let if_expr = {
  63                     let cond = invert_boolean_expression(cond_expr);
  64                     let then_branch = make::block_expr(
  65                         once(make::expr_stmt(make::expr_break(None)).into()),
  66                         None,
  67                     );
  68
  69                     make::expr_if(make::condition(cond, None), then_branch, None)
  70                 };
  71
  72                 let if_expr = if_expr.indent(while_indent_level);
  73                 let stmts = once(make::expr_stmt(if_expr).into()).chain(while_body.statements());
  74
  75                 let block_expr = make::block_expr(stmts, while_body.tail_expr());
  76
  77                 let block_expr = block_expr.indent(while_indent_level);
  78
  79                 make::expr_loop(block_expr)
  80             };
  81
  82             edit.replace(target, replacement.syntax().text())
  83         },
  84     )
  85 }
  86
  87 #[cfg(test)]
  88 mod tests {
  89     use crate::tests::{check_assist, check_assist_not_applicable};
  90
  91     use super::*;
  92
  93     #[test]
  94     fn convert_inside_fn() {
  95         check_assist(
  96             convert_while_to_loop,
  97             r#"
  98 fn main() {
  99     while$0 cond {
 100         foo();
 101     }
 102 }
 103 "#,
 104             r#"
 105 fn main() {
 106     loop {
 107         if !cond {
 108             break;
 109         }
 110         foo();
 111     }
 112 }
 113 "#,
 114         );
 115     }
 116
 117     #[test]
 118     fn convert_busy_wait() {
 119         check_assist(
 120             convert_while_to_loop,
 121             r#"
 122 fn main() {
 123     while$0 cond() {}
 124 }
 125 "#,
 126             r#"
 127 fn main() {
 128     loop {
 129         if !cond() {
 130             break;
 131         }
 132     }
 133 }
 134 "#,
 135         );
 136     }
 137
 138     #[test]
 139     fn convert_trailing_expr() {
 140         check_assist(
 141             convert_while_to_loop,
 142             r#"
 143 fn main() {
 144     while$0 cond() {
 145         bar()
 146     }
 147 }
 148 "#,
 149             r#"
 150 fn main() {
 151     loop {
 152         if !cond() {
 153             break;
 154         }
 155         bar()
 156     }
 157 }
 158 "#,
 159         );
 160     }
 161
 162     #[test]
 163     fn ignore_while_let() {
 164         check_assist_not_applicable(
 165             convert_while_to_loop,
 166             r#"
 167 fn main() {
 168     while$0 let Some(_) = foo() {
 169         bar();
 170     }
 171 }
 172 "#,
 173         );
 174     }
 175
 176     #[test]
 177     fn ignore_cursor_in_body() {
 178         check_assist_not_applicable(
 179             convert_while_to_loop,
 180             r#"
 181 fn main() {
 182     while cond {$0
 183         bar();
 184     }
 185 }
 186 "#,
 187         );
 188     }
 189 }