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