Modify around add_trait_assoc_items_to_impl to migrate add_missing_impl_members

Signed-off-by: Hayashi Mikihiro <34ttrweoewiwe28@gmail.com>

Author: Hmikihiro

crates/ide-assists/src/handlers/add_missing_impl_members.rs

@@ -2,7 +2,7 @@ use hir::HasSource;
 use syntax::{
     Edition,
     ast::{self, AstNode, make},
-    ted,
+    syntax_editor::{Position, SyntaxEditor},
 };
 
 use crate::{
@@ -148,31 +148,62 @@ fn add_missing_impl_members_inner(
 
     let target = impl_def.syntax().text_range();
     acc.add(AssistId::quick_fix(assist_id), label, target, |edit| {
-        let new_impl_def = edit.make_mut(impl_def.clone());
-        let first_new_item = add_trait_assoc_items_to_impl(
+        let new_item = add_trait_assoc_items_to_impl(
             &ctx.sema,
             ctx.config,
             &missing_items,
             trait_,
-            &new_impl_def,
+            &impl_def,
             &target_scope,
         );
 
+        let Some((first_new_item, other_items)) = new_item.split_first() else {
+            return;
+        };
+
+        let mut first_new_item = if let DefaultMethods::No = mode
+            && let ast::AssocItem::Fn(func) = &first_new_item
+            && let Some(body) = try_gen_trait_body(
+                ctx,
+                func,
+                trait_ref,
+                &impl_def,
+                target_scope.krate().edition(ctx.sema.db),
+            )
+            && let Some(func_body) = func.body()
+        {
+            let mut func_editor = SyntaxEditor::new(first_new_item.syntax().clone_subtree());
+            func_editor.replace(func_body.syntax(), body.syntax());
+            ast::AssocItem::cast(func_editor.finish().new_root().clone())
+        } else {
+            Some(first_new_item.clone())
+        };
+
+        let new_assoc_items = first_new_item
+            .clone()
+            .into_iter()
+            .chain(other_items.iter().cloned())
+            .map(either::Either::Right)
+            .collect::<Vec<_>>();
+
+        let mut editor = edit.make_editor(impl_def.syntax());
+        if let Some(assoc_item_list) = impl_def.assoc_item_list() {
+            let items = new_assoc_items.into_iter().filter_map(either::Either::right).collect();
+            assoc_item_list.add_items(&mut editor, items);
+        } else {
+            let assoc_item_list = make::assoc_item_list(Some(new_assoc_items)).clone_for_update();
+            editor.insert_all(
+                Position::after(impl_def.syntax()),
+                vec![make::tokens::whitespace(" ").into(), assoc_item_list.syntax().clone().into()],
+            );
+            first_new_item = assoc_item_list.assoc_items().next();
+        }
+
         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 let Some(body) = try_gen_trait_body(
-                        ctx,
-                        func,
-                        trait_ref,
-                        &impl_def,
-                        target_scope.krate().edition(ctx.sema.db),
-                    ) && 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 let Some(ast::AssocItem::Fn(func)) = &first_new_item {
+                    if let Some(m) = func.syntax().descendants().find_map(ast::MacroCall::cast)
                         && m.syntax().text() == "todo!()"
                     {
                         placeholder = Some(m);
@@ -181,11 +212,14 @@ fn add_missing_impl_members_inner(
             }
 
             if let Some(macro_call) = placeholder {
-                edit.add_placeholder_snippet(cap, macro_call);
-            } else {
-                edit.add_tabstop_before(cap, first_new_item);
+                let placeholder = edit.make_placeholder_snippet(cap);
+                editor.add_annotation(macro_call.syntax(), placeholder);
+            } else if let Some(first_new_item) = first_new_item {
+                let tabstop = edit.make_tabstop_before(cap);
+                editor.add_annotation(first_new_item.syntax(), tabstop);
             };
         };
+        edit.add_file_edits(ctx.vfs_file_id(), editor);
     })
 }
 
@@ -322,7 +356,7 @@ impl Foo for S {
     }
 
     #[test]
-    fn test_impl_def_without_braces() {
+    fn test_impl_def_without_braces_macro() {
         check_assist(
             add_missing_impl_members,
             r#"
@@ -340,6 +374,33 @@ impl Foo for S {
         );
     }
 
+    #[test]
+    fn test_impl_def_without_braces_tabstop_first_item() {
+        check_assist(
+            add_missing_impl_members,
+            r#"
+trait Foo {
+    type Output;
+    fn foo(&self);
+}
+struct S;
+impl Foo for S { $0 }"#,
+            r#"
+trait Foo {
+    type Output;
+    fn foo(&self);
+}
+struct S;
+impl Foo for S {
+    $0type Output;
+
+    fn foo(&self) {
+        todo!()
+    }
+}"#,
+        );
+    }
+
     #[test]
     fn fill_in_type_params_1() {
         check_assist(

crates/ide-assists/src/handlers/generate_impl.rs

@@ -167,33 +167,44 @@ pub(crate) fn generate_impl_trait(acc: &mut Assists, ctx: &AssistContext<'_>) ->
                 DefaultMethods::No,
                 IgnoreAssocItems::DocHiddenAttrPresent,
             );
-            let impl_ = make::impl_trait(
-                trait_.unsafe_token().is_some(),
-                None,
-                trait_.generic_param_list().map(|list| {
-                    make::generic_arg_list(list.generic_params().map(|_| holder_arg.clone()))
-                }),
-                None,
-                None,
-                false,
-                make::ty(&name.text()),
-                make::ty_placeholder(),
-                None,
-                None,
-                None,
-            )
-            .clone_for_update();
-
-            if !missing_items.is_empty() {
-                utils::add_trait_assoc_items_to_impl(
+
+            let trait_gen_args = trait_.generic_param_list().map(|list| {
+                make::generic_arg_list(list.generic_params().map(|_| holder_arg.clone()))
+            });
+
+            let make_impl_ = |body| {
+                make::impl_trait(
+                    trait_.unsafe_token().is_some(),
+                    None,
+                    trait_gen_args.clone(),
+                    None,
+                    None,
+                    false,
+                    make::ty(&name.text()),
+                    make::ty_placeholder(),
+                    None,
+                    None,
+                    body,
+                )
+                .clone_for_update()
+            };
+
+            let impl_ = if missing_items.is_empty() {
+                make_impl_(None)
+            } else {
+                let impl_ = make_impl_(None);
+                let assoc_items = utils::add_trait_assoc_items_to_impl(
                     &ctx.sema,
                     ctx.config,
                     &missing_items,
                     hir_trait,
                     &impl_,
                     &target_scope,
                 );
-            }
+                let assoc_items = assoc_items.into_iter().map(either::Either::Right).collect();
+                let assoc_item_list = make::assoc_item_list(Some(assoc_items));
+                make_impl_(Some(assoc_item_list))
+            };
 
             if let Some(cap) = ctx.config.snippet_cap {
                 if let Some(generics) = impl_.trait_().and_then(|it| it.generic_arg_list()) {

crates/ide-assists/src/handlers/replace_derive_with_manual_impl.rs

@@ -206,19 +206,35 @@ fn impl_def_from_trait(
     }
     let impl_def = generate_trait_impl(impl_is_unsafe, adt, make::ty_path(trait_path.clone()));
 
-    let _ =
+    let assoc_items =
         add_trait_assoc_items_to_impl(sema, config, &trait_items, trait_, &impl_def, &target_scope);
-    let impl_def = impl_def.clone_subtree();
-    let mut editor = SyntaxEditor::new(impl_def.syntax().clone());
-    let first_assoc_item = impl_def.assoc_item_list().and_then(|item| item.assoc_items().next())?;
-    // Generate a default `impl` function body for the derived trait.
-    if let ast::AssocItem::Fn(ref func) = first_assoc_item {
-        if let Some(body) = gen_trait_fn_body(func, trait_path, adt, None)
+    let assoc_item_list = if let Some((first, other)) =
+        assoc_items.split_first().map(|(first, other)| (first.clone_subtree(), other))
+    {
+        let first_item = if let ast::AssocItem::Fn(ref func) = first
+            && let Some(body) = gen_trait_fn_body(func, trait_path, adt, None)
             && let Some(func_body) = func.body()
         {
+            let mut editor = SyntaxEditor::new(first.syntax().clone());
             editor.replace(func_body.syntax(), body.syntax());
-        }
-    };
+            ast::AssocItem::cast(editor.finish().new_root().clone())
+        } else {
+            Some(first.clone())
+        };
+        let items = first_item
+            .into_iter()
+            .chain(other.iter().cloned())
+            .map(either::Either::Right)
+            .collect();
+        make::assoc_item_list(Some(items))
+    } else {
+        make::assoc_item_list(None)
+    }
+    .clone_for_update();
+
+    let impl_def = impl_def.clone_subtree();
+    let mut editor = SyntaxEditor::new(impl_def.syntax().clone());
+    editor.replace(impl_def.assoc_item_list()?.syntax(), assoc_item_list.syntax());
     let impl_def = ast::Impl::cast(editor.finish().new_root().clone())?;
     Some(impl_def)
 }

crates/ide-assists/src/utils.rs

@@ -23,7 +23,7 @@ use syntax::{
     ast::{
         self, HasArgList, HasAttrs, HasGenericParams, HasName, HasTypeBounds, Whitespace,
         edit::{AstNodeEdit, IndentLevel},
-        edit_in_place::{AttrsOwnerEdit, Indent, Removable},
+        edit_in_place::{AttrsOwnerEdit, Removable},
         make,
         syntax_factory::SyntaxFactory,
     },
@@ -178,78 +178,73 @@ pub fn filter_assoc_items(
 /// [`filter_assoc_items()`]), clones each item for update and applies path transformation to it,
 /// then inserts into `impl_`. Returns the modified `impl_` and the first associated item that got
 /// inserted.
+#[must_use]
 pub fn add_trait_assoc_items_to_impl(
     sema: &Semantics<'_, RootDatabase>,
     config: &AssistConfig,
     original_items: &[InFile<ast::AssocItem>],
     trait_: hir::Trait,
     impl_: &ast::Impl,
     target_scope: &hir::SemanticsScope<'_>,
-) -> ast::AssocItem {
+) -> Vec<ast::AssocItem> {
     let new_indent_level = IndentLevel::from_node(impl_.syntax()) + 1;
-    let items = original_items.iter().map(|InFile { file_id, value: original_item }| {
-        let cloned_item = {
-            if let Some(macro_file) = file_id.macro_file() {
-                let span_map = sema.db.expansion_span_map(macro_file);
-                let item_prettified = prettify_macro_expansion(
-                    sema.db,
-                    original_item.syntax().clone(),
-                    &span_map,
-                    target_scope.krate().into(),
-                );
-                if let Some(formatted) = ast::AssocItem::cast(item_prettified) {
-                    return formatted;
-                } else {
-                    stdx::never!("formatted `AssocItem` could not be cast back to `AssocItem`");
+    original_items
+        .iter()
+        .map(|InFile { file_id, value: original_item }| {
+            let cloned_item = {
+                if let Some(macro_file) = file_id.macro_file() {
+                    let span_map = sema.db.expansion_span_map(macro_file);
+                    let item_prettified = prettify_macro_expansion(
+                        sema.db,
+                        original_item.syntax().clone(),
+                        &span_map,
+                        target_scope.krate().into(),
+                    );
+                    if let Some(formatted) = ast::AssocItem::cast(item_prettified) {
+                        return formatted;
+                    } else {
+                        stdx::never!("formatted `AssocItem` could not be cast back to `AssocItem`");
+                    }
                 }
-            }
-            original_item.clone_for_update()
-        };
-
-        if let Some(source_scope) = sema.scope(original_item.syntax()) {
-            // FIXME: Paths in nested macros are not handled well. See
-            // `add_missing_impl_members::paths_in_nested_macro_should_get_transformed` test.
-            let transform =
-                PathTransform::trait_impl(target_scope, &source_scope, trait_, impl_.clone());
-            transform.apply(cloned_item.syntax());
-        }
-        cloned_item.remove_attrs_and_docs();
-        cloned_item.reindent_to(new_indent_level);
-        cloned_item
-    });
+                original_item.clone_for_update()
+            };
 
-    let assoc_item_list = impl_.get_or_create_assoc_item_list();
-
-    let mut first_item = None;
-    for item in items {
-        first_item.get_or_insert_with(|| item.clone());
-        match &item {
-            ast::AssocItem::Fn(fn_) if fn_.body().is_none() => {
-                let body = AstNodeEdit::indent(
-                    &make::block_expr(
-                        None,
-                        Some(match config.expr_fill_default {
-                            ExprFillDefaultMode::Todo => make::ext::expr_todo(),
-                            ExprFillDefaultMode::Underscore => make::ext::expr_underscore(),
-                            ExprFillDefaultMode::Default => make::ext::expr_todo(),
-                        }),
-                    ),
-                    new_indent_level,
-                );
-                ted::replace(fn_.get_or_create_body().syntax(), body.syntax())
+            if let Some(source_scope) = sema.scope(original_item.syntax()) {
+                // FIXME: Paths in nested macros are not handled well. See
+                // `add_missing_impl_members::paths_in_nested_macro_should_get_transformed` test.
+                let transform =
+                    PathTransform::trait_impl(target_scope, &source_scope, trait_, impl_.clone());
+                transform.apply(cloned_item.syntax());
             }
-            ast::AssocItem::TypeAlias(type_alias) => {
-                if let Some(type_bound_list) = type_alias.type_bound_list() {
-                    type_bound_list.remove()
+            cloned_item.remove_attrs_and_docs();
+            cloned_item.reset_indent()
+        })
+        .map(|item| {
+            match &item {
+                ast::AssocItem::Fn(fn_) if fn_.body().is_none() => {
+                    let body = AstNodeEdit::indent(
+                        &make::block_expr(
+                            None,
+                            Some(match config.expr_fill_default {
+                                ExprFillDefaultMode::Todo => make::ext::expr_todo(),
+                                ExprFillDefaultMode::Underscore => make::ext::expr_underscore(),
+                                ExprFillDefaultMode::Default => make::ext::expr_todo(),
+                            }),
+                        ),
+                        IndentLevel::single(),
+                    );
+                    ted::replace(fn_.get_or_create_body().syntax(), body.syntax());
                 }
+                ast::AssocItem::TypeAlias(type_alias) => {
+                    if let Some(type_bound_list) = type_alias.type_bound_list() {
+                        type_bound_list.remove()
+                    }
+                }
+                _ => {}
             }
-            _ => {}
-        }
-
-        assoc_item_list.add_item(item)
-    }
-
-    first_item.unwrap()
+            AstNodeEdit::indent(&item, new_indent_level)
+        })
+        .collect()
 }
 
 pub(crate) fn vis_offset(node: &SyntaxNode) -> TextSize {