Merge #9830

9830: Enable more assists to generate default trait body impls r=Veykril a=yoshuawuyts

Enable more assists to benefit from trait body generation. Follow-up to #9825 and #9814.

__edit:__ I'd like to move the existing tests to this new file too, but I'll do that in a follow-up PR.

Co-authored-by: Yoshua Wuyts <yoshuawuyts@gmail.com>

Author: bors[bot]

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_fn_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::ext::ident_path(&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_fn_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

@@ -2,8 +2,6 @@ use hir::ModuleDef;
 use ide_db::helpers::{import_assets::NameToImport, mod_path_to_ast};
 use ide_db::items_locator;
 use itertools::Itertools;
-use syntax::ast::edit::AstNodeEdit;
-use syntax::ted;
 use syntax::{
     ast::{self, make, AstNode, NameOwner},
     SyntaxKind::{IDENT, WHITESPACE},
@@ -12,8 +10,8 @@ use syntax::{
 use crate::{
     assist_context::{AssistBuilder, AssistContext, Assists},
     utils::{
-        add_trait_assoc_items_to_impl, filter_assoc_items, generate_trait_impl_text,
-        render_snippet, Cursor, DefaultMethods,
+        add_trait_assoc_items_to_impl, filter_assoc_items, gen_trait_fn_body,
+        generate_trait_impl_text, render_snippet, Cursor, DefaultMethods,
     },
     AssistId, AssistKind,
 };
@@ -169,158 +167,12 @@ 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_impl(func, trait_path, adt, annotated_name);
+        let _ = gen_trait_fn_body(func, trait_path, adt);
     };
 
     Some((impl_def, first_assoc_item))
 }
 
-/// Generate custom trait bodies where possible.
-///
-/// 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.
-fn gen_trait_body_impl(
-    func: &ast::Fn,
-    trait_path: &ast::Path,
-    adt: &ast::Adt,
-    annotated_name: &ast::Name,
-) -> Option<()> {
-    match trait_path.segment()?.name_ref()?.text().as_str() {
-        "Debug" => gen_debug_impl(adt, func, annotated_name),
-        "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<()> {
-    match adt {
-        // `Debug` cannot be derived for unions, so no default impl can be provided.
-        ast::Adt::Union(_) => None,
-
-        // => match self { Self::Variant => write!(f, "Variant") }
-        ast::Adt::Enum(enum_) => {
-            let list = enum_.variant_list()?;
-            let mut arms = vec![];
-            for variant in list.variants() {
-                let name = variant.name()?;
-                let left = make::ext::ident_path("Self");
-                let right = make::ext::ident_path(&format!("{}", name));
-                let variant_name = make::path_pat(make::path_concat(left, right));
-
-                let target = make::expr_path(make::ext::ident_path("f").into());
-                let fmt_string = make::expr_literal(&(format!("\"{}\"", name))).into();
-                let args = make::arg_list(vec![target, fmt_string]);
-                let macro_name = make::expr_path(make::ext::ident_path("write"));
-                let macro_call = make::expr_macro_call(macro_name, args);
-
-                arms.push(make::match_arm(Some(variant_name.into()), None, macro_call.into()));
-            }
-
-            let match_target = make::expr_path(make::ext::ident_path("self"));
-            let list = make::match_arm_list(arms).indent(ast::edit::IndentLevel(1));
-            let match_expr = make::expr_match(match_target, list);
-
-            let body = make::block_expr(None, Some(match_expr));
-            let body = body.indent(ast::edit::IndentLevel(1));
-            ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
-            Some(())
-        }
-
-        ast::Adt::Struct(strukt) => {
-            let name = format!("\"{}\"", annotated_name);
-            let args = make::arg_list(Some(make::expr_literal(&name).into()));
-            let target = make::expr_path(make::ext::ident_path("f"));
-
-            let expr = match strukt.field_list() {
-                // => f.debug_struct("Name").finish()
-                None => make::expr_method_call(target, make::name_ref("debug_struct"), args),
-
-                // => f.debug_struct("Name").field("foo", &self.foo).finish()
-                Some(ast::FieldList::RecordFieldList(field_list)) => {
-                    let method = make::name_ref("debug_struct");
-                    let mut expr = make::expr_method_call(target, method, args);
-                    for field in field_list.fields() {
-                        let name = field.name()?;
-                        let f_name = make::expr_literal(&(format!("\"{}\"", name))).into();
-                        let f_path = make::expr_path(make::ext::ident_path("self"));
-                        let f_path = make::expr_ref(f_path, false);
-                        let f_path = make::expr_field(f_path, &format!("{}", name)).into();
-                        let args = make::arg_list(vec![f_name, f_path]);
-                        expr = make::expr_method_call(expr, make::name_ref("field"), args);
-                    }
-                    expr
-                }
-
-                // => f.debug_tuple("Name").field(self.0).finish()
-                Some(ast::FieldList::TupleFieldList(field_list)) => {
-                    let method = make::name_ref("debug_tuple");
-                    let mut expr = make::expr_method_call(target, method, args);
-                    for (idx, _) in field_list.fields().enumerate() {
-                        let f_path = make::expr_path(make::ext::ident_path("self"));
-                        let f_path = make::expr_ref(f_path, false);
-                        let f_path = make::expr_field(f_path, &format!("{}", idx)).into();
-                        let method = make::name_ref("field");
-                        expr = make::expr_method_call(expr, method, make::arg_list(Some(f_path)));
-                    }
-                    expr
-                }
-            };
-
-            let method = make::name_ref("finish");
-            let expr = make::expr_method_call(expr, method, make::arg_list(None));
-            let body = make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1));
-            ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
-            Some(())
-        }
-    }
-}
-
-/// 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_call() -> ast::Expr {
-        let trait_name = make::ext::ident_path("Default");
-        let method_name = make::ext::ident_path("default");
-        let fn_name = make::expr_path(make::path_concat(trait_name, method_name));
-        make::expr_call(fn_name, make::arg_list(None))
-    }
-    match adt {
-        // `Debug` cannot be derived for unions, so no default impl can be provided.
-        ast::Adt::Union(_) => None,
-        // Deriving `Debug` for enums is not stable yet.
-        ast::Adt::Enum(_) => None,
-        ast::Adt::Struct(strukt) => {
-            let expr = match strukt.field_list() {
-                Some(ast::FieldList::RecordFieldList(field_list)) => {
-                    let mut fields = vec![];
-                    for field in field_list.fields() {
-                        let method_call = gen_default_call();
-                        let name_ref = make::name_ref(&field.name()?.to_string());
-                        let field = make::record_expr_field(name_ref, Some(method_call));
-                        fields.push(field);
-                    }
-                    let struct_name = make::ext::ident_path("Self");
-                    let fields = make::record_expr_field_list(fields);
-                    make::record_expr(struct_name, fields).into()
-                }
-                Some(ast::FieldList::TupleFieldList(field_list)) => {
-                    let struct_name = make::expr_path(make::ext::ident_path("Self"));
-                    let fields = field_list.fields().map(|_| gen_default_call());
-                    make::expr_call(struct_name, make::arg_list(fields))
-                }
-                None => {
-                    let struct_name = make::ext::ident_path("Self");
-                    let fields = make::record_expr_field_list(None);
-                    make::record_expr(struct_name, fields).into()
-                }
-            };
-            let body = make::block_expr(None, Some(expr)).indent(ast::edit::IndentLevel(1));
-            ted::replace(func.body()?.syntax(), body.clone_for_update().syntax());
-            Some(())
-        }
-    }
-}
 fn update_attribute(
     builder: &mut AssistBuilder,
     input: &ast::TokenTree,

crates/ide_assists/src/utils.rs

@@ -1,6 +1,7 @@
 //! Assorted functions shared by several assists.
 
 pub(crate) mod suggest_name;
+mod gen_trait_fn_body;
 
 use std::ops;
 
@@ -25,6 +26,8 @@ use syntax::{
 
 use crate::assist_context::{AssistBuilder, AssistContext};
 
+pub(crate) use gen_trait_fn_body::gen_trait_fn_body;
+
 pub(crate) fn unwrap_trivial_block(block: ast::BlockExpr) -> ast::Expr {
     extract_trivial_expression(&block)
         .filter(|expr| !expr.syntax().text().contains_char('\n'))