Merge pull request #20270 from Hmikihiro/migrate_generate_new

ShoyuVanilla · web-flow · commit 7cbc5a3d57a3 · 2025-07-22T03:41:03.000Z
Migrate `generate new` assist to use `SyntaxEditor`
diff --git a/crates/ide-assists/src/handlers/generate_mut_trait_impl.rs b/crates/ide-assists/src/handlers/generate_mut_trait_impl.rs
@@ -94,7 +94,7 @@ pub(crate) fn generate_mut_trait_impl(acc: &mut Assists, ctx: &AssistContext<'_>
     })?;
     let _ = process_ref_mut(&fn_);
 
-    let assoc_list = make::assoc_item_list().clone_for_update();
+    let assoc_list = make::assoc_item_list(None).clone_for_update();
     ted::replace(impl_def.assoc_item_list()?.syntax(), assoc_list.syntax());
     impl_def.get_or_create_assoc_item_list().add_item(syntax::ast::AssocItem::Fn(fn_));
 
diff --git a/crates/ide-assists/src/handlers/generate_new.rs b/crates/ide-assists/src/handlers/generate_new.rs
@@ -4,12 +4,12 @@ use ide_db::{
 };
 use syntax::{
     ast::{self, AstNode, HasName, HasVisibility, StructKind, edit_in_place::Indent, make},
-    ted,
+    syntax_editor::Position,
 };
 
 use crate::{
     AssistContext, AssistId, Assists,
-    utils::{find_struct_impl, generate_impl},
+    utils::{find_struct_impl, generate_impl_with_item},
 };
 
 // Assist: generate_new
@@ -149,7 +149,53 @@ pub(crate) fn generate_new(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option
         .clone_for_update();
         fn_.indent(1.into());
 
-        if let Some(cap) = ctx.config.snippet_cap {
+        let mut editor = builder.make_editor(strukt.syntax());
+
+        // Get the node for set annotation
+        let contain_fn = if let Some(impl_def) = impl_def {
+            fn_.indent(impl_def.indent_level());
+
+            if let Some(l_curly) = impl_def.assoc_item_list().and_then(|list| list.l_curly_token())
+            {
+                editor.insert_all(
+                    Position::after(l_curly),
+                    vec![
+                        make::tokens::whitespace(&format!("\n{}", impl_def.indent_level() + 1))
+                            .into(),
+                        fn_.syntax().clone().into(),
+                        make::tokens::whitespace("\n").into(),
+                    ],
+                );
+                fn_.syntax().clone()
+            } else {
+                let items = vec![either::Either::Right(ast::AssocItem::Fn(fn_))];
+                let list = make::assoc_item_list(Some(items));
+                editor.insert(Position::after(impl_def.syntax()), list.syntax());
+                list.syntax().clone()
+            }
+        } else {
+            // Generate a new impl to add the method to
+            let indent_level = strukt.indent_level();
+            let body = vec![either::Either::Right(ast::AssocItem::Fn(fn_))];
+            let list = make::assoc_item_list(Some(body));
+            let impl_def = generate_impl_with_item(&ast::Adt::Struct(strukt.clone()), Some(list));
+
+            impl_def.indent(strukt.indent_level());
+
+            // Insert it after the adt
+            editor.insert_all(
+                Position::after(strukt.syntax()),
+                vec![
+                    make::tokens::whitespace(&format!("\n\n{indent_level}")).into(),
+                    impl_def.syntax().clone().into(),
+                ],
+            );
+            impl_def.syntax().clone()
+        };
+
+        if let Some(fn_) = contain_fn.descendants().find_map(ast::Fn::cast)
+            && let Some(cap) = ctx.config.snippet_cap
+        {
             match strukt.kind() {
                 StructKind::Tuple(_) => {
                     let struct_args = fn_
@@ -168,8 +214,8 @@ pub(crate) fn generate_new(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option
                         for (struct_arg, fn_param) in struct_args.zip(fn_params.params()) {
                             if let Some(fn_pat) = fn_param.pat() {
                                 let fn_pat = fn_pat.syntax().clone();
-                                builder
-                                    .add_placeholder_snippet_group(cap, vec![struct_arg, fn_pat]);
+                                let placeholder = builder.make_placeholder_snippet(cap);
+                                editor.add_annotation_all(vec![struct_arg, fn_pat], placeholder)
                             }
                         }
                     }
@@ -179,36 +225,12 @@ pub(crate) fn generate_new(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option
 
             // Add a tabstop before the name
             if let Some(name) = fn_.name() {
-                builder.add_tabstop_before(cap, name);
+                let tabstop_before = builder.make_tabstop_before(cap);
+                editor.add_annotation(name.syntax(), tabstop_before);
             }
         }
 
-        // Get the mutable version of the impl to modify
-        let impl_def = if let Some(impl_def) = impl_def {
-            fn_.indent(impl_def.indent_level());
-            builder.make_mut(impl_def)
-        } else {
-            // Generate a new impl to add the method to
-            let impl_def = generate_impl(&ast::Adt::Struct(strukt.clone()));
-            let indent_level = strukt.indent_level();
-            fn_.indent(indent_level);
-
-            // Insert it after the adt
-            let strukt = builder.make_mut(strukt.clone());
-
-            ted::insert_all_raw(
-                ted::Position::after(strukt.syntax()),
-                vec![
-                    make::tokens::whitespace(&format!("\n\n{indent_level}")).into(),
-                    impl_def.syntax().clone().into(),
-                ],
-            );
-
-            impl_def
-        };
-
-        // Add the `new` method at the start of the impl
-        impl_def.get_or_create_assoc_item_list().add_item_at_start(fn_.into());
+        builder.add_file_edits(ctx.vfs_file_id(), editor);
     })
 }
 
diff --git a/crates/ide-assists/src/utils.rs b/crates/ide-assists/src/utils.rs
@@ -664,30 +664,38 @@ fn generate_impl_text_inner(
 
 /// Generates the corresponding `impl Type {}` including type and lifetime
 /// parameters.
+pub(crate) fn generate_impl_with_item(
+    adt: &ast::Adt,
+    body: Option<ast::AssocItemList>,
+) -> ast::Impl {
+    generate_impl_inner(adt, None, true, body)
+}
+
 pub(crate) fn generate_impl(adt: &ast::Adt) -> ast::Impl {
-    generate_impl_inner(adt, None, true)
+    generate_impl_inner(adt, None, true, None)
 }
 
 /// Generates the corresponding `impl <trait> for Type {}` including type
 /// and lifetime parameters, with `<trait>` appended to `impl`'s generic parameters' bounds.
 ///
 /// This is useful for traits like `PartialEq`, since `impl<T> PartialEq for U<T>` often requires `T: PartialEq`.
 pub(crate) fn generate_trait_impl(adt: &ast::Adt, trait_: ast::Type) -> ast::Impl {
-    generate_impl_inner(adt, Some(trait_), true)
+    generate_impl_inner(adt, Some(trait_), true, None)
 }
 
 /// Generates the corresponding `impl <trait> for Type {}` including type
 /// and lifetime parameters, with `impl`'s generic parameters' bounds kept as-is.
 ///
 /// This is useful for traits like `From<T>`, since `impl<T> From<T> for U<T>` doesn't require `T: From<T>`.
 pub(crate) fn generate_trait_impl_intransitive(adt: &ast::Adt, trait_: ast::Type) -> ast::Impl {
-    generate_impl_inner(adt, Some(trait_), false)
+    generate_impl_inner(adt, Some(trait_), false, None)
 }
 
 fn generate_impl_inner(
     adt: &ast::Adt,
     trait_: Option<ast::Type>,
     trait_is_transitive: bool,
+    body: Option<ast::AssocItemList>,
 ) -> ast::Impl {
     // Ensure lifetime params are before type & const params
     let generic_params = adt.generic_param_list().map(|generic_params| {
@@ -737,9 +745,9 @@ fn generate_impl_inner(
             ty,
             None,
             adt.where_clause(),
-            None,
+            body,
         ),
-        None => make::impl_(generic_params, generic_args, ty, adt.where_clause(), None),
+        None => make::impl_(generic_params, generic_args, ty, adt.where_clause(), body),
     }
     .clone_for_update();
 
diff --git a/crates/syntax/src/ast/edit_in_place.rs b/crates/syntax/src/ast/edit_in_place.rs
@@ -644,7 +644,7 @@ impl Removable for ast::Use {
 impl ast::Impl {
     pub fn get_or_create_assoc_item_list(&self) -> ast::AssocItemList {
         if self.assoc_item_list().is_none() {
-            let assoc_item_list = make::assoc_item_list().clone_for_update();
+            let assoc_item_list = make::assoc_item_list(None).clone_for_update();
             ted::append_child(self.syntax(), assoc_item_list.syntax());
         }
         self.assoc_item_list().unwrap()
diff --git a/crates/syntax/src/ast/make.rs b/crates/syntax/src/ast/make.rs
@@ -229,8 +229,18 @@ pub fn ty_fn_ptr<I: Iterator<Item = Param>>(
     }
 }
 
-pub fn assoc_item_list() -> ast::AssocItemList {
-    ast_from_text("impl C for D {}")
+pub fn assoc_item_list(
+    body: Option<Vec<either::Either<ast::Attr, ast::AssocItem>>>,
+) -> ast::AssocItemList {
+    let is_break_braces = body.is_some();
+    let body_newline = if is_break_braces { "\n".to_owned() } else { String::new() };
+    let body_indent = if is_break_braces { "    ".to_owned() } else { String::new() };
+
+    let body = match body {
+        Some(bd) => bd.iter().map(|elem| elem.to_string()).join(""),
+        None => String::new(),
+    };
+    ast_from_text(&format!("impl C for D {{{body_newline}{body_indent}{body}{body_newline}}}"))
 }
 
 fn merge_gen_params(
@@ -273,28 +283,21 @@ pub fn impl_(
     generic_args: Option<ast::GenericArgList>,
     path_type: ast::Type,
     where_clause: Option<ast::WhereClause>,
-    body: Option<Vec<either::Either<ast::Attr, ast::AssocItem>>>,
+    body: Option<ast::AssocItemList>,
 ) -> ast::Impl {
     let gen_args = generic_args.map_or_else(String::new, |it| it.to_string());
 
     let gen_params = generic_params.map_or_else(String::new, |it| it.to_string());
 
     let body_newline =
         if where_clause.is_some() && body.is_none() { "\n".to_owned() } else { String::new() };
-
     let where_clause = match where_clause {
         Some(pr) => format!("\n{pr}\n"),
         None => " ".to_owned(),
     };
 
-    let body = match body {
-        Some(bd) => bd.iter().map(|elem| elem.to_string()).join(""),
-        None => String::new(),
-    };
-
-    ast_from_text(&format!(
-        "impl{gen_params} {path_type}{gen_args}{where_clause}{{{body_newline}{body}}}"
-    ))
+    let body = body.map_or_else(|| format!("{{{body_newline}}}"), |it| it.to_string());
+    ast_from_text(&format!("impl{gen_params} {path_type}{gen_args}{where_clause}{body}"))
 }
 
 pub fn impl_trait(
@@ -308,7 +311,7 @@ pub fn impl_trait(
     ty: ast::Type,
     trait_where_clause: Option<ast::WhereClause>,
     ty_where_clause: Option<ast::WhereClause>,
-    body: Option<Vec<either::Either<ast::Attr, ast::AssocItem>>>,
+    body: Option<ast::AssocItemList>,
 ) -> ast::Impl {
     let is_unsafe = if is_unsafe { "unsafe " } else { "" };
 
@@ -330,13 +333,10 @@ pub fn impl_trait(
     let where_clause = merge_where_clause(ty_where_clause, trait_where_clause)
         .map_or_else(|| " ".to_owned(), |wc| format!("\n{wc}\n"));
 
-    let body = match body {
-        Some(bd) => bd.iter().map(|elem| elem.to_string()).join(""),
-        None => String::new(),
-    };
+    let body = body.map_or_else(|| format!("{{{body_newline}}}"), |it| it.to_string());
 
     ast_from_text(&format!(
-        "{is_unsafe}impl{gen_params} {is_negative}{path_type}{trait_gen_args} for {ty}{type_gen_args}{where_clause}{{{body_newline}{body}}}"
+        "{is_unsafe}impl{gen_params} {is_negative}{path_type}{trait_gen_args} for {ty}{type_gen_args}{where_clause}{body}"
     ))
 }
 
diff --git a/crates/syntax/src/syntax_editor.rs b/crates/syntax/src/syntax_editor.rs
@@ -5,7 +5,7 @@
 //! [`SyntaxEditor`]: https://github.com/dotnet/roslyn/blob/43b0b05cc4f492fd5de00f6f6717409091df8daa/src/Workspaces/Core/Portable/Editing/SyntaxEditor.cs
 
 use std::{
-    fmt,
+    fmt, iter,
     num::NonZeroU32,
     ops::RangeInclusive,
     sync::atomic::{AtomicU32, Ordering},
@@ -41,6 +41,15 @@ impl SyntaxEditor {
         self.annotations.push((element.syntax_element(), annotation))
     }
 
+    pub fn add_annotation_all(
+        &mut self,
+        elements: Vec<impl Element>,
+        annotation: SyntaxAnnotation,
+    ) {
+        self.annotations
+            .extend(elements.into_iter().map(|e| e.syntax_element()).zip(iter::repeat(annotation)));
+    }
+
     pub fn merge(&mut self, mut other: SyntaxEditor) {
         debug_assert!(
             self.root == other.root || other.root.ancestors().any(|node| node == self.root),

Original file line number	Diff line number	Diff line change
`@@ -644,7 +644,7 @@ impl Removable for ast::Use {`
`644`	`644`	`impl ast::Impl {`
`645`	`645`	`pub fn get_or_create_assoc_item_list(&self) -> ast::AssocItemList {`
`646`	`646`	`if self.assoc_item_list().is_none() {`
`647`		`- let assoc_item_list = make::assoc_item_list().clone_for_update();`
	`647`	`+ let assoc_item_list = make::assoc_item_list(None).clone_for_update();`
`648`	`648`	`ted::append_child(self.syntax(), assoc_item_list.syntax());`
`649`	`649`	`}`
`650`	`650`	`self.assoc_item_list().unwrap()`