Add missing AssocItems in add_custom_impl assist

1 files changed, 15 insertions, 77 deletions

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};