Add trait codegen to add_missing_impl_members assist

Author: yoshuawuyts

crates/hir/src/has_source.rs

@@ -10,8 +10,8 @@ use hir_expand::InFile;
 use syntax::ast;
 
 use crate::{
-    db::HirDatabase, Const, ConstParam, Enum, Field, FieldSource, Function, Impl, LifetimeParam,
-    MacroDef, Module, Static, Struct, Trait, TypeAlias, TypeParam, Union, Variant,
+    db::HirDatabase, Adt, Const, ConstParam, Enum, Field, FieldSource, Function, Impl,
+    LifetimeParam, MacroDef, Module, Static, Struct, Trait, TypeAlias, TypeParam, Union, Variant,
 };
 
 pub trait HasSource {
@@ -56,6 +56,16 @@ impl HasSource for Field {
         Some(field_source)
     }
 }
+impl HasSource for Adt {
+    type Ast = ast::Adt;
+    fn source(self, db: &dyn HirDatabase) -> Option<InFile<Self::Ast>> {
+        match self {
+            Adt::Struct(s) => Some(s.source(db)?.map(|s| ast::Adt::Struct(s))),
+            Adt::Union(u) => Some(u.source(db)?.map(|u| ast::Adt::Union(u))),
+            Adt::Enum(e) => Some(e.source(db)?.map(|e| ast::Adt::Enum(e))),
+        }
+    }
+}
 impl HasSource for Struct {
     type Ast = ast::Struct;
     fn source(self, db: &dyn HirDatabase) -> Option<InFile<Self::Ast>> {

crates/ide_assists/src/handlers/add_missing_impl_members.rs

@@ -1,10 +1,12 @@
+use hir::HasSource;
 use ide_db::traits::resolve_target_trait;
-use syntax::ast::{self, AstNode};
+use syntax::ast::{self, make, AstNode};
 
 use crate::{
     assist_context::{AssistContext, Assists},
     utils::{
-        add_trait_assoc_items_to_impl, filter_assoc_items, render_snippet, Cursor, DefaultMethods,
+        add_trait_assoc_items_to_impl, filter_assoc_items, gen_trait_body, render_snippet, Cursor,
+        DefaultMethods,
     },
     AssistId, AssistKind,
 };
@@ -115,18 +117,26 @@ fn add_missing_impl_members_inner(
     let target = impl_def.syntax().text_range();
     acc.add(AssistId(assist_id, AssistKind::QuickFix), label, target, |builder| {
         let target_scope = ctx.sema.scope(impl_def.syntax());
-        let (new_impl_def, first_new_item) =
-            add_trait_assoc_items_to_impl(&ctx.sema, missing_items, trait_, impl_def, target_scope);
+        let (new_impl_def, first_new_item) = add_trait_assoc_items_to_impl(
+            &ctx.sema,
+            missing_items,
+            trait_,
+            impl_def.clone(),
+            target_scope,
+        );
         match ctx.config.snippet_cap {
             None => builder.replace(target, new_impl_def.to_string()),
             Some(cap) => {
                 let mut cursor = Cursor::Before(first_new_item.syntax());
                 let placeholder;
                 if let ast::AssocItem::Fn(func) = &first_new_item {
-                    if let Some(m) = func.syntax().descendants().find_map(ast::MacroCall::cast) {
-                        if m.syntax().text() == "todo!()" {
-                            placeholder = m;
-                            cursor = Cursor::Replace(placeholder.syntax());
+                    if try_gen_trait_body(ctx, func, &trait_, &impl_def).is_none() {
+                        if let Some(m) = func.syntax().descendants().find_map(ast::MacroCall::cast)
+                        {
+                            if m.syntax().text() == "todo!()" {
+                                placeholder = m;
+                                cursor = Cursor::Replace(placeholder.syntax());
+                            }
                         }
                     }
                 }
@@ -140,6 +150,18 @@ fn add_missing_impl_members_inner(
     })
 }
 
+fn try_gen_trait_body(
+    ctx: &AssistContext,
+    func: &ast::Fn,
+    trait_: &hir::Trait,
+    impl_def: &ast::Impl,
+) -> Option<()> {
+    let trait_path = make::path_from_text(&trait_.name(ctx.db()).to_string());
+    let hir_ty = ctx.sema.resolve_type(&impl_def.self_ty()?)?;
+    let adt = hir_ty.as_adt()?.source(ctx.db())?;
+    gen_trait_body(func, &trait_path, &adt.value)
+}
+
 #[cfg(test)]
 mod tests {
     use crate::tests::{check_assist, check_assist_not_applicable};
@@ -847,4 +869,28 @@ impl T for () {
 ",
         );
     }
+
+    #[test]
+    fn test_default_body_generation() {
+        check_assist(
+            add_missing_impl_members,
+            r#"
+//- minicore: default
+struct Foo(usize);
+
+impl Default for Foo {
+    $0
+}
+"#,
+            r#"
+struct Foo(usize);
+
+impl Default for Foo {
+    $0fn default() -> Self {
+        Self(Default::default())
+    }
+}
+"#,
+        )
+    }
 }

crates/ide_assists/src/handlers/replace_derive_with_manual_impl.rs

@@ -168,7 +168,7 @@ fn impl_def_from_trait(
 
     // Generate a default `impl` function body for the derived trait.
     if let ast::AssocItem::Fn(ref func) = first_assoc_item {
-        let _ = gen_trait_body(func, trait_path, adt, annotated_name);
+        let _ = gen_trait_body(func, trait_path, adt);
     };
 
     Some((impl_def, first_assoc_item))

crates/ide_assists/src/utils/gen_trait_body.rs

@@ -1,3 +1,5 @@
+//! This module contains functions to generate default trait impl function bodies where possible.
+
 use syntax::ast::{self, edit::AstNodeEdit, make, AstNode, NameOwner};
 use syntax::ted;
 
@@ -6,21 +8,17 @@ use syntax::ted;
 /// Returns `Option` so that we can use `?` rather than `if let Some`. Returning
 /// `None` means that generating a custom trait body failed, and the body will remain
 /// as `todo!` instead.
-pub(crate) fn gen_trait_body(
-    func: &ast::Fn,
-    trait_path: &ast::Path,
-    adt: &ast::Adt,
-    annotated_name: &ast::Name,
-) -> Option<()> {
+pub(crate) fn gen_trait_body(func: &ast::Fn, trait_path: &ast::Path, adt: &ast::Adt) -> Option<()> {
     match trait_path.segment()?.name_ref()?.text().as_str() {
-        "Debug" => gen_debug_impl(adt, func, annotated_name),
+        "Debug" => gen_debug_impl(adt, func),
         "Default" => gen_default_impl(adt, func),
         _ => None,
     }
 }
 
 /// Generate a `Debug` impl based on the fields and members of the target type.
-fn gen_debug_impl(adt: &ast::Adt, func: &ast::Fn, annotated_name: &ast::Name) -> Option<()> {
+fn gen_debug_impl(adt: &ast::Adt, func: &ast::Fn) -> Option<()> {
+    let annotated_name = adt.name()?;
     match adt {
         // `Debug` cannot be derived for unions, so no default impl can be provided.
         ast::Adt::Union(_) => None,

crates/rust-analyzer/tests/slow-tests/tidy.rs

@@ -280,6 +280,7 @@ fn check_todo(path: &Path, text: &str) {
         "ast/make.rs",
         // The documentation in string literals may contain anything for its own purposes
         "ide_db/src/helpers/generated_lints.rs",
+        "ide_assists/src/utils/gen_trait_body.rs",
         "ide_assists/src/tests/generated.rs",
     ];
     if need_todo.iter().any(|p| path.ends_with(p)) {