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 = while_kw.parent().and_then(ast::WhileExpr::cast)?;
  43     let while_body = while_expr.loop_body()?;
  44     let while_cond = while_expr.condition()?;
  45
  46     let if_cond = match while_cond.pat() {
  47         Some(_) => while_expr.condition()?,
  48         None => make::condition(invert_boolean_expression(while_cond.expr()?), None),
  49     };
  50     let target = while_expr.syntax().text_range();
  51     acc.add(
  52         AssistId("convert_while_to_loop", AssistKind::RefactorRewrite),
  53         "Convert while to loop",
  54         target,
  55         |edit| {
  56             let while_indent_level = IndentLevel::from_node(while_expr.syntax());
  57
  58             let break_block =
  59                 make::block_expr(once(make::expr_stmt(make::expr_break(None)).into()), None)
  60                     .indent(while_indent_level);
  61             let block_expr = match while_cond.pat() {
  62                 Some(_) => {
  63                     let if_expr = make::expr_if(if_cond, while_body, Some(break_block.into()));
  64                     let stmts = once(make::expr_stmt(if_expr).into());
  65                     make::block_expr(stmts, None)
  66                 }
  67                 None => {
  68                     let if_expr = make::expr_if(if_cond, break_block, None);
  69                     let stmts =
  70                         once(make::expr_stmt(if_expr).into()).chain(while_body.statements());
  71                     make::block_expr(stmts, while_body.tail_expr())
  72                 }
  73             };
  74
  75             let replacement = make::expr_loop(block_expr.indent(while_indent_level));
  76             edit.replace(target, replacement.syntax().text())
  77         },
  78     )
  79 }
  80
  81 #[cfg(test)]
  82 mod tests {
  83     use crate::tests::{check_assist, check_assist_not_applicable};
  84
  85     use super::*;
  86
  87     #[test]
  88     fn convert_inside_fn() {
  89         check_assist(
  90             convert_while_to_loop,
  91             r#"
  92 fn main() {
  93     while$0 cond {
  94         foo();
  95     }
  96 }
  97 "#,
  98             r#"
  99 fn main() {
 100     loop {
 101         if !cond {
 102             break;
 103         }
 104         foo();
 105     }
 106 }
 107 "#,
 108         );
 109     }
 110
 111     #[test]
 112     fn convert_busy_wait() {
 113         check_assist(
 114             convert_while_to_loop,
 115             r#"
 116 fn main() {
 117     while$0 cond() {}
 118 }
 119 "#,
 120             r#"
 121 fn main() {
 122     loop {
 123         if !cond() {
 124             break;
 125         }
 126     }
 127 }
 128 "#,
 129         );
 130     }
 131
 132     #[test]
 133     fn convert_trailing_expr() {
 134         check_assist(
 135             convert_while_to_loop,
 136             r#"
 137 fn main() {
 138     while$0 cond() {
 139         bar()
 140     }
 141 }
 142 "#,
 143             r#"
 144 fn main() {
 145     loop {
 146         if !cond() {
 147             break;
 148         }
 149         bar()
 150     }
 151 }
 152 "#,
 153         );
 154     }
 155
 156     #[test]
 157     fn convert_while_let() {
 158         check_assist(
 159             convert_while_to_loop,
 160             r#"
 161 fn main() {
 162     while$0 let Some(_) = foo() {
 163         bar();
 164     }
 165 }
 166 "#,
 167             r#"
 168 fn main() {
 169     loop {
 170         if let Some(_) = foo() {
 171             bar();
 172         } else {
 173             break;
 174         }
 175     }
 176 }
 177 "#,
 178         );
 179     }
 180
 181     #[test]
 182     fn ignore_cursor_in_body() {
 183         check_assist_not_applicable(
 184             convert_while_to_loop,
 185             r#"
 186 fn main() {
 187     while cond {$0
 188         bar();
 189     }
 190 }
 191 "#,
 192         );
 193     }
 194 }