support rename after adding label

Author: Natural-selection1

crates/ide-assists/src/handlers/add_label_to_loop.rs

@@ -1,7 +1,10 @@
-use ide_db::syntax_helpers::node_ext::for_each_break_and_continue_expr;
+use ide_db::{
+    source_change::SourceChangeBuilder, syntax_helpers::node_ext::for_each_break_and_continue_expr,
+};
 use syntax::{
-    T,
-    ast::{self, AstNode, HasLoopBody},
+    SyntaxToken, T,
+    ast::{self, AstNode, HasLoopBody, make, syntax_factory::SyntaxFactory},
+    syntax_editor::{Position, SyntaxEditor},
 };
 
 use crate::{AssistContext, AssistId, Assists};
@@ -39,30 +42,57 @@ pub(crate) fn add_label_to_loop(acc: &mut Assists, ctx: &AssistContext<'_>) -> O
         "Add Label",
         loop_expr.syntax().text_range(),
         |builder| {
-            builder.insert(loop_kw.text_range().start(), "'l: ");
+            let make = SyntaxFactory::with_mappings();
+            let mut editor = builder.make_editor(loop_expr.syntax());
+
+            let label = make.lifetime("'l");
+            let elements = vec![
+                label.syntax().clone().into(),
+                make::token(syntax::SyntaxKind::COLON).into(),
+                make::token(syntax::SyntaxKind::WHITESPACE).into(),
+            ];
+            editor.insert_all(Position::before(&loop_kw), elements);
+
+            if let Some(cap) = ctx.config.snippet_cap {
+                editor.add_annotation(label.syntax(), builder.make_placeholder_snippet(cap));
+            }
 
             let loop_body = loop_expr.loop_body().and_then(|it| it.stmt_list());
-            for_each_break_and_continue_expr(
-                loop_expr.label(),
-                loop_body,
-                &mut |expr| match expr {
-                    ast::Expr::BreakExpr(break_expr) => {
-                        if let Some(break_token) = break_expr.break_token() {
-                            builder.insert(break_token.text_range().end(), " 'l")
-                        }
-                    }
-                    ast::Expr::ContinueExpr(continue_expr) => {
-                        if let Some(continue_token) = continue_expr.continue_token() {
-                            builder.insert(continue_token.text_range().end(), " 'l")
-                        }
-                    }
-                    _ => {}
-                },
-            );
+            for_each_break_and_continue_expr(loop_expr.label(), loop_body, &mut |expr| {
+                let token = match expr {
+                    ast::Expr::BreakExpr(break_expr) => break_expr.break_token(),
+                    ast::Expr::ContinueExpr(continue_expr) => continue_expr.continue_token(),
+                    _ => return,
+                };
+                if let Some(token) = token {
+                    insert_lifetime_after_token(&mut editor, &make, &token, ctx, builder);
+                }
+            });
+
+            editor.add_mappings(make.finish_with_mappings());
+            builder.add_file_edits(ctx.vfs_file_id(), editor);
+            builder.rename();
         },
     )
 }
 
+fn insert_lifetime_after_token(
+    editor: &mut SyntaxEditor,
+    make: &SyntaxFactory,
+    token: &SyntaxToken,
+    ctx: &AssistContext<'_>,
+    builder: &mut SourceChangeBuilder,
+) {
+    let label = make.lifetime("'l");
+    let elements =
+        vec![make::token(syntax::SyntaxKind::WHITESPACE).into(), label.syntax().clone().into()];
+    editor.insert_all(Position::after(token), elements);
+
+    if let Some(cap) = ctx.config.snippet_cap {
+        editor.add_annotation(label.syntax(), builder.make_placeholder_snippet(cap));
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::tests::{check_assist, check_assist_not_applicable};
@@ -82,9 +112,9 @@ fn main() {
 }"#,
             r#"
 fn main() {
-    'l: loop {
-        break 'l;
-        continue 'l;
+    ${1:'l}: loop {
+        break ${2:'l};
+        continue ${0:'l};
     }
 }"#,
         );
@@ -107,9 +137,9 @@ fn main() {
 }"#,
             r#"
 fn main() {
-    'l: loop {
-        break 'l;
-        continue 'l;
+    ${1:'l}: loop {
+        break ${2:'l};
+        continue ${0:'l};
         loop {
             break;
             continue;
@@ -139,9 +169,9 @@ fn main() {
     loop {
         break;
         continue;
-        'l: loop {
-            break 'l;
-            continue 'l;
+        ${1:'l}: loop {
+            break ${2:'l};
+            continue ${0:'l};
         }
     }
 }"#,

crates/ide-assists/src/tests/generated.rs

@@ -183,9 +183,9 @@ fn main() {
 "#####,
         r#####"
 fn main() {
-    'l: loop {
-        break 'l;
-        continue 'l;
+    ${1:'l}: loop {
+        break ${2:'l};
+        continue ${0:'l};
     }
 }
 "#####,