Add missing AssocItems in add_custom_impl assist

5 files changed, 240 insertions, 110 deletions

diff --git a/crates/assists/src/handlers/add_custom_impl.rs b/crates/assists/src/handlers/add_custom_impl.rs
index 669dd9b21..c13493fd8 100644
--- a/crates/assists/src/handlers/add_custom_impl.rs
+++ b/crates/assists/src/handlers/add_custom_impl.rs
@@ -4,13 +4,15 @@ use syntax::{
     ast::{self, make, AstNode},
     Direction, SmolStr,
     SyntaxKind::{IDENT, WHITESPACE},
-    TextRange, TextSize,
+    TextSize,
 };
 
 use crate::{
-    assist_config::SnippetCap,
     assist_context::{AssistBuilder, AssistContext, Assists},
-    utils::mod_path_to_ast,
+    utils::{
+        add_trait_assoc_items_to_impl, filter_assoc_items, mod_path_to_ast, render_snippet, Cursor,
+        DefaultMethods,
+    },
     AssistId, AssistKind,
 };
 
@@ -47,11 +49,10 @@ pub(crate) fn add_custom_impl(acc: &mut Assists, ctx: &AssistContext) -> Option<
         ctx.token_at_offset().find(|t| t.kind() == IDENT && *t.text() != attr_name)?;
     let trait_path = make::path_unqualified(make::path_segment(make::name_ref(trait_token.text())));
 
-    let annotated = attr.syntax().siblings(Direction::Next).find_map(ast::Name::cast)?;
-    let annotated_name = annotated.syntax().text().to_string();
-    let insert_pos = annotated.syntax().parent()?.text_range().end();
+    let annotated_name = attr.syntax().siblings(Direction::Next).find_map(ast::Name::cast)?;
+    let insert_pos = annotated_name.syntax().parent()?.text_range().end();
 
-    let current_module = ctx.sema.scope(annotated.syntax()).module()?;
+    let current_module = ctx.sema.scope(annotated_name.syntax()).module()?;
     let current_crate = current_module.krate();
 
     let found_traits = imports_locator::find_imports(&ctx.sema, current_crate, trait_token.text())
@@ -69,21 +70,22 @@ pub(crate) fn add_custom_impl(acc: &mut Assists, ctx: &AssistContext) -> Option<
         });
 
     let mut no_traits_found = true;
-    for (trait_path, _trait) in found_traits.inspect(|_| no_traits_found = false) {
-        add_assist(acc, ctx.config.snippet_cap, &attr, &trait_path, &annotated_name, insert_pos)?;
+    for (trait_path, trait_) in found_traits.inspect(|_| no_traits_found = false) {
+        add_assist(acc, ctx, &attr, &trait_path, Some(trait_), &annotated_name, insert_pos)?;
     }
     if no_traits_found {
-        add_assist(acc, ctx.config.snippet_cap, &attr, &trait_path, &annotated_name, insert_pos)?;
+        add_assist(acc, ctx, &attr, &trait_path, None, &annotated_name, insert_pos)?;
     }
     Some(())
 }
 
 fn add_assist(
     acc: &mut Assists,
-    snippet_cap: Option<SnippetCap>,
+    ctx: &AssistContext,
     attr: &ast::Attr,
     trait_path: &ast::Path,
-    annotated_name: &str,
+    trait_: Option<hir::Trait>,
+    annotated_name: &ast::Name,
     insert_pos: TextSize,
 ) -> Option<()> {
     let target = attr.syntax().text_range();
@@ -92,25 +94,62 @@ fn add_assist(
     let trait_name = trait_path.segment().and_then(|seg| seg.name_ref())?;
 
     acc.add(AssistId("add_custom_impl", AssistKind::Refactor), label, target, |builder| {
+        let impl_def_with_items =
+            impl_def_from_trait(&ctx.sema, annotated_name, trait_, trait_path);
         update_attribute(builder, &input, &trait_name, &attr);
-        match snippet_cap {
-            Some(cap) => {
+        match (ctx.config.snippet_cap, impl_def_with_items) {
+            (None, _) => builder.insert(
+                insert_pos,
+                format!("\n\nimpl {} for {} {{\n\n}}", trait_path, annotated_name),
+            ),
+            (Some(cap), None) => builder.insert_snippet(
+                cap,
+                insert_pos,
+                format!("\n\nimpl {} for {} {{\n    $0\n}}", trait_path, annotated_name),
+            ),
+            (Some(cap), Some((impl_def, first_assoc_item))) => {
+                let mut cursor = Cursor::Before(first_assoc_item.syntax());
+                let placeholder;
+                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!()" {
+                            placeholder = m;
+                            cursor = Cursor::Replace(placeholder.syntax());
+                        }
+                    }
+                }
+
                 builder.insert_snippet(
                     cap,
                     insert_pos,
-                    format!("\n\nimpl {} for {} {{\n    $0\n}}", trait_path, annotated_name),
-                );
-            }
-            None => {
-                builder.insert(
-                    insert_pos,
-                    format!("\n\nimpl {} for {} {{\n\n}}", trait_path, annotated_name),
-                );
+                    format!("\n\n{}", render_snippet(cap, impl_def.syntax(), cursor)),
+                )
             }
-        }
+        };
     })
 }
 
+fn impl_def_from_trait(
+    sema: &hir::Semantics<ide_db::RootDatabase>,
+    annotated_name: &ast::Name,
+    trait_: Option<hir::Trait>,
+    trait_path: &ast::Path,
+) -> Option<(ast::Impl, ast::AssocItem)> {
+    let trait_ = trait_?;
+    let target_scope = sema.scope(annotated_name.syntax());
+    let trait_items = filter_assoc_items(sema.db, &trait_.items(sema.db), DefaultMethods::No);
+    if trait_items.is_empty() {
+        return None;
+    }
+    let impl_def = make::impl_trait(
+        trait_path.clone(),
+        make::path_unqualified(make::path_segment(make::name_ref(annotated_name.text()))),
+    );
+    let (impl_def, first_assoc_item) =
+        add_trait_assoc_items_to_impl(sema, trait_items, trait_, impl_def, target_scope);
+    Some((impl_def, first_assoc_item))
+}
+
 fn update_attribute(
     builder: &mut AssistBuilder,
     input: &ast::TokenTree,
@@ -133,13 +172,14 @@ fn update_attribute(
         let attr_range = attr.syntax().text_range();
         builder.delete(attr_range);
 
-        let line_break_range = attr
+        if let Some(line_break_range) = attr
             .syntax()
             .next_sibling_or_token()
             .filter(|t| t.kind() == WHITESPACE)
             .map(|t| t.text_range())
-            .unwrap_or_else(|| TextRange::new(TextSize::from(0), TextSize::from(0)));
-        builder.delete(line_break_range);
+        {
+            builder.delete(line_break_range);
+        }
     }
 }
 
@@ -150,12 +190,17 @@ mod tests {
     use super::*;
 
     #[test]
-    fn add_custom_impl_qualified() {
+    fn add_custom_impl_debug() {
         check_assist(
             add_custom_impl,
             "
 mod fmt {
-    pub trait Debug {}
+    pub struct Error;
+    pub type Result = Result<(), Error>;
+    pub struct Formatter<'a>;
+    pub trait Debug {
+        fn fmt(&self, f: &mut Formatter<'_>) -> Result;
+    }
 }
 
 #[derive(Debu<|>g)]
@@ -165,7 +210,12 @@ struct Foo {
 ",
             "
 mod fmt {
-    pub trait Debug {}
+    pub struct Error;
+    pub type Result = Result<(), Error>;
+    pub struct Formatter<'a>;
+    pub trait Debug {
+        fn fmt(&self, f: &mut Formatter<'_>) -> Result;
+    }
 }
 
 struct Foo {
@@ -173,7 +223,58 @@ struct Foo {
 }
 
 impl fmt::Debug for Foo {
-    $0
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        ${0:todo!()}
+    }
+}
+",
+        )
+    }
+    #[test]
+    fn add_custom_impl_all() {
+        check_assist(
+            add_custom_impl,
+            "
+mod foo {
+    pub trait Bar {
+        type Qux;
+        const Baz: usize = 42;
+        const Fez: usize;
+        fn foo();
+        fn bar() {}
+    }
+}
+
+#[derive(<|>Bar)]
+struct Foo {
+    bar: String,
+}
+",
+            "
+mod foo {
+    pub trait Bar {
+        type Qux;
+        const Baz: usize = 42;
+        const Fez: usize;
+        fn foo();
+        fn bar() {}
+    }
+}
+
+struct Foo {
+    bar: String,
+}
+
+impl foo::Bar for Foo {
+    $0type Qux;
+
+    const Baz: usize = 42;
+
+    const Fez: usize;
+
+    fn foo() {
+        todo!()
+    }
 }
 ",
         )
diff --git a/crates/assists/src/handlers/add_missing_impl_members.rs b/crates/assists/src/handlers/add_missing_impl_members.rs
index b82fb30ad..bbb71e261 100644
--- a/crates/assists/src/handlers/add_missing_impl_members.rs
+++ b/crates/assists/src/handlers/add_missing_impl_members.rs
@@ -1,27 +1,14 @@
-use hir::HasSource;
-use ide_db::traits::{get_missing_assoc_items, resolve_target_trait};
-use syntax::{
-    ast::{
-        self,
-        edit::{self, AstNodeEdit, IndentLevel},
-        make, AstNode, NameOwner,
-    },
-    SmolStr,
-};
+use ide_db::traits::resolve_target_trait;
+use syntax::ast::{self, AstNode};
 
 use crate::{
     assist_context::{AssistContext, Assists},
-    ast_transform::{self, AstTransform, QualifyPaths, SubstituteTypeParams},
-    utils::{render_snippet, Cursor},
+    utils::add_trait_assoc_items_to_impl,
+    utils::DefaultMethods,
+    utils::{filter_assoc_items, render_snippet, Cursor},
     AssistId, AssistKind,
 };
 
-#[derive(PartialEq)]
-enum AddMissingImplMembersMode {
-    DefaultMethodsOnly,
-    NoDefaultMethods,
-}
-
 // Assist: add_impl_missing_members
 //
 // Adds scaffold for required impl members.
@@ -55,7 +42,7 @@ pub(crate) fn add_missing_impl_members(acc: &mut Assists, ctx: &AssistContext) -
     add_missing_impl_members_inner(
         acc,
         ctx,
-        AddMissingImplMembersMode::NoDefaultMethods,
+        DefaultMethods::No,
         "add_impl_missing_members",
         "Implement missing members",
     )
@@ -97,7 +84,7 @@ pub(crate) fn add_missing_default_members(acc: &mut Assists, ctx: &AssistContext
     add_missing_impl_members_inner(
         acc,
         ctx,
-        AddMissingImplMembersMode::DefaultMethodsOnly,
+        DefaultMethods::Only,
         "add_impl_default_members",
         "Implement default members",
     )
@@ -106,7 +93,7 @@ pub(crate) fn add_missing_default_members(acc: &mut Assists, ctx: &AssistContext
 fn add_missing_impl_members_inner(
     acc: &mut Assists,
     ctx: &AssistContext,
-    mode: AddMissingImplMembersMode,
+    mode: DefaultMethods,
     assist_id: &'static str,
     label: &'static str,
 ) -> Option<()> {
@@ -114,32 +101,11 @@ fn add_missing_impl_members_inner(
     let impl_def = ctx.find_node_at_offset::<ast::Impl>()?;
     let trait_ = resolve_target_trait(&ctx.sema, &impl_def)?;
 
-    let def_name = |item: &ast::AssocItem| -> Option<SmolStr> {
-        match item {
-            ast::AssocItem::Fn(def) => def.name(),
-            ast::AssocItem::TypeAlias(def) => def.name(),
-            ast::AssocItem::Const(def) => def.name(),
-            ast::AssocItem::MacroCall(_) => None,
-        }
-        .map(|it| it.text().clone())
-    };
-
-    let missing_items = get_missing_assoc_items(&ctx.sema, &impl_def)
-        .iter()
-        .map(|i| match i {
-            hir::AssocItem::Function(i) => ast::AssocItem::Fn(i.source(ctx.db()).value),
-            hir::AssocItem::TypeAlias(i) => ast::AssocItem::TypeAlias(i.source(ctx.db()).value),
-            hir::AssocItem::Const(i) => ast::AssocItem::Const(i.source(ctx.db()).value),
-        })
-        .filter(|t| def_name(&t).is_some())
-        .filter(|t| match t {
-            ast::AssocItem::Fn(def) => match mode {
-                AddMissingImplMembersMode::DefaultMethodsOnly => def.body().is_some(),
-                AddMissingImplMembersMode::NoDefaultMethods => def.body().is_none(),
-            },
-            _ => mode == AddMissingImplMembersMode::NoDefaultMethods,
-        })
-        .collect::<Vec<_>>();
+    let missing_items = filter_assoc_items(
+        ctx.db(),
+        &ide_db::traits::get_missing_assoc_items(&ctx.sema, &impl_def),
+        mode,
+    );
 
     if missing_items.is_empty() {
         return None;
@@ -147,29 +113,9 @@ fn add_missing_impl_members_inner(
 
     let target = impl_def.syntax().text_range();
     acc.add(AssistId(assist_id, AssistKind::QuickFix), label, target, |builder| {
-        let impl_item_list = impl_def.assoc_item_list().unwrap_or_else(make::assoc_item_list);
-
-        let n_existing_items = impl_item_list.assoc_items().count();
-        let source_scope = ctx.sema.scope_for_def(trait_);
         let target_scope = ctx.sema.scope(impl_def.syntax());
-        let ast_transform = QualifyPaths::new(&target_scope, &source_scope)
-            .or(SubstituteTypeParams::for_trait_impl(&source_scope, trait_, impl_def.clone()));
-
-        let items = missing_items
-            .into_iter()
-            .map(|it| ast_transform::apply(&*ast_transform, it))
-            .map(|it| match it {
-                ast::AssocItem::Fn(def) => ast::AssocItem::Fn(add_body(def)),
-                ast::AssocItem::TypeAlias(def) => ast::AssocItem::TypeAlias(def.remove_bounds()),
-                _ => it,
-            })
-            .map(|it| edit::remove_attrs_and_docs(&it));
-
-        let new_impl_item_list = impl_item_list.append_items(items);
-        let new_impl_def = impl_def.with_assoc_item_list(new_impl_item_list);
-        let first_new_item =
-            new_impl_def.assoc_item_list().unwrap().assoc_items().nth(n_existing_items).unwrap();
-
+        let (new_impl_def, first_new_item) =
+            add_trait_assoc_items_to_impl(&ctx.sema, missing_items, trait_, impl_def, target_scope);
         match ctx.config.snippet_cap {
             None => builder.replace(target, new_impl_def.to_string()),
             Some(cap) => {
@@ -193,14 +139,6 @@ fn add_missing_impl_members_inner(
     })
 }
 
-fn add_body(fn_def: ast::Fn) -> ast::Fn {
-    if fn_def.body().is_some() {
-        return fn_def;
-    }
-    let body = make::block_expr(None, Some(make::expr_todo())).indent(IndentLevel(1));
-    fn_def.with_body(body)
-}
-
 #[cfg(test)]
 mod tests {
     use crate::tests::{check_assist, check_assist_not_applicable};
diff --git a/crates/assists/src/utils.rs b/crates/assists/src/utils.rs
index 56f925ee6..7071fe96b 100644
--- a/crates/assists/src/utils.rs
+++ b/crates/assists/src/utils.rs
@@ -4,17 +4,22 @@ pub(crate) mod import_assets;
 
 use std::ops;
 
-use hir::{Crate, Enum, Module, ScopeDef, Semantics, Trait};
+use hir::{Crate, Enum, HasSource, Module, ScopeDef, Semantics, Trait};
 use ide_db::RootDatabase;
 use itertools::Itertools;
 use syntax::{
-    ast::{self, make, ArgListOwner},
+    ast::edit::AstNodeEdit,
+    ast::NameOwner,
+    ast::{self, edit, make, ArgListOwner},
     AstNode, Direction,
     SyntaxKind::*,
     SyntaxNode, TextSize, T,
 };
 
-use crate::assist_config::SnippetCap;
+use crate::{
+    assist_config::SnippetCap,
+    ast_transform::{self, AstTransform, QualifyPaths, SubstituteTypeParams},
+};
 
 pub use insert_use::MergeBehaviour;
 pub(crate) use insert_use::{insert_use, ImportScope};
@@ -77,6 +82,87 @@ pub fn extract_trivial_expression(block: &ast::BlockExpr) -> Option<ast::Expr> {
     None
 }
 
+#[derive(Copy, Clone, PartialEq)]
+pub enum DefaultMethods {
+    Only,
+    No,
+}
+
+pub fn filter_assoc_items(
+    db: &RootDatabase,
+    items: &[hir::AssocItem],
+    default_methods: DefaultMethods,
+) -> Vec<ast::AssocItem> {
+    fn has_def_name(item: &ast::AssocItem) -> bool {
+        match item {
+            ast::AssocItem::Fn(def) => def.name(),
+            ast::AssocItem::TypeAlias(def) => def.name(),
+            ast::AssocItem::Const(def) => def.name(),
+            ast::AssocItem::MacroCall(_) => None,
+        }
+        .is_some()
+    };
+
+    items
+        .iter()
+        .map(|i| match i {
+            hir::AssocItem::Function(i) => ast::AssocItem::Fn(i.source(db).value),
+            hir::AssocItem::TypeAlias(i) => ast::AssocItem::TypeAlias(i.source(db).value),
+            hir::AssocItem::Const(i) => ast::AssocItem::Const(i.source(db).value),
+        })
+        .filter(has_def_name)
+        .filter(|it| match it {
+            ast::AssocItem::Fn(def) => matches!(
+                (default_methods, def.body()),
+                (DefaultMethods::Only, Some(_)) | (DefaultMethods::No, None)
+            ),
+            _ => default_methods == DefaultMethods::No,
+        })
+        .collect::<Vec<_>>()
+}
+
+pub fn add_trait_assoc_items_to_impl(
+    sema: &hir::Semantics<ide_db::RootDatabase>,
+    items: Vec<ast::AssocItem>,
+    trait_: hir::Trait,
+    impl_def: ast::Impl,
+    target_scope: hir::SemanticsScope,
+) -> (ast::Impl, ast::AssocItem) {
+    let impl_item_list = impl_def.assoc_item_list().unwrap_or_else(make::assoc_item_list);
+
+    let n_existing_items = impl_item_list.assoc_items().count();
+    let source_scope = sema.scope_for_def(trait_);
+    let ast_transform = QualifyPaths::new(&target_scope, &source_scope)
+        .or(SubstituteTypeParams::for_trait_impl(&source_scope, trait_, impl_def.clone()));
+
+    let items = items
+        .into_iter()
+        .map(|it| ast_transform::apply(&*ast_transform, it))
+        .map(|it| match it {
+            ast::AssocItem::Fn(def) => ast::AssocItem::Fn(add_body(def)),
+            ast::AssocItem::TypeAlias(def) => ast::AssocItem::TypeAlias(def.remove_bounds()),
+            _ => it,
+        })
+        .map(|it| edit::remove_attrs_and_docs(&it));
+
+    let new_impl_item_list = impl_item_list.append_items(items);
+    let new_impl_def = impl_def.with_assoc_item_list(new_impl_item_list);
+    let first_new_item =
+        new_impl_def.assoc_item_list().unwrap().assoc_items().nth(n_existing_items).unwrap();
+    return (new_impl_def, first_new_item);
+
+    fn add_body(fn_def: ast::Fn) -> ast::Fn {
+        match fn_def.body() {
+            Some(_) => fn_def,
+            None => {
+                let body =
+                    make::block_expr(None, Some(make::expr_todo())).indent(edit::IndentLevel(1));
+                fn_def.with_body(body)
+            }
+        }
+    }
+}
+
 #[derive(Clone, Copy, Debug)]
 pub(crate) enum Cursor<'a> {
     Replace(&'a SyntaxNode),
diff --git a/crates/syntax/src/ast/make.rs b/crates/syntax/src/ast/make.rs
index b1578820f..876659a2b 100644
--- a/crates/syntax/src/ast/make.rs
+++ b/crates/syntax/src/ast/make.rs
@@ -25,6 +25,10 @@ pub fn assoc_item_list() -> ast::AssocItemList {
     ast_from_text("impl C for D {};")
 }
 
+pub fn impl_trait(trait_: ast::Path, ty: ast::Path) -> ast::Impl {
+    ast_from_text(&format!("impl {} for {} {{}}", trait_, ty))
+}
+
 pub fn path_segment(name_ref: ast::NameRef) -> ast::PathSegment {
     ast_from_text(&format!("use {};", name_ref))
 }
diff --git a/xtask/tests/tidy.rs b/xtask/tests/tidy.rs
index 9de60c76c..4c7db8405 100644
--- a/xtask/tests/tidy.rs
+++ b/xtask/tests/tidy.rs
@@ -215,6 +215,7 @@ fn check_todo(path: &Path, text: &str) {
         "tests/cli.rs",
         // Some of our assists generate `todo!()`.
         "tests/generated.rs",
+        "handlers/add_custom_impl.rs",
         "handlers/add_missing_impl_members.rs",
         "handlers/add_turbo_fish.rs",
         "handlers/generate_function.rs",