src/tools/rust-analyzer/crates/ide-assists/src/handlers/add_label_to_loop.rs

   1 use ide_db::syntax_helpers::node_ext::for_each_break_and_continue_expr;
   2 use syntax::{
   3     ast::{self, AstNode, HasLoopBody},
   4     T,
   5 };
   6
   7 use crate::{AssistContext, AssistId, AssistKind, Assists};
   8
   9 // Assist: add_label_to_loop
  10 //
  11 // Adds a label to a loop.
  12 //
  13 // ```
  14 // fn main() {
  15 //     loop$0 {
  16 //         break;
  17 //         continue;
  18 //     }
  19 // }
  20 // ```
  21 // ->
  22 // ```
  23 // fn main() {
  24 //     'l: loop {
  25 //         break 'l;
  26 //         continue 'l;
  27 //     }
  28 // }
  29 // ```
  30 pub(crate) fn add_label_to_loop(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
  31     let loop_kw = ctx.find_token_syntax_at_offset(T![loop])?;
  32     let loop_expr = loop_kw.parent().and_then(ast::LoopExpr::cast)?;
  33     if loop_expr.label().is_some() {
  34         return None;
  35     }
  36
  37     acc.add(
  38         AssistId("add_label_to_loop", AssistKind::Generate),
  39         "Add Label",
  40         loop_expr.syntax().text_range(),
  41         |builder| {
  42             builder.insert(loop_kw.text_range().start(), "'l: ");
  43
  44             let loop_body = loop_expr.loop_body().and_then(|it| it.stmt_list());
  45             for_each_break_and_continue_expr(
  46                 loop_expr.label(),
  47                 loop_body,
  48                 &mut |expr| match expr {
  49                     ast::Expr::BreakExpr(break_expr) => {
  50                         if let Some(break_token) = break_expr.break_token() {
  51                             builder.insert(break_token.text_range().end(), " 'l")
  52                         }
  53                     }
  54                     ast::Expr::ContinueExpr(continue_expr) => {
  55                         if let Some(continue_token) = continue_expr.continue_token() {
  56                             builder.insert(continue_token.text_range().end(), " 'l")
  57                         }
  58                     }
  59                     _ => {}
  60                 },
  61             );
  62         },
  63     )
  64 }
  65
  66 #[cfg(test)]
  67 mod tests {
  68     use crate::tests::{check_assist, check_assist_not_applicable};
  69
  70     use super::*;
  71
  72     #[test]
  73     fn add_label() {
  74         check_assist(
  75             add_label_to_loop,
  76             r#"
  77 fn main() {
  78     loop$0 {
  79         break;
  80         continue;
  81     }
  82 }"#,
  83             r#"
  84 fn main() {
  85     'l: loop {
  86         break 'l;
  87         continue 'l;
  88     }
  89 }"#,
  90         );
  91     }
  92
  93     #[test]
  94     fn add_label_to_outer_loop() {
  95         check_assist(
  96             add_label_to_loop,
  97             r#"
  98 fn main() {
  99     loop$0 {
 100         break;
 101         continue;
 102         loop {
 103             break;
 104             continue;
 105         }
 106     }
 107 }"#,
 108             r#"
 109 fn main() {
 110     'l: loop {
 111         break 'l;
 112         continue 'l;
 113         loop {
 114             break;
 115             continue;
 116         }
 117     }
 118 }"#,
 119         );
 120     }
 121
 122     #[test]
 123     fn add_label_to_inner_loop() {
 124         check_assist(
 125             add_label_to_loop,
 126             r#"
 127 fn main() {
 128     loop {
 129         break;
 130         continue;
 131         loop$0 {
 132             break;
 133             continue;
 134         }
 135     }
 136 }"#,
 137             r#"
 138 fn main() {
 139     loop {
 140         break;
 141         continue;
 142         'l: loop {
 143             break 'l;
 144             continue 'l;
 145         }
 146     }
 147 }"#,
 148         );
 149     }
 150
 151     #[test]
 152     fn do_not_add_label_if_exists() {
 153         check_assist_not_applicable(
 154             add_label_to_loop,
 155             r#"
 156 fn main() {
 157     'l: loop$0 {
 158         break 'l;
 159         continue 'l;
 160     }
 161 }"#,
 162         );
 163     }
 164 }