split ted from gen_trait_fn_body

Hmikihiro · Hmikihiro · commit 4a0527f78c4f · 2025-07-25T19:50:44.000+09:00
diff --git a/crates/ide-assists/src/handlers/add_missing_impl_members.rs b/crates/ide-assists/src/handlers/add_missing_impl_members.rs
@@ -2,6 +2,7 @@ use hir::HasSource;
 use syntax::{
     Edition,
     ast::{self, AstNode, make},
+    ted,
 };
 
 use crate::{
@@ -157,19 +158,21 @@ fn add_missing_impl_members_inner(
             &target_scope,
         );
 
+        let mut editor = edit.make_editor(impl_def.syntax());
         if let Some(cap) = ctx.config.snippet_cap {
             let mut placeholder = None;
             if let DefaultMethods::No = mode {
                 if let ast::AssocItem::Fn(func) = &first_new_item {
-                    if try_gen_trait_body(
+                    if let Some(body) = try_gen_trait_body(
                         ctx,
                         func,
                         trait_ref,
                         &impl_def,
                         target_scope.krate().edition(ctx.sema.db),
-                    )
-                    .is_none()
+                    ) && let Some(func_body) = func.body()
                     {
+                        ted::replace(func_body.syntax(), body.syntax());
+                    } else {
                         if let Some(m) = func.syntax().descendants().find_map(ast::MacroCall::cast)
                         {
                             if m.syntax().text() == "todo!()" {
@@ -195,7 +198,7 @@ fn try_gen_trait_body(
     trait_ref: hir::TraitRef<'_>,
     impl_def: &ast::Impl,
     edition: Edition,
-) -> Option<()> {
+) -> Option<ast::BlockExpr> {
     let trait_path = make::ext::ident_path(
         &trait_ref.trait_().name(ctx.db()).display(ctx.db(), edition).to_string(),
     );
diff --git a/crates/ide-assists/src/handlers/replace_derive_with_manual_impl.rs b/crates/ide-assists/src/handlers/replace_derive_with_manual_impl.rs
@@ -135,47 +135,29 @@ fn add_assist(
             &annotated_name,
             trait_,
             replace_trait_path,
+            impl_is_unsafe,
         );
         update_attribute(builder, old_derives, old_tree, old_trait_path, attr);
 
         let trait_path = make::ty_path(replace_trait_path.clone());
 
         match (ctx.config.snippet_cap, impl_def_with_items) {
             (None, None) => {
-                let impl_def = generate_trait_impl(adt, trait_path);
-                if impl_is_unsafe {
-                    ted::insert(
-                        Position::first_child_of(impl_def.syntax()),
-                        make::token(T![unsafe]),
-                    );
-                }
+                let impl_def = generate_trait_impl(impl_is_unsafe, adt, trait_path);
 
                 ted::insert_all(
                     insert_after,
                     vec![make::tokens::blank_line().into(), impl_def.syntax().clone().into()],
                 );
             }
             (None, Some((impl_def, _))) => {
-                if impl_is_unsafe {
-                    ted::insert(
-                        Position::first_child_of(impl_def.syntax()),
-                        make::token(T![unsafe]),
-                    );
-                }
                 ted::insert_all(
                     insert_after,
                     vec![make::tokens::blank_line().into(), impl_def.syntax().clone().into()],
                 );
             }
             (Some(cap), None) => {
-                let impl_def = generate_trait_impl(adt, trait_path);
-
-                if impl_is_unsafe {
-                    ted::insert(
-                        Position::first_child_of(impl_def.syntax()),
-                        make::token(T![unsafe]),
-                    );
-                }
+                let impl_def = generate_trait_impl(impl_is_unsafe, adt, trait_path);
 
                 if let Some(l_curly) = impl_def.assoc_item_list().and_then(|it| it.l_curly_token())
                 {
@@ -188,26 +170,13 @@ fn add_assist(
                 );
             }
             (Some(cap), Some((impl_def, first_assoc_item))) => {
-                let mut added_snippet = false;
-
-                if impl_is_unsafe {
-                    ted::insert(
-                        Position::first_child_of(impl_def.syntax()),
-                        make::token(T![unsafe]),
-                    );
-                }
-
-                if let ast::AssocItem::Fn(ref func) = first_assoc_item {
-                    if let Some(m) = func.syntax().descendants().find_map(ast::MacroCall::cast) {
-                        if m.syntax().text() == "todo!()" {
-                            // Make the `todo!()` a placeholder
-                            builder.add_placeholder_snippet(cap, m);
-                            added_snippet = true;
-                        }
-                    }
-                }
-
-                if !added_snippet {
+                if let ast::AssocItem::Fn(ref func) = first_assoc_item
+                    && let Some(m) = func.syntax().descendants().find_map(ast::MacroCall::cast)
+                    && m.syntax().text() == "todo!()"
+                {
+                    // Make the `todo!()` a placeholder
+                    builder.add_placeholder_snippet(cap, m);
+                } else {
                     // If we haven't already added a snippet, add a tabstop before the generated function
                     builder.add_tabstop_before(cap, first_assoc_item);
                 }
@@ -228,6 +197,7 @@ fn impl_def_from_trait(
     annotated_name: &ast::Name,
     trait_: Option<hir::Trait>,
     trait_path: &ast::Path,
+    impl_is_unsafe: bool,
 ) -> Option<(ast::Impl, ast::AssocItem)> {
     let trait_ = trait_?;
     let target_scope = sema.scope(annotated_name.syntax())?;
@@ -245,14 +215,18 @@ fn impl_def_from_trait(
     if trait_items.is_empty() {
         return None;
     }
-    let impl_def = generate_trait_impl(adt, make::ty_path(trait_path.clone()));
+    let impl_def = generate_trait_impl(impl_is_unsafe, adt, make::ty_path(trait_path.clone()));
 
     let first_assoc_item =
         add_trait_assoc_items_to_impl(sema, config, &trait_items, trait_, &impl_def, &target_scope);
 
     // Generate a default `impl` function body for the derived trait.
     if let ast::AssocItem::Fn(ref func) = first_assoc_item {
-        let _ = gen_trait_fn_body(func, trait_path, adt, None);
+        if let Some(body) = gen_trait_fn_body(func, trait_path, adt, None)
+            && let Some(func_body) = func.body()
+        {
+            ted::replace(func_body.syntax(), body.syntax());
+        }
     };
 
     Some((impl_def, first_assoc_item))
diff --git a/crates/ide-assists/src/utils.rs b/crates/ide-assists/src/utils.rs
@@ -668,30 +668,31 @@ pub(crate) fn generate_impl_with_item(
     adt: &ast::Adt,
     body: Option<ast::AssocItemList>,
 ) -> ast::Impl {
-    generate_impl_inner(adt, None, true, body)
+    generate_impl_inner(false, adt, None, true, body)
 }
 
 pub(crate) fn generate_impl(adt: &ast::Adt) -> ast::Impl {
-    generate_impl_inner(adt, None, true, None)
+    generate_impl_inner(false, 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, None)
+pub(crate) fn generate_trait_impl(is_unsafe: bool, adt: &ast::Adt, trait_: ast::Type) -> ast::Impl {
+    generate_impl_inner(is_unsafe, 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, None)
+    generate_impl_inner(false, adt, Some(trait_), false, None)
 }
 
 fn generate_impl_inner(
+    is_unsafe: bool,
     adt: &ast::Adt,
     trait_: Option<ast::Type>,
     trait_is_transitive: bool,
@@ -735,7 +736,7 @@ fn generate_impl_inner(
 
     let impl_ = match trait_ {
         Some(trait_) => make::impl_trait(
-            false,
+            is_unsafe,
             None,
             None,
             generic_params,
diff --git a/crates/ide-assists/src/utils/gen_trait_fn_body.rs b/crates/ide-assists/src/utils/gen_trait_fn_body.rs
@@ -1,10 +1,7 @@
 //! This module contains functions to generate default trait impl function bodies where possible.
 
 use hir::TraitRef;
-use syntax::{
-    ast::{self, AstNode, BinaryOp, CmpOp, HasName, LogicOp, edit::AstNodeEdit, make},
-    ted,
-};
+use syntax::ast::{self, AstNode, BinaryOp, CmpOp, HasName, LogicOp, edit::AstNodeEdit, make};
 
 /// Generate custom trait bodies without default implementation where possible.
 ///
@@ -18,21 +15,33 @@ pub(crate) fn gen_trait_fn_body(
     trait_path: &ast::Path,
     adt: &ast::Adt,
     trait_ref: Option<TraitRef<'_>>,
-) -> Option<()> {
+) -> Option<ast::BlockExpr> {
+    let _ = func.body()?;
     match trait_path.segment()?.name_ref()?.text().as_str() {
-        "Clone" => gen_clone_impl(adt, func),
-        "Debug" => gen_debug_impl(adt, func),
-        "Default" => gen_default_impl(adt, func),
-        "Hash" => gen_hash_impl(adt, func),
-        "PartialEq" => gen_partial_eq(adt, func, trait_ref),
-        "PartialOrd" => gen_partial_ord(adt, func, trait_ref),
+        "Clone" => {
+            stdx::always!(func.name().is_some_and(|name| name.text() == "clone"));
+            gen_clone_impl(adt)
+        }
+        "Debug" => gen_debug_impl(adt),
+        "Default" => gen_default_impl(adt),
+        "Hash" => {
+            stdx::always!(func.name().is_some_and(|name| name.text() == "hash"));
+            gen_hash_impl(adt)
+        }
+        "PartialEq" => {
+            stdx::always!(func.name().is_some_and(|name| name.text() == "eq"));
+            gen_partial_eq(adt, trait_ref)
+        }
+        "PartialOrd" => {
+            stdx::always!(func.name().is_some_and(|name| name.text() == "partial_cmp"));
+            gen_partial_ord(adt, trait_ref)
+        }
         _ => None,
     }
 }
 
 /// Generate a `Clone` impl based on the fields and members of the target type.
-fn gen_clone_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
-    stdx::always!(func.name().is_some_and(|name| name.text() == "clone"));
+fn gen_clone_impl(adt: &ast::Adt) -> Option<ast::BlockExpr> {
     fn gen_clone_call(target: ast::Expr) -> ast::Expr {
         let method = make::name_ref("clone");
         make::expr_method_call(target, method, make::arg_list(None)).into()
@@ -139,12 +148,11 @@ fn gen_clone_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
         }
     };
     let body = make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1));
-    ted::replace(func.body()?.syntax(), body.syntax());
-    Some(())
+    Some(body)
 }
 
 /// Generate a `Debug` impl based on the fields and members of the target type.
-fn gen_debug_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
+fn gen_debug_impl(adt: &ast::Adt) -> Option<ast::BlockExpr> {
     let annotated_name = adt.name()?;
     match adt {
         // `Debug` cannot be derived for unions, so no default impl can be provided.
@@ -248,8 +256,7 @@ fn gen_debug_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
 
             let body = make::block_expr(None, Some(match_expr.into()));
             let body = body.indent(ast::edit::IndentLevel(1));
-            ted::replace(func.body()?.syntax(), body.syntax());
-            Some(())
+            Some(body)
         }
 
         ast::Adt::Struct(strukt) => {
@@ -296,14 +303,13 @@ fn gen_debug_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
             let method = make::name_ref("finish");
             let expr = make::expr_method_call(expr, method, make::arg_list(None)).into();
             let body = make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1));
-            ted::replace(func.body()?.syntax(), body.syntax());
-            Some(())
+            Some(body)
         }
     }
 }
 
 /// Generate a `Debug` impl based on the fields and members of the target type.
-fn gen_default_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
+fn gen_default_impl(adt: &ast::Adt) -> Option<ast::BlockExpr> {
     fn gen_default_call() -> Option<ast::Expr> {
         let fn_name = make::ext::path_from_idents(["Default", "default"])?;
         Some(make::expr_call(make::expr_path(fn_name), make::arg_list(None)).into())
@@ -342,15 +348,13 @@ fn gen_default_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
                 }
             };
             let body = make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1));
-            ted::replace(func.body()?.syntax(), body.syntax());
-            Some(())
+            Some(body)
         }
     }
 }
 
 /// Generate a `Hash` impl based on the fields and members of the target type.
-fn gen_hash_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
-    stdx::always!(func.name().is_some_and(|name| name.text() == "hash"));
+fn gen_hash_impl(adt: &ast::Adt) -> Option<ast::BlockExpr> {
     fn gen_hash_call(target: ast::Expr) -> ast::Stmt {
         let method = make::name_ref("hash");
         let arg = make::expr_path(make::ext::ident_path("state"));
@@ -400,13 +404,11 @@ fn gen_hash_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
         },
     };
 
-    ted::replace(func.body()?.syntax(), body.syntax());
-    Some(())
+    Some(body)
 }
 
 /// Generate a `PartialEq` impl based on the fields and members of the target type.
-fn gen_partial_eq(adt: &ast::Adt, func: &ast::Fn, trait_ref: Option<TraitRef<'_>>) -> Option<()> {
-    stdx::always!(func.name().is_some_and(|name| name.text() == "eq"));
+fn gen_partial_eq(adt: &ast::Adt, trait_ref: Option<TraitRef<'_>>) -> Option<ast::BlockExpr> {
     fn gen_eq_chain(expr: Option<ast::Expr>, cmp: ast::Expr) -> Option<ast::Expr> {
         match expr {
             Some(expr) => Some(make::expr_bin_op(expr, BinaryOp::LogicOp(LogicOp::And), cmp)),
@@ -595,12 +597,10 @@ fn gen_partial_eq(adt: &ast::Adt, func: &ast::Fn, trait_ref: Option<TraitRef<'_>
         },
     };
 
-    ted::replace(func.body()?.syntax(), body.syntax());
-    Some(())
+    Some(body)
 }
 
-fn gen_partial_ord(adt: &ast::Adt, func: &ast::Fn, trait_ref: Option<TraitRef<'_>>) -> Option<()> {
-    stdx::always!(func.name().is_some_and(|name| name.text() == "partial_cmp"));
+fn gen_partial_ord(adt: &ast::Adt, trait_ref: Option<TraitRef<'_>>) -> Option<ast::BlockExpr> {
     fn gen_partial_eq_match(match_target: ast::Expr) -> Option<ast::Stmt> {
         let mut arms = vec![];
 
@@ -686,8 +686,7 @@ fn gen_partial_ord(adt: &ast::Adt, func: &ast::Fn, trait_ref: Option<TraitRef<'_
         },
     };
 
-    ted::replace(func.body()?.syntax(), body.syntax());
-    Some(())
+    Some(body)
 }
 
 fn make_discriminant() -> Option<ast::Expr> {
