Merge #8837

8837: internal: rewrite assoc item manipulaion to use mutable trees r=matklad a=matklad

bors r+
🤖

Co-authored-by: Aleksey Kladov <[email protected]>

7 files changed, 187 insertions, 221 deletions

diff --git a/crates/ide_assists/src/handlers/add_missing_impl_members.rs b/crates/ide_assists/src/handlers/add_missing_impl_members.rs
index 0148635f9..8225ae22c 100644
--- a/crates/ide_assists/src/handlers/add_missing_impl_members.rs
+++ b/crates/ide_assists/src/handlers/add_missing_impl_members.rs
@@ -64,7 +64,6 @@ pub(crate) fn add_missing_impl_members(acc: &mut Assists, ctx: &AssistContext) -
 // impl Trait for () {
 //     type X = ();
 //     fn foo(&self) {}$0
-//
 // }
 // ```
 // ->
@@ -195,6 +194,7 @@ impl Foo for S {
     fn baz(&self) {
         todo!()
     }
+
 }"#,
         );
     }
@@ -231,6 +231,7 @@ impl Foo for S {
     fn foo(&self) {
         ${0:todo!()}
     }
+
 }"#,
         );
     }
diff --git a/crates/ide_assists/src/handlers/generate_new.rs b/crates/ide_assists/src/handlers/generate_new.rs
index 8ce5930b7..959a1f86c 100644
--- a/crates/ide_assists/src/handlers/generate_new.rs
+++ b/crates/ide_assists/src/handlers/generate_new.rs
@@ -1,4 +1,3 @@
-use ast::Adt;
 use itertools::Itertools;
 use stdx::format_to;
 use syntax::ast::{self, AstNode, NameOwner, StructKind, VisibilityOwner};
@@ -37,7 +36,7 @@ pub(crate) fn generate_new(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
     };
 
     // Return early if we've found an existing new fn
-    let impl_def = find_struct_impl(&ctx, &Adt::Struct(strukt.clone()), "new")?;
+    let impl_def = find_struct_impl(&ctx, &ast::Adt::Struct(strukt.clone()), "new")?;
 
     let target = strukt.syntax().text_range();
     acc.add(AssistId("generate_new", AssistKind::Generate), "Generate `new`", target, |builder| {
@@ -60,7 +59,7 @@ pub(crate) fn generate_new(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
         let start_offset = impl_def
             .and_then(|impl_def| find_impl_block_start(impl_def, &mut buf))
             .unwrap_or_else(|| {
-                buf = generate_impl_text(&Adt::Struct(strukt.clone()), &buf);
+                buf = generate_impl_text(&ast::Adt::Struct(strukt.clone()), &buf);
                 strukt.syntax().text_range().end()
             });
 
@@ -81,101 +80,132 @@ mod tests {
     use super::*;
 
     #[test]
-    #[rustfmt::skip]
     fn test_generate_new() {
-        // Check output of generation
         check_assist(
             generate_new,
-"struct Foo {$0}",
-"struct Foo {}
+            r#"
+struct Foo {$0}
+"#,
+            r#"
+struct Foo {}
 
 impl Foo {
     fn $0new() -> Self { Self {  } }
-}",
+}
+"#,
         );
         check_assist(
             generate_new,
-"struct Foo<T: Clone> {$0}",
-"struct Foo<T: Clone> {}
+            r#"
+struct Foo<T: Clone> {$0}
+"#,
+            r#"
+struct Foo<T: Clone> {}
 
 impl<T: Clone> Foo<T> {
     fn $0new() -> Self { Self {  } }
-}",
+}
+"#,
         );
         check_assist(
             generate_new,
-"struct Foo<'a, T: Foo<'a>> {$0}",
-"struct Foo<'a, T: Foo<'a>> {}
+            r#"
+struct Foo<'a, T: Foo<'a>> {$0}
+"#,
+            r#"
+struct Foo<'a, T: Foo<'a>> {}
 
 impl<'a, T: Foo<'a>> Foo<'a, T> {
     fn $0new() -> Self { Self {  } }
-}",
+}
+"#,
         );
         check_assist(
             generate_new,
-"struct Foo { baz: String $0}",
-"struct Foo { baz: String }
+            r#"
+struct Foo { baz: String $0}
+"#,
+            r#"
+struct Foo { baz: String }
 
 impl Foo {
     fn $0new(baz: String) -> Self { Self { baz } }
-}",
+}
+"#,
         );
         check_assist(
             generate_new,
-"struct Foo { baz: String, qux: Vec<i32> $0}",
-"struct Foo { baz: String, qux: Vec<i32> }
+            r#"
+struct Foo { baz: String, qux: Vec<i32> $0}
+"#,
+            r#"
+struct Foo { baz: String, qux: Vec<i32> }
 
 impl Foo {
     fn $0new(baz: String, qux: Vec<i32>) -> Self { Self { baz, qux } }
-}",
+}
+"#,
         );
+    }
 
-        // Check that visibility modifiers don't get brought in for fields
+    #[test]
+    fn check_that_visibility_modifiers_dont_get_brought_in() {
         check_assist(
             generate_new,
-"struct Foo { pub baz: String, pub qux: Vec<i32> $0}",
-"struct Foo { pub baz: String, pub qux: Vec<i32> }
+            r#"
+struct Foo { pub baz: String, pub qux: Vec<i32> $0}
+"#,
+            r#"
+struct Foo { pub baz: String, pub qux: Vec<i32> }
 
 impl Foo {
     fn $0new(baz: String, qux: Vec<i32>) -> Self { Self { baz, qux } }
-}",
+}
+"#,
         );
+    }
 
-        // Check that it reuses existing impls
+    #[test]
+    fn check_it_reuses_existing_impls() {
         check_assist(
             generate_new,
-"struct Foo {$0}
+            r#"
+struct Foo {$0}
 
 impl Foo {}
-",
-"struct Foo {}
+"#,
+            r#"
+struct Foo {}
 
 impl Foo {
     fn $0new() -> Self { Self {  } }
 }
-",
+"#,
         );
         check_assist(
             generate_new,
-"struct Foo {$0}
+            r#"
+struct Foo {$0}
 
 impl Foo {
     fn qux(&self) {}
 }
-",
-"struct Foo {}
+"#,
+            r#"
+struct Foo {}
 
 impl Foo {
     fn $0new() -> Self { Self {  } }
 
     fn qux(&self) {}
 }
-",
+"#,
         );
 
         check_assist(
             generate_new,
-"struct Foo {$0}
+            r#"
+struct Foo {$0}
 
 impl Foo {
     fn qux(&self) {}
@@ -183,8 +213,9 @@ impl Foo {
         5
     }
 }
-",
-"struct Foo {}
+"#,
+            r#"
+struct Foo {}
 
 impl Foo {
     fn $0new() -> Self { Self {  } }
@@ -194,27 +225,37 @@ impl Foo {
         5
     }
 }
-",
+"#,
         );
+    }
 
-        // Check visibility of new fn based on struct
+    #[test]
+    fn check_visibility_of_new_fn_based_on_struct() {
         check_assist(
             generate_new,
-"pub struct Foo {$0}",
-"pub struct Foo {}
+            r#"
+pub struct Foo {$0}
+"#,
+            r#"
+pub struct Foo {}
 
 impl Foo {
     pub fn $0new() -> Self { Self {  } }
-}",
+}
+"#,
         );
         check_assist(
             generate_new,
-"pub(crate) struct Foo {$0}",
-"pub(crate) struct Foo {}
+            r#"
+pub(crate) struct Foo {$0}
+"#,
+            r#"
+pub(crate) struct Foo {}
 
 impl Foo {
     pub(crate) fn $0new() -> Self { Self {  } }
-}",
+}
+"#,
         );
     }
 
@@ -222,26 +263,28 @@ impl Foo {
     fn generate_new_not_applicable_if_fn_exists() {
         check_assist_not_applicable(
             generate_new,
-            "
+            r#"
 struct Foo {$0}
 
 impl Foo {
     fn new() -> Self {
         Self
     }
-}",
+}
+"#,
         );
 
         check_assist_not_applicable(
             generate_new,
-            "
+            r#"
 struct Foo {$0}
 
 impl Foo {
     fn New() -> Self {
         Self
     }
-}",
+}
+"#,
         );
     }
 
@@ -249,12 +292,12 @@ impl Foo {
     fn generate_new_target() {
         check_assist_target(
             generate_new,
-            "
+            r#"
 struct SomeThingIrrelevant;
 /// Has a lifetime parameter
 struct Foo<'a, T: Foo<'a>> {$0}
 struct EvenMoreIrrelevant;
-",
+"#,
             "/// Has a lifetime parameter
 struct Foo<'a, T: Foo<'a>> {}",
         );
@@ -264,7 +307,7 @@ struct Foo<'a, T: Foo<'a>> {}",
     fn test_unrelated_new() {
         check_assist(
             generate_new,
-            r##"
+            r#"
 pub struct AstId<N: AstNode> {
     file_id: HirFileId,
     file_ast_id: FileAstId<N>,
@@ -285,8 +328,9 @@ impl<T> Source<T> {
     pub fn map<F: FnOnce(T) -> U, U>(self, f: F) -> Source<U> {
         Source { file_id: self.file_id, ast: f(self.ast) }
     }
-}"##,
-            r##"
+}
+"#,
+            r#"
 pub struct AstId<N: AstNode> {
     file_id: HirFileId,
     file_ast_id: FileAstId<N>,
@@ -309,7 +353,8 @@ impl<T> Source<T> {
     pub fn map<F: FnOnce(T) -> U, U>(self, f: F) -> Source<U> {
         Source { file_id: self.file_id, ast: f(self.ast) }
     }
-}"##,
+}
+"#,
         );
     }
 }
diff --git a/crates/ide_assists/src/tests/generated.rs b/crates/ide_assists/src/tests/generated.rs
index 49c1a9776..4406406a2 100644
--- a/crates/ide_assists/src/tests/generated.rs
+++ b/crates/ide_assists/src/tests/generated.rs
@@ -50,7 +50,6 @@ trait Trait {
 impl Trait for () {
     type X = ();
     fn foo(&self) {}$0
-
 }
 "#####,
         r#####"
diff --git a/crates/ide_assists/src/utils.rs b/crates/ide_assists/src/utils.rs
index 0dcf20c61..7d562d1d4 100644
--- a/crates/ide_assists/src/utils.rs
+++ b/crates/ide_assists/src/utils.rs
@@ -128,15 +128,12 @@ 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,
+    impl_: 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()));
+        .or(SubstituteTypeParams::for_trait_impl(&source_scope, trait_, impl_.clone()));
 
     let items = items
         .into_iter()
@@ -147,13 +144,18 @@ pub fn add_trait_assoc_items_to_impl(
             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);
+        .map(|it| edit::remove_attrs_and_docs(&it).clone_subtree().clone_for_update());
+
+    let res = impl_.clone_for_update();
+    let assoc_item_list = res.get_or_create_assoc_item_list();
+    let mut first_item = None;
+    for item in items {
+        if first_item.is_none() {
+            first_item = Some(item.clone())
+        }
+        assoc_item_list.add_item(item)
+    }
+    return (res, first_item.unwrap());
 
     fn add_body(fn_def: ast::Fn) -> ast::Fn {
         match fn_def.body() {
diff --git a/crates/syntax/src/ast/edit.rs b/crates/syntax/src/ast/edit.rs
index 10ec94cd2..7e4b8252e 100644
--- a/crates/syntax/src/ast/edit.rs
+++ b/crates/syntax/src/ast/edit.rs
@@ -80,81 +80,6 @@ where
     }
 }
 
-impl ast::Impl {
-    #[must_use]
-    pub fn with_assoc_item_list(&self, items: ast::AssocItemList) -> ast::Impl {
-        let mut to_insert: ArrayVec<SyntaxElement, 2> = ArrayVec::new();
-        if let Some(old_items) = self.assoc_item_list() {
-            let to_replace: SyntaxElement = old_items.syntax().clone().into();
-            to_insert.push(items.syntax().clone().into());
-            self.replace_children(single_node(to_replace), to_insert)
-        } else {
-            to_insert.push(make::tokens::single_space().into());
-            to_insert.push(items.syntax().clone().into());
-            self.insert_children(InsertPosition::Last, to_insert)
-        }
-    }
-}
-
-impl ast::AssocItemList {
-    #[must_use]
-    pub fn append_items(
-        &self,
-        items: impl IntoIterator<Item = ast::AssocItem>,
-    ) -> ast::AssocItemList {
-        let mut res = self.clone();
-        if !self.syntax().text().contains_char('\n') {
-            res = make_multiline(res);
-        }
-        items.into_iter().for_each(|it| res = res.append_item(it));
-        res.fixup_trailing_whitespace().unwrap_or(res)
-    }
-
-    #[must_use]
-    pub fn append_item(&self, item: ast::AssocItem) -> ast::AssocItemList {
-        let (indent, position, whitespace) = match self.assoc_items().last() {
-            Some(it) => (
-                leading_indent(it.syntax()).unwrap_or_default().to_string(),
-                InsertPosition::After(it.syntax().clone().into()),
-                "\n\n",
-            ),
-            None => match self.l_curly_token() {
-                Some(it) => (
-                    "    ".to_string() + &leading_indent(self.syntax()).unwrap_or_default(),
-                    InsertPosition::After(it.into()),
-                    "\n",
-                ),
-                None => return self.clone(),
-            },
-        };
-        let ws = tokens::WsBuilder::new(&format!("{}{}", whitespace, indent));
-        let to_insert: ArrayVec<SyntaxElement, 2> =
-            [ws.ws().into(), item.syntax().clone().into()].into();
-        self.insert_children(position, to_insert)
-    }
-
-    /// Remove extra whitespace between last item and closing curly brace.
-    fn fixup_trailing_whitespace(&self) -> Option<ast::AssocItemList> {
-        let first_token_after_items =
-            self.assoc_items().last()?.syntax().next_sibling_or_token()?;
-        let last_token_before_curly = self.r_curly_token()?.prev_sibling_or_token()?;
-        if last_token_before_curly != first_token_after_items {
-            // there is something more between last item and
-            // right curly than just whitespace - bail out
-            return None;
-        }
-        let whitespace =
-            last_token_before_curly.clone().into_token().and_then(ast::Whitespace::cast)?;
-        let text = whitespace.syntax().text();
-        let newline = text.rfind('\n')?;
-        let keep = tokens::WsBuilder::new(&text[newline..]);
-        Some(self.replace_children(
-            first_token_after_items..=last_token_before_curly,
-            std::iter::once(keep.ws().into()),
-        ))
-    }
-}
-
 impl ast::RecordExprFieldList {
     #[must_use]
     pub fn append_field(&self, field: &ast::RecordExprField) -> ast::RecordExprFieldList {
@@ -246,21 +171,6 @@ impl ast::TypeAlias {
     }
 }
 
-impl ast::TypeParam {
-    #[must_use]
-    pub fn remove_bounds(&self) -> ast::TypeParam {
-        let colon = match self.colon_token() {
-            Some(it) => it,
-            None => return self.clone(),
-        };
-        let end = match self.type_bound_list() {
-            Some(it) => it.syntax().clone().into(),
-            None => colon.clone().into(),
-        };
-        self.replace_children(colon.into()..=end, iter::empty())
-    }
-}
-
 impl ast::Path {
     #[must_use]
     pub fn with_segment(&self, segment: ast::PathSegment) -> ast::Path {
@@ -411,61 +321,6 @@ impl ast::MatchArmList {
     }
 }
 
-impl ast::GenericParamList {
-    #[must_use]
-    pub fn append_params(
-        &self,
-        params: impl IntoIterator<Item = ast::GenericParam>,
-    ) -> ast::GenericParamList {
-        let mut res = self.clone();
-        params.into_iter().for_each(|it| res = res.append_param(it));
-        res
-    }
-
-    #[must_use]
-    pub fn append_param(&self, item: ast::GenericParam) -> ast::GenericParamList {
-        let space = tokens::single_space();
-
-        let mut to_insert: ArrayVec<SyntaxElement, 4> = ArrayVec::new();
-        if self.generic_params().next().is_some() {
-            to_insert.push(space.into());
-        }
-        to_insert.push(item.syntax().clone().into());
-
-        macro_rules! after_l_angle {
-            () => {{
-                let anchor = match self.l_angle_token() {
-                    Some(it) => it.into(),
-                    None => return self.clone(),
-                };
-                InsertPosition::After(anchor)
-            }};
-        }
-
-        macro_rules! after_field {
-            ($anchor:expr) => {
-                if let Some(comma) = $anchor
-                    .syntax()
-                    .siblings_with_tokens(Direction::Next)
-                    .find(|it| it.kind() == T![,])
-                {
-                    InsertPosition::After(comma)
-                } else {
-                    to_insert.insert(0, make::token(T![,]).into());
-                    InsertPosition::After($anchor.syntax().clone().into())
-                }
-            };
-        }
-
-        let position = match self.generic_params().last() {
-            Some(it) => after_field!(it),
-            None => after_l_angle!(),
-        };
-
-        self.insert_children(position, to_insert)
-    }
-}
-
 #[must_use]
 pub fn remove_attrs_and_docs<N: ast::AttrsOwner>(node: &N) -> N {
     N::cast(remove_attrs_and_docs_inner(node.syntax().clone())).unwrap()
@@ -516,6 +371,12 @@ impl ops::Add<u8> for IndentLevel {
 }
 
 impl IndentLevel {
+    pub fn single() -> IndentLevel {
+        IndentLevel(0)
+    }
+    pub fn is_zero(&self) -> bool {
+        self.0 == 0
+    }
     pub fn from_element(element: &SyntaxElement) -> IndentLevel {
         match element {
             rowan::NodeOrToken::Node(it) => IndentLevel::from_node(it),
diff --git a/crates/syntax/src/ast/edit_in_place.rs b/crates/syntax/src/ast/edit_in_place.rs
index 168355555..9812e00c9 100644
--- a/crates/syntax/src/ast/edit_in_place.rs
+++ b/crates/syntax/src/ast/edit_in_place.rs
@@ -2,11 +2,16 @@
 
 use std::iter::empty;
 
-use parser::T;
+use parser::{SyntaxKind, T};
+use rowan::SyntaxElement;
 
 use crate::{
     algo::neighbor,
-    ast::{self, edit::AstNodeEdit, make, GenericParamsOwner, WhereClause},
+    ast::{
+        self,
+        edit::{AstNodeEdit, IndentLevel},
+        make, GenericParamsOwner,
+    },
     ted::{self, Position},
     AstNode, AstToken, Direction,
 };
@@ -37,7 +42,7 @@ impl GenericParamsOwnerEdit for ast::Fn {
         }
     }
 
-    fn get_or_create_where_clause(&self) -> WhereClause {
+    fn get_or_create_where_clause(&self) -> ast::WhereClause {
         if self.where_clause().is_none() {
             let position = if let Some(ty) = self.ret_type() {
                 Position::after(ty.syntax())
@@ -67,7 +72,7 @@ impl GenericParamsOwnerEdit for ast::Impl {
         }
     }
 
-    fn get_or_create_where_clause(&self) -> WhereClause {
+    fn get_or_create_where_clause(&self) -> ast::WhereClause {
         if self.where_clause().is_none() {
             let position = if let Some(items) = self.assoc_item_list() {
                 Position::before(items.syntax())
@@ -97,7 +102,7 @@ impl GenericParamsOwnerEdit for ast::Trait {
         }
     }
 
-    fn get_or_create_where_clause(&self) -> WhereClause {
+    fn get_or_create_where_clause(&self) -> ast::WhereClause {
         if self.where_clause().is_none() {
             let position = if let Some(items) = self.assoc_item_list() {
                 Position::before(items.syntax())
@@ -127,7 +132,7 @@ impl GenericParamsOwnerEdit for ast::Struct {
         }
     }
 
-    fn get_or_create_where_clause(&self) -> WhereClause {
+    fn get_or_create_where_clause(&self) -> ast::WhereClause {
         if self.where_clause().is_none() {
             let tfl = self.field_list().and_then(|fl| match fl {
                 ast::FieldList::RecordFieldList(_) => None,
@@ -165,7 +170,7 @@ impl GenericParamsOwnerEdit for ast::Enum {
         }
     }
 
-    fn get_or_create_where_clause(&self) -> WhereClause {
+    fn get_or_create_where_clause(&self) -> ast::WhereClause {
         if self.where_clause().is_none() {
             let position = if let Some(gpl) = self.generic_param_list() {
                 Position::after(gpl.syntax())
@@ -272,6 +277,59 @@ impl ast::Use {
     }
 }
 
+impl ast::Impl {
+    pub fn get_or_create_assoc_item_list(&self) -> ast::AssocItemList {
+        if self.assoc_item_list().is_none() {
+            let assoc_item_list = make::assoc_item_list().clone_for_update();
+            ted::append_child(self.syntax(), assoc_item_list.syntax());
+        }
+        self.assoc_item_list().unwrap()
+    }
+}
+
+impl ast::AssocItemList {
+    pub fn add_item(&self, item: ast::AssocItem) {
+        let (indent, position, whitespace) = match self.assoc_items().last() {
+            Some(last_item) => (
+                IndentLevel::from_node(last_item.syntax()),
+                Position::after(last_item.syntax()),
+                "\n\n",
+            ),
+            None => match self.l_curly_token() {
+                Some(l_curly) => {
+                    self.normalize_ws_between_braces();
+                    (IndentLevel::from_token(&l_curly) + 1, Position::after(&l_curly), "\n")
+                }
+                None => (IndentLevel::single(), Position::last_child_of(self.syntax()), "\n"),
+            },
+        };
+        let elements: Vec<SyntaxElement<_>> = vec![
+            make::tokens::whitespace(&format!("{}{}", whitespace, indent)).into(),
+            item.syntax().clone().into(),
+        ];
+        ted::insert_all(position, elements);
+    }
+
+    fn normalize_ws_between_braces(&self) -> Option<()> {
+        let l = self.l_curly_token()?;
+        let r = self.r_curly_token()?;
+        let indent = IndentLevel::from_node(self.syntax());
+
+        match l.next_sibling_or_token() {
+            Some(ws) if ws.kind() == SyntaxKind::WHITESPACE => {
+                if ws.next_sibling_or_token()?.into_token()? == r {
+                    ted::replace(ws, make::tokens::whitespace(&format!("\n{}", indent)));
+                }
+            }
+            Some(ws) if ws.kind() == T!['}'] => {
+                ted::insert(Position::after(l), make::tokens::whitespace(&format!("\n{}", indent)));
+            }
+            _ => (),
+        }
+        Some(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use std::fmt;
diff --git a/crates/syntax/src/ast/make.rs b/crates/syntax/src/ast/make.rs
index de04c8620..d13926ded 100644
--- a/crates/syntax/src/ast/make.rs
+++ b/crates/syntax/src/ast/make.rs
@@ -99,7 +99,7 @@ fn ty_from_text(text: &str) -> ast::Type {
 }
 
 pub fn assoc_item_list() -> ast::AssocItemList {
-    ast_from_text("impl C for D {};")
+    ast_from_text("impl C for D {}")
 }
 
 pub fn impl_trait(trait_: ast::Path, ty: ast::Path) -> ast::Impl {