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