Merge #3640

3640: Merge imports assist r=matklad a=matklad

Work towards #2220



bors r+
🤖

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

8 files changed, 274 insertions, 72 deletions

diff --git a/crates/ra_assists/src/doc_tests/generated.rs b/crates/ra_assists/src/doc_tests/generated.rs
index 3f56dd508..aef6793e8 100644
--- a/crates/ra_assists/src/doc_tests/generated.rs
+++ b/crates/ra_assists/src/doc_tests/generated.rs
@@ -418,6 +418,20 @@ fn main() {
 }
 
 #[test]
+fn doctest_merge_imports() {
+    check(
+        "merge_imports",
+        r#####"
+use std::<|>fmt::Formatter;
+use std::io;
+"#####,
+        r#####"
+use std::{fmt::Formatter, io};
+"#####,
+    )
+}
+
+#[test]
 fn doctest_merge_match_arms() {
     check(
         "merge_match_arms",
diff --git a/crates/ra_assists/src/handlers/merge_imports.rs b/crates/ra_assists/src/handlers/merge_imports.rs
new file mode 100644
index 000000000..96b1ab86a
--- /dev/null
+++ b/crates/ra_assists/src/handlers/merge_imports.rs
@@ -0,0 +1,154 @@
+use std::iter::successors;
+
+use ast::{edit::AstNodeEdit, make};
+use ra_syntax::{ast, AstNode, AstToken, Direction, InsertPosition, SyntaxElement, T};
+
+use crate::{Assist, AssistCtx, AssistId};
+
+// Assist: merge_imports
+//
+// Merges two imports with a common prefix.
+//
+// ```
+// use std::<|>fmt::Formatter;
+// use std::io;
+// ```
+// ->
+// ```
+// use std::{fmt::Formatter, io};
+// ```
+pub(crate) fn merge_imports(ctx: AssistCtx) -> Option<Assist> {
+    let tree: ast::UseTree = ctx.find_node_at_offset()?;
+    let use_item = tree.syntax().parent().and_then(ast::UseItem::cast)?;
+    let (merged, to_delete) = [Direction::Prev, Direction::Next]
+        .iter()
+        .copied()
+        .filter_map(|dir| next_use_item(&use_item, dir))
+        .filter_map(|it| Some((it.clone(), it.use_tree()?)))
+        .find_map(|(use_item, use_tree)| {
+            Some((try_merge_trees(&tree, &use_tree)?, use_item.clone()))
+        })?;
+    let mut offset = ctx.frange.range.start();
+    ctx.add_assist(AssistId("merge_imports"), "Merge imports", |edit| {
+        edit.replace_ast(tree, merged);
+
+        let mut range = to_delete.syntax().text_range();
+        let next_ws = to_delete
+            .syntax()
+            .next_sibling_or_token()
+            .and_then(|it| it.into_token())
+            .and_then(ast::Whitespace::cast);
+        if let Some(ws) = next_ws {
+            range = range.extend_to(&ws.syntax().text_range())
+        }
+        edit.delete(range);
+        if range.end() <= offset {
+            offset -= range.len();
+        }
+        edit.set_cursor(offset);
+    })
+}
+
+fn next_use_item(this_use_item: &ast::UseItem, direction: Direction) -> Option<ast::UseItem> {
+    this_use_item.syntax().siblings(direction).skip(1).find_map(ast::UseItem::cast)
+}
+
+fn try_merge_trees(old: &ast::UseTree, new: &ast::UseTree) -> Option<ast::UseTree> {
+    let lhs_path = old.path()?;
+    let rhs_path = new.path()?;
+
+    let (lhs_prefix, rhs_prefix) = common_prefix(&lhs_path, &rhs_path)?;
+
+    let lhs = old.split_prefix(&lhs_prefix);
+    let rhs = new.split_prefix(&rhs_prefix);
+
+    let mut to_insert: Vec<SyntaxElement> = Vec::new();
+    to_insert.push(make::token(T![,]).into());
+    to_insert.push(make::tokens::single_space().into());
+    to_insert.extend(
+        rhs.use_tree_list()?
+            .syntax()
+            .children_with_tokens()
+            .filter(|it| it.kind() != T!['{'] && it.kind() != T!['}']),
+    );
+    let use_tree_list = lhs.use_tree_list()?;
+    let pos = InsertPosition::Before(use_tree_list.r_curly()?.into());
+    let use_tree_list = use_tree_list.insert_children(pos, to_insert);
+    Some(lhs.with_use_tree_list(use_tree_list))
+}
+
+fn common_prefix(lhs: &ast::Path, rhs: &ast::Path) -> Option<(ast::Path, ast::Path)> {
+    let mut res = None;
+    let mut lhs_curr = first_path(&lhs);
+    let mut rhs_curr = first_path(&rhs);
+    loop {
+        match (lhs_curr.segment(), rhs_curr.segment()) {
+            (Some(lhs), Some(rhs)) if lhs.syntax().text() == rhs.syntax().text() => (),
+            _ => break,
+        }
+        res = Some((lhs_curr.clone(), rhs_curr.clone()));
+
+        match (lhs_curr.parent_path(), rhs_curr.parent_path()) {
+            (Some(lhs), Some(rhs)) => {
+                lhs_curr = lhs;
+                rhs_curr = rhs;
+            }
+            _ => break,
+        }
+    }
+
+    res
+}
+
+fn first_path(path: &ast::Path) -> ast::Path {
+    successors(Some(path.clone()), |it| it.qualifier()).last().unwrap()
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::helpers::check_assist;
+
+    use super::*;
+
+    #[test]
+    fn test_merge_first() {
+        check_assist(
+            merge_imports,
+            r"
+use std::fmt<|>::Debug;
+use std::fmt::Display;
+",
+            r"
+use std::fmt<|>::{Debug, Display};
+",
+        )
+    }
+
+    #[test]
+    fn test_merge_second() {
+        check_assist(
+            merge_imports,
+            r"
+use std::fmt::Debug;
+use std::fmt<|>::Display;
+",
+            r"
+use std::fmt<|>::{Display, Debug};
+",
+        )
+    }
+
+    #[test]
+    #[ignore]
+    fn test_merge_nested() {
+        check_assist(
+            merge_imports,
+            r"
+use std::{fmt<|>::Debug, fmt::Display};
+",
+            r"
+use std::{fmt::{Debug, Display}};
+",
+        )
+    }
+}
diff --git a/crates/ra_assists/src/handlers/move_bounds.rs b/crates/ra_assists/src/handlers/move_bounds.rs
index 0b501f3e5..342a770ec 100644
--- a/crates/ra_assists/src/handlers/move_bounds.rs
+++ b/crates/ra_assists/src/handlers/move_bounds.rs
@@ -1,5 +1,5 @@
 use ra_syntax::{
-    ast::{self, edit, make, AstNode, NameOwner, TypeBoundsOwner},
+    ast::{self, edit::AstNodeEdit, make, AstNode, NameOwner, TypeBoundsOwner},
     SyntaxElement,
     SyntaxKind::*,
 };
@@ -54,7 +54,7 @@ pub(crate) fn move_bounds_to_where_clause(ctx: AssistCtx) -> Option<Assist> {
                 (type_param, without_bounds)
             });
 
-        let new_type_param_list = edit::replace_descendants(&type_param_list, new_params);
+        let new_type_param_list = type_param_list.replace_descendants(new_params);
         edit.replace_ast(type_param_list.clone(), new_type_param_list);
 
         let where_clause = {
diff --git a/crates/ra_assists/src/handlers/split_import.rs b/crates/ra_assists/src/handlers/split_import.rs
index 292c39f59..d9244f22d 100644
--- a/crates/ra_assists/src/handlers/split_import.rs
+++ b/crates/ra_assists/src/handlers/split_import.rs
@@ -1,9 +1,6 @@
-use std::iter::{once, successors};
+use std::iter::successors;
 
-use ra_syntax::{
-    ast::{self, make},
-    AstNode, T,
-};
+use ra_syntax::{ast, AstNode, T};
 
 use crate::{Assist, AssistCtx, AssistId};
 
@@ -25,7 +22,10 @@ pub(crate) fn split_import(ctx: AssistCtx) -> Option<Assist> {
 
     let use_tree = top_path.syntax().ancestors().find_map(ast::UseTree::cast)?;
 
-    let new_tree = split_use_tree_prefix(&use_tree, &path)?;
+    let new_tree = use_tree.split_prefix(&path);
+    if new_tree == use_tree {
+        return None;
+    }
     let cursor = ctx.frange.range.start();
 
     ctx.add_assist(AssistId("split_import"), "Split import", |edit| {
@@ -35,23 +35,6 @@ pub(crate) fn split_import(ctx: AssistCtx) -> Option<Assist> {
     })
 }
 
-fn split_use_tree_prefix(use_tree: &ast::UseTree, prefix: &ast::Path) -> Option<ast::UseTree> {
-    let suffix = split_path_prefix(&prefix)?;
-    let use_tree = make::use_tree(suffix.clone(), use_tree.use_tree_list(), use_tree.alias());
-    let nested = make::use_tree_list(once(use_tree));
-    let res = make::use_tree(prefix.clone(), Some(nested), None);
-    Some(res)
-}
-
-fn split_path_prefix(prefix: &ast::Path) -> Option<ast::Path> {
-    let parent = prefix.parent_path()?;
-    let mut res = make::path_unqualified(parent.segment()?);
-    for p in successors(parent.parent_path(), |it| it.parent_path()) {
-        res = make::path_qualified(res, p.segment()?);
-    }
-    Some(res)
-}
-
 #[cfg(test)]
 mod tests {
     use crate::helpers::{check_assist, check_assist_target};
diff --git a/crates/ra_assists/src/lib.rs b/crates/ra_assists/src/lib.rs
index 260ada716..b8704ea7d 100644
--- a/crates/ra_assists/src/lib.rs
+++ b/crates/ra_assists/src/lib.rs
@@ -110,6 +110,7 @@ mod handlers {
     mod inline_local_variable;
     mod introduce_variable;
     mod invert_if;
+    mod merge_imports;
     mod merge_match_arms;
     mod move_bounds;
     mod move_guard;
@@ -140,6 +141,7 @@ mod handlers {
             inline_local_variable::inline_local_variable,
             introduce_variable::introduce_variable,
             invert_if::invert_if,
+            merge_imports::merge_imports,
             merge_match_arms::merge_match_arms,
             move_bounds::move_bounds_to_where_clause,
             move_guard::move_arm_cond_to_match_guard,
diff --git a/crates/ra_syntax/src/ast/edit.rs b/crates/ra_syntax/src/ast/edit.rs
index 40a04b9c5..68dae008f 100644
--- a/crates/ra_syntax/src/ast/edit.rs
+++ b/crates/ra_syntax/src/ast/edit.rs
@@ -23,7 +23,7 @@ impl ast::BinExpr {
     pub fn replace_op(&self, op: SyntaxKind) -> Option<ast::BinExpr> {
         let op_node: SyntaxElement = self.op_details()?.0.into();
         let to_insert: Option<SyntaxElement> = Some(make::token(op).into());
-        Some(replace_children(self, single_node(op_node), to_insert))
+        Some(self.replace_children(single_node(op_node), to_insert))
     }
 }
 
@@ -39,10 +39,10 @@ impl ast::FnDef {
         } else {
             to_insert.push(make::tokens::single_space().into());
             to_insert.push(body.syntax().clone().into());
-            return insert_children(self, InsertPosition::Last, to_insert);
+            return self.insert_children(InsertPosition::Last, to_insert);
         };
         to_insert.push(body.syntax().clone().into());
-        replace_children(self, single_node(old_body_or_semi), to_insert)
+        self.replace_children(single_node(old_body_or_semi), to_insert)
     }
 }
 
@@ -75,7 +75,7 @@ impl ast::ItemList {
         let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
         let to_insert: ArrayVec<[SyntaxElement; 2]> =
             [ws.ws().into(), item.syntax().clone().into()].into();
-        insert_children(self, position, to_insert)
+        self.insert_children(position, to_insert)
     }
 
     fn l_curly(&self) -> Option<SyntaxElement> {
@@ -106,8 +106,8 @@ impl ast::ItemList {
         let ws = tokens::WsBuilder::new(&format!("\n{}", indent));
         let to_insert = iter::once(ws.ws().into());
         match existing_ws {
-            None => insert_children(self, InsertPosition::After(l_curly), to_insert),
-            Some(ws) => replace_children(self, single_node(ws), to_insert),
+            None => self.insert_children(InsertPosition::After(l_curly), to_insert),
+            Some(ws) => self.replace_children(single_node(ws), to_insert),
         }
     }
 }
@@ -184,7 +184,7 @@ impl ast::RecordFieldList {
             InsertPosition::After(anchor) => after_field!(anchor),
         };
 
-        insert_children(self, position, to_insert)
+        self.insert_children(position, to_insert)
     }
 
     fn l_curly(&self) -> Option<SyntaxElement> {
@@ -203,7 +203,7 @@ impl ast::TypeParam {
             Some(it) => it.syntax().clone().into(),
             None => colon.clone().into(),
         };
-        replace_children(self, colon.into()..=end, iter::empty())
+        self.replace_children(colon.into()..=end, iter::empty())
     }
 }
 
@@ -211,8 +211,7 @@ impl ast::Path {
     #[must_use]
     pub fn with_segment(&self, segment: ast::PathSegment) -> ast::Path {
         if let Some(old) = self.segment() {
-            return replace_children(
-                self,
+            return self.replace_children(
                 single_node(old.syntax().clone()),
                 iter::once(segment.syntax().clone().into()),
             );
@@ -234,8 +233,7 @@ impl ast::PathSegment {
 
     fn _with_type_args(&self, type_args: ast::TypeArgList, turbo: bool) -> ast::PathSegment {
         if let Some(old) = self.type_arg_list() {
-            return replace_children(
-                self,
+            return self.replace_children(
                 single_node(old.syntax().clone()),
                 iter::once(type_args.syntax().clone().into()),
             );
@@ -245,7 +243,7 @@ impl ast::PathSegment {
             to_insert.push(make::token(T![::]).into());
         }
         to_insert.push(type_args.syntax().clone().into());
-        insert_children(self, InsertPosition::Last, to_insert)
+        self.insert_children(InsertPosition::Last, to_insert)
     }
 }
 
@@ -253,7 +251,7 @@ impl ast::UseItem {
     #[must_use]
     pub fn with_use_tree(&self, use_tree: ast::UseTree) -> ast::UseItem {
         if let Some(old) = self.use_tree() {
-            return replace_descendants(self, iter::once((old, use_tree)));
+            return self.replace_descendants(iter::once((old, use_tree)));
         }
         self.clone()
     }
@@ -263,7 +261,7 @@ impl ast::UseTree {
     #[must_use]
     pub fn with_path(&self, path: ast::Path) -> ast::UseTree {
         if let Some(old) = self.path() {
-            return replace_descendants(self, iter::once((old, path)));
+            return self.replace_descendants(iter::once((old, path)));
         }
         self.clone()
     }
@@ -271,10 +269,30 @@ impl ast::UseTree {
     #[must_use]
     pub fn with_use_tree_list(&self, use_tree_list: ast::UseTreeList) -> ast::UseTree {
         if let Some(old) = self.use_tree_list() {
-            return replace_descendants(self, iter::once((old, use_tree_list)));
+            return self.replace_descendants(iter::once((old, use_tree_list)));
         }
         self.clone()
     }
+
+    #[must_use]
+    pub fn split_prefix(&self, prefix: &ast::Path) -> ast::UseTree {
+        let suffix = match split_path_prefix(&prefix) {
+            Some(it) => it,
+            None => return self.clone(),
+        };
+        let use_tree = make::use_tree(suffix.clone(), self.use_tree_list(), self.alias());
+        let nested = make::use_tree_list(iter::once(use_tree));
+        return make::use_tree(prefix.clone(), Some(nested), None);
+
+        fn split_path_prefix(prefix: &ast::Path) -> Option<ast::Path> {
+            let parent = prefix.parent_path()?;
+            let mut res = make::path_unqualified(parent.segment()?);
+            for p in iter::successors(parent.parent_path(), |it| it.parent_path()) {
+                res = make::path_qualified(res, p.segment()?);
+            }
+            Some(res)
+        }
+    }
 }
 
 #[must_use]
@@ -295,19 +313,6 @@ fn strip_attrs_and_docs_inner(mut node: SyntaxNode) -> SyntaxNode {
     node
 }
 
-#[must_use]
-pub fn replace_descendants<N: AstNode, D: AstNode>(
-    parent: &N,
-    replacement_map: impl IntoIterator<Item = (D, D)>,
-) -> N {
-    let map = replacement_map
-        .into_iter()
-        .map(|(from, to)| (from.syntax().clone().into(), to.syntax().clone().into()))
-        .collect::<FxHashMap<SyntaxElement, _>>();
-    let new_syntax = algo::replace_descendants(parent.syntax(), |n| map.get(n).cloned());
-    N::cast(new_syntax).unwrap()
-}
-
 #[derive(Debug, Clone, Copy)]
 pub struct IndentLevel(pub u8);
 
@@ -411,31 +416,48 @@ fn prev_tokens(token: SyntaxToken) -> impl Iterator<Item = SyntaxToken> {
     iter::successors(Some(token), |token| token.prev_token())
 }
 
-#[must_use]
-fn insert_children<N: AstNode>(
-    parent: &N,
-    position: InsertPosition<SyntaxElement>,
-    to_insert: impl IntoIterator<Item = SyntaxElement>,
-) -> N {
-    let new_syntax = algo::insert_children(parent.syntax(), position, to_insert);
-    N::cast(new_syntax).unwrap()
+pub trait AstNodeEdit: AstNode + Sized {
+    #[must_use]
+    fn insert_children(
+        &self,
+        position: InsertPosition<SyntaxElement>,
+        to_insert: impl IntoIterator<Item = SyntaxElement>,
+    ) -> Self {
+        let new_syntax = algo::insert_children(self.syntax(), position, to_insert);
+        Self::cast(new_syntax).unwrap()
+    }
+
+    #[must_use]
+    fn replace_children(
+        &self,
+        to_replace: RangeInclusive<SyntaxElement>,
+        to_insert: impl IntoIterator<Item = SyntaxElement>,
+    ) -> Self {
+        let new_syntax = algo::replace_children(self.syntax(), to_replace, to_insert);
+        Self::cast(new_syntax).unwrap()
+    }
+
+    #[must_use]
+    fn replace_descendants<D: AstNode>(
+        &self,
+        replacement_map: impl IntoIterator<Item = (D, D)>,
+    ) -> Self {
+        let map = replacement_map
+            .into_iter()
+            .map(|(from, to)| (from.syntax().clone().into(), to.syntax().clone().into()))
+            .collect::<FxHashMap<SyntaxElement, _>>();
+        let new_syntax = algo::replace_descendants(self.syntax(), |n| map.get(n).cloned());
+        Self::cast(new_syntax).unwrap()
+    }
 }
 
+impl<N: AstNode> AstNodeEdit for N {}
+
 fn single_node(element: impl Into<SyntaxElement>) -> RangeInclusive<SyntaxElement> {
     let element = element.into();
     element.clone()..=element
 }
 
-#[must_use]
-fn replace_children<N: AstNode>(
-    parent: &N,
-    to_replace: RangeInclusive<SyntaxElement>,
-    to_insert: impl IntoIterator<Item = SyntaxElement>,
-) -> N {
-    let new_syntax = algo::replace_children(parent.syntax(), to_replace, to_insert);
-    N::cast(new_syntax).unwrap()
-}
-
 #[test]
 fn test_increase_indent() {
     let arm_list = {
diff --git a/crates/ra_syntax/src/ast/extensions.rs b/crates/ra_syntax/src/ast/extensions.rs
index d5986e8b4..c3ae8f90e 100644
--- a/crates/ra_syntax/src/ast/extensions.rs
+++ b/crates/ra_syntax/src/ast/extensions.rs
@@ -167,6 +167,20 @@ impl ast::UseTreeList {
             .and_then(ast::UseTree::cast)
             .expect("UseTreeLists are always nested in UseTrees")
     }
+    pub fn l_curly(&self) -> Option<SyntaxToken> {
+        self.token(T!['{'])
+    }
+
+    pub fn r_curly(&self) -> Option<SyntaxToken> {
+        self.token(T!['}'])
+    }
+
+    fn token(&self, kind: SyntaxKind) -> Option<SyntaxToken> {
+        self.syntax()
+            .children_with_tokens()
+            .filter_map(|it| it.into_token())
+            .find(|it| it.kind() == kind)
+    }
 }
 
 impl ast::ImplDef {
diff --git a/docs/user/assists.md b/docs/user/assists.md
index 9861332af..e2850b4dd 100644
--- a/docs/user/assists.md
+++ b/docs/user/assists.md
@@ -402,6 +402,19 @@ fn main() {
 }
 ```
 
+## `merge_imports`
+
+Merges two imports with a common prefix.
+
+```rust
+// BEFORE
+use std::┃fmt::Formatter;
+use std::io;
+
+// AFTER
+use std::{fmt::Formatter, io};
+```
+
 ## `merge_match_arms`
 
 Merges identical match arms.