Make extract_expressions_from_format_string only use snippets when available

Author: DropDemBits

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

@@ -1,4 +1,4 @@
-use crate::{AssistContext, Assists};
+use crate::{utils, AssistContext, Assists};
 use hir::DescendPreference;
 use ide_db::{
     assists::{AssistId, AssistKind},
@@ -8,8 +8,12 @@ use ide_db::{
     },
 };
 use itertools::Itertools;
-use stdx::format_to;
-use syntax::{ast, AstNode, AstToken, NodeOrToken, SyntaxKind::COMMA, TextRange};
+use syntax::{
+    ast::{self, make},
+    ted, AstNode, AstToken, NodeOrToken,
+    SyntaxKind::WHITESPACE,
+    T,
+};
 
 // Assist: extract_expressions_from_format_string
 //
@@ -34,6 +38,7 @@ pub(crate) fn extract_expressions_from_format_string(
 ) -> Option<()> {
     let fmt_string = ctx.find_token_at_offset::<ast::String>()?;
     let tt = fmt_string.syntax().parent().and_then(ast::TokenTree::cast)?;
+    let tt_delimiter = tt.left_delimiter_token()?.kind();
 
     let expanded_t = ast::String::cast(
         ctx.sema
@@ -61,81 +66,95 @@ pub(crate) fn extract_expressions_from_format_string(
         "Extract format expressions",
         tt.syntax().text_range(),
         |edit| {
-            let fmt_range = fmt_string.syntax().text_range();
-
-            // Replace old format string with new format string whose arguments have been extracted
-            edit.replace(fmt_range, new_fmt);
-
-            // Insert cursor at end of format string
-            edit.insert(fmt_range.end(), "$0");
+            let tt = edit.make_mut(tt);
 
             // Extract existing arguments in macro
-            let tokens =
-                tt.token_trees_and_tokens().collect_vec();
-
-            let mut existing_args: Vec<String> = vec![];
+            let tokens = tt.token_trees_and_tokens().collect_vec();
 
-            let mut current_arg = String::new();
-            if let [_opening_bracket, NodeOrToken::Token(format_string), _args_start_comma, tokens @ .., NodeOrToken::Token(end_bracket)] =
+            let existing_args = if let [_opening_bracket, NodeOrToken::Token(_format_string), _args_start_comma, tokens @ .., NodeOrToken::Token(_end_bracket)] =
                 tokens.as_slice()
             {
-                for t in tokens {
-                    match t {
-                        NodeOrToken::Node(n) => {
-                            format_to!(current_arg, "{n}");
-                        },
-                        NodeOrToken::Token(t) if t.kind() == COMMA => {
-                            existing_args.push(current_arg.trim().into());
-                            current_arg.clear();
-                        },
-                        NodeOrToken::Token(t) => {
-                            current_arg.push_str(t.text());
-                        },
-                    }
-                }
-                existing_args.push(current_arg.trim().into());
+                let args = tokens.split(|it| matches!(it, NodeOrToken::Token(t) if t.kind() == T![,])).map(|arg| {
+                    // Strip off leading and trailing whitespace tokens
+                    let arg = match arg.split_first() {
+                        Some((NodeOrToken::Token(t), rest)) if t.kind() == WHITESPACE => rest,
+                        _ => arg,
+                    };
+                    let arg = match arg.split_last() {
+                        Some((NodeOrToken::Token(t), rest)) if t.kind() == WHITESPACE => rest,
+                        _ => arg,
+                    };
+                    arg
+                });
 
-                // delete everything after the format string till end bracket
-                // we're going to insert the new arguments later
-                edit.delete(TextRange::new(
-                    format_string.text_range().end(),
-                    end_bracket.text_range().start(),
-                ));
-            }
+                args.collect()
+            } else {
+                vec![]
+            };
 
             // Start building the new args
             let mut existing_args = existing_args.into_iter();
-            let mut args = String::new();
+            let mut new_tt_bits = vec![NodeOrToken::Token(make::tokens::literal(&new_fmt))];
+            let mut placeholder_indexes = vec![];
 
-            let mut placeholder_idx = 1;
+            for arg in extracted_args {
+                if matches!(arg, Arg::Expr(_) | Arg::Placeholder) {
+                    // insert ", " before each arg
+                    new_tt_bits.extend_from_slice(&[
+                        NodeOrToken::Token(make::token(T![,])),
+                        NodeOrToken::Token(make::tokens::single_space()),
+                    ]);
+                }
 
-            for extracted_args in extracted_args {
-                match extracted_args {
-                    Arg::Expr(s)=> {
-                        args.push_str(", ");
+                match arg {
+                    Arg::Expr(s) => {
                         // insert arg
-                        args.push_str(&s);
+                        // FIXME: use the crate's edition for parsing
+                        let expr = ast::Expr::parse(&s, syntax::Edition::CURRENT).syntax_node();
+                        let mut expr_tt = utils::tt_from_syntax(expr);
+                        new_tt_bits.append(&mut expr_tt);
                     }
                     Arg::Placeholder => {
-                        args.push_str(", ");
                         // try matching with existing argument
                         match existing_args.next() {
-                            Some(ea) => {
-                                args.push_str(&ea);
+                            Some(arg) => {
+                                new_tt_bits.extend_from_slice(arg);
                             }
                             None => {
-                                // insert placeholder
-                                args.push_str(&format!("${placeholder_idx}"));
-                                placeholder_idx += 1;
+                                placeholder_indexes.push(new_tt_bits.len());
+                                new_tt_bits.push(NodeOrToken::Token(make::token(T![_])));
                             }
                         }
                     }
                     Arg::Ident(_s) => (),
                 }
             }
 
+
             // Insert new args
-            edit.insert(fmt_range.end(), args);
+            let new_tt = make::token_tree(tt_delimiter, new_tt_bits).clone_for_update();
+            ted::replace(tt.syntax(), new_tt.syntax());
+
+            if let Some(cap) = ctx.config.snippet_cap {
+                // Add placeholder snippets over placeholder args
+                for pos in placeholder_indexes {
+                    // Skip the opening delimiter
+                    let Some(NodeOrToken::Token(placeholder)) =
+                        new_tt.token_trees_and_tokens().skip(1).nth(pos)
+                    else {
+                        continue;
+                    };
+
+                    if stdx::always!(placeholder.kind() == T![_]) {
+                        edit.add_placeholder_snippet_token(cap, placeholder);
+                    }
+                }
+
+                // Add the final tabstop after the format literal
+                if let Some(NodeOrToken::Token(literal)) = new_tt.token_trees_and_tokens().nth(1) {
+                    edit.add_tabstop_after_token(cap, literal);
+                }
+            }
         },
     );
 
@@ -145,7 +164,7 @@ pub(crate) fn extract_expressions_from_format_string(
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::tests::check_assist;
+    use crate::tests::{check_assist, check_assist_no_snippet_cap};
 
     #[test]
     fn multiple_middle_arg() {
@@ -195,7 +214,7 @@ fn main() {
 "#,
             r#"
 fn main() {
-    print!("{} {:b} {} {}"$0, y + 2, x + 1, 2, $1);
+    print!("{} {:b} {} {}"$0, y + 2, x + 1, 2, ${1:_});
 }
 "#,
         );
@@ -292,4 +311,22 @@ fn main() {
             "#,
         );
     }
+
+    #[test]
+    fn without_snippets() {
+        check_assist_no_snippet_cap(
+            extract_expressions_from_format_string,
+            r#"
+//- minicore: fmt
+fn main() {
+    print!("{} {x + 1:b} {} {}$0", y + 2, 2);
+}
+"#,
+            r#"
+fn main() {
+    print!("{} {:b} {} {}", y + 2, x + 1, 2, _);
+}
+"#,
+        );
+    }
 }

src/tools/rust-analyzer/crates/ide-db/src/source_change.rs

@@ -338,6 +338,12 @@ impl SourceChangeBuilder {
         self.add_snippet(PlaceSnippet::Over(node.syntax().clone().into()))
     }
 
+    /// Adds a snippet to move the cursor selected over `token`
+    pub fn add_placeholder_snippet_token(&mut self, _cap: SnippetCap, token: SyntaxToken) {
+        assert!(token.parent().is_some());
+        self.add_snippet(PlaceSnippet::Over(token.into()))
+    }
+
     /// Adds a snippet to move the cursor selected over `nodes`
     ///
     /// This allows for renaming newly generated items without having to go

src/tools/rust-analyzer/crates/syntax/src/ast/make.rs

@@ -1159,7 +1159,7 @@ pub mod tokens {
 
     pub(super) static SOURCE_FILE: Lazy<Parse<SourceFile>> = Lazy::new(|| {
         SourceFile::parse(
-            "const C: <()>::Item = ( true && true , true || true , 1 != 1, 2 == 2, 3 < 3, 4 <= 4, 5 > 5, 6 >= 6, !true, *p, &p , &mut p, { let a @ [] })\n;\n\nimpl A for B where: {}", Edition::CURRENT,
+            "const C: <()>::Item = ( true && true , true || true , 1 != 1, 2 == 2, 3 < 3, 4 <= 4, 5 > 5, 6 >= 6, !true, *p, &p , &mut p, { let _ @ [] })\n;\n\nimpl A for B where: {}", Edition::CURRENT,
         )
     });