ide_assists/handlers/
add_label_to_loop.rs

1use ide_db::{
2    source_change::SourceChangeBuilder, syntax_helpers::node_ext::for_each_break_and_continue_expr,
3};
4use syntax::{
5    SyntaxToken, T,
6    ast::{
7        self, AstNode, HasLoopBody,
8        make::{self, tokens},
9        syntax_factory::SyntaxFactory,
10    },
11    syntax_editor::{Position, SyntaxEditor},
12};
13
14use crate::{AssistContext, AssistId, Assists};
15
16// Assist: add_label_to_loop
17//
18// Adds a label to a loop.
19//
20// ```
21// fn main() {
22//     loop$0 {
23//         break;
24//         continue;
25//     }
26// }
27// ```
28// ->
29// ```
30// fn main() {
31//     ${1:'l}: loop {
32//         break ${2:'l};
33//         continue ${0:'l};
34//     }
35// }
36// ```
37pub(crate) fn add_label_to_loop(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
38    let loop_kw = ctx.find_token_syntax_at_offset(T![loop])?;
39    let loop_expr = loop_kw.parent().and_then(ast::LoopExpr::cast)?;
40    if loop_expr.label().is_some() {
41        return None;
42    }
43
44    acc.add(
45        AssistId::generate("add_label_to_loop"),
46        "Add Label",
47        loop_expr.syntax().text_range(),
48        |builder| {
49            let make = SyntaxFactory::with_mappings();
50            let mut editor = builder.make_editor(loop_expr.syntax());
51
52            let label = make.lifetime("'l");
53            let elements = vec![
54                label.syntax().clone().into(),
55                make::token(T![:]).into(),
56                tokens::single_space().into(),
57            ];
58            editor.insert_all(Position::before(&loop_kw), elements);
59
60            if let Some(cap) = ctx.config.snippet_cap {
61                editor.add_annotation(label.syntax(), builder.make_placeholder_snippet(cap));
62            }
63
64            let loop_body = loop_expr.loop_body().and_then(|it| it.stmt_list());
65            for_each_break_and_continue_expr(loop_expr.label(), loop_body, &mut |expr| {
66                let token = match expr {
67                    ast::Expr::BreakExpr(break_expr) => break_expr.break_token(),
68                    ast::Expr::ContinueExpr(continue_expr) => continue_expr.continue_token(),
69                    _ => return,
70                };
71                if let Some(token) = token {
72                    insert_label_after_token(&mut editor, &make, &token, ctx, builder);
73                }
74            });
75
76            editor.add_mappings(make.finish_with_mappings());
77            builder.add_file_edits(ctx.vfs_file_id(), editor);
78            builder.rename();
79        },
80    )
81}
82
83fn insert_label_after_token(
84    editor: &mut SyntaxEditor,
85    make: &SyntaxFactory,
86    token: &SyntaxToken,
87    ctx: &AssistContext<'_>,
88    builder: &mut SourceChangeBuilder,
89) {
90    let label = make.lifetime("'l");
91    let elements = vec![tokens::single_space().into(), label.syntax().clone().into()];
92    editor.insert_all(Position::after(token), elements);
93
94    if let Some(cap) = ctx.config.snippet_cap {
95        editor.add_annotation(label.syntax(), builder.make_placeholder_snippet(cap));
96    }
97}
98
99#[cfg(test)]
100mod tests {
101    use crate::tests::{check_assist, check_assist_not_applicable};
102
103    use super::*;
104
105    #[test]
106    fn add_label() {
107        check_assist(
108            add_label_to_loop,
109            r#"
110fn main() {
111    loop$0 {
112        break;
113        continue;
114    }
115}"#,
116            r#"
117fn main() {
118    ${1:'l}: loop {
119        break ${2:'l};
120        continue ${0:'l};
121    }
122}"#,
123        );
124    }
125
126    #[test]
127    fn add_label_to_outer_loop() {
128        check_assist(
129            add_label_to_loop,
130            r#"
131fn main() {
132    loop$0 {
133        break;
134        continue;
135        loop {
136            break;
137            continue;
138        }
139    }
140}"#,
141            r#"
142fn main() {
143    ${1:'l}: loop {
144        break ${2:'l};
145        continue ${0:'l};
146        loop {
147            break;
148            continue;
149        }
150    }
151}"#,
152        );
153    }
154
155    #[test]
156    fn add_label_to_inner_loop() {
157        check_assist(
158            add_label_to_loop,
159            r#"
160fn main() {
161    loop {
162        break;
163        continue;
164        loop$0 {
165            break;
166            continue;
167        }
168    }
169}"#,
170            r#"
171fn main() {
172    loop {
173        break;
174        continue;
175        ${1:'l}: loop {
176            break ${2:'l};
177            continue ${0:'l};
178        }
179    }
180}"#,
181        );
182    }
183
184    #[test]
185    fn do_not_add_label_if_exists() {
186        check_assist_not_applicable(
187            add_label_to_loop,
188            r#"
189fn main() {
190    'l: loop$0 {
191        break 'l;
192        continue 'l;
193    }
194}"#,
195        );
196    }
197}
ide_assists/handlers/add_label_to_loop.rs

ide_assists/handlers/
add_label_to_loop.rs
