Auto merge of rust-lang#12452 - iDawer:assist.merge_selected_imports, r=Veykril

feature: `Merge imports` assist can merge multiple selected imports.

The selected imports have to have a common prefix in paths.

Select imports or use trees to merge:
```rust
$0use std::fmt::Display;
use std::fmt::Debug;
use std::fmt::Write;$0
```
Apply `Merge imports`:
```rust
use std::fmt::{Display, Debug, Write};
```

Closes rust-lang#12426

Author: bors

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

@@ -1,12 +1,15 @@
+use either::Either;
 use ide_db::imports::merge_imports::{try_merge_imports, try_merge_trees, MergeBehavior};
-use syntax::{algo::neighbor, ast, ted, AstNode};
+use syntax::{algo::neighbor, ast, match_ast, ted, AstNode, SyntaxElement, SyntaxNode};
 
 use crate::{
     assist_context::{AssistContext, Assists},
     utils::next_prev,
     AssistId, AssistKind,
 };
 
+use Edit::*;
+
 // Assist: merge_imports
 //
 // Merges two imports with a common prefix.
@@ -20,51 +23,115 @@ use crate::{
 // use std::{fmt::Formatter, io};
 // ```
 pub(crate) fn merge_imports(acc: &mut Assists, ctx: &AssistContext) -> Option<()> {
-    let tree: ast::UseTree = ctx.find_node_at_offset()?;
-
-    let mut imports = None;
-    let mut uses = None;
-    if let Some(use_item) = tree.syntax().parent().and_then(ast::Use::cast) {
-        let (merged, to_remove) =
-            next_prev().filter_map(|dir| neighbor(&use_item, dir)).find_map(|use_item2| {
-                try_merge_imports(&use_item, &use_item2, MergeBehavior::Crate).zip(Some(use_item2))
-            })?;
-
-        imports = Some((use_item, merged, to_remove));
+    let (target, edits) = if ctx.has_empty_selection() {
+        // Merge a neighbor
+        let tree: ast::UseTree = ctx.find_node_at_offset()?;
+        let target = tree.syntax().text_range();
+
+        let edits = if let Some(use_item) = tree.syntax().parent().and_then(ast::Use::cast) {
+            let mut neighbor = next_prev().find_map(|dir| neighbor(&use_item, dir)).into_iter();
+            use_item.try_merge_from(&mut neighbor)
+        } else {
+            let mut neighbor = next_prev().find_map(|dir| neighbor(&tree, dir)).into_iter();
+            tree.try_merge_from(&mut neighbor)
+        };
+        (target, edits?)
     } else {
-        let (merged, to_remove) =
-            next_prev().filter_map(|dir| neighbor(&tree, dir)).find_map(|use_tree| {
-                try_merge_trees(&tree, &use_tree, MergeBehavior::Crate).zip(Some(use_tree))
-            })?;
-
-        uses = Some((tree.clone(), merged, to_remove))
+        // Merge selected
+        let selection_range = ctx.selection_trimmed();
+        let parent_node = match ctx.covering_element() {
+            SyntaxElement::Node(n) => n,
+            SyntaxElement::Token(t) => t.parent()?,
+        };
+        let mut selected_nodes =
+            parent_node.children().filter(|it| selection_range.contains_range(it.text_range()));
+
+        let first_selected = selected_nodes.next()?;
+        let edits = match_ast! {
+            match first_selected {
+                ast::Use(use_item) => {
+                    use_item.try_merge_from(&mut selected_nodes.filter_map(ast::Use::cast))
+                },
+                ast::UseTree(use_tree) => {
+                    use_tree.try_merge_from(&mut selected_nodes.filter_map(ast::UseTree::cast))
+                },
+                _ => return None,
+            }
+        };
+        (selection_range, edits?)
     };
 
-    let target = tree.syntax().text_range();
     acc.add(
         AssistId("merge_imports", AssistKind::RefactorRewrite),
         "Merge imports",
         target,
         |builder| {
-            if let Some((to_replace, replacement, to_remove)) = imports {
-                let to_replace = builder.make_mut(to_replace);
-                let to_remove = builder.make_mut(to_remove);
-
-                ted::replace(to_replace.syntax(), replacement.syntax());
-                to_remove.remove();
-            }
-
-            if let Some((to_replace, replacement, to_remove)) = uses {
-                let to_replace = builder.make_mut(to_replace);
-                let to_remove = builder.make_mut(to_remove);
-
-                ted::replace(to_replace.syntax(), replacement.syntax());
-                to_remove.remove()
+            let edits_mut: Vec<Edit> = edits
+                .into_iter()
+                .map(|it| match it {
+                    Remove(Either::Left(it)) => Remove(Either::Left(builder.make_mut(it))),
+                    Remove(Either::Right(it)) => Remove(Either::Right(builder.make_mut(it))),
+                    Replace(old, new) => Replace(builder.make_syntax_mut(old), new),
+                })
+                .collect();
+            for edit in edits_mut {
+                match edit {
+                    Remove(it) => it.as_ref().either(ast::Use::remove, ast::UseTree::remove),
+                    Replace(old, new) => ted::replace(old, new),
+                }
             }
         },
     )
 }
 
+trait Merge: AstNode + Clone {
+    fn try_merge_from(self, items: &mut dyn Iterator<Item = Self>) -> Option<Vec<Edit>> {
+        let mut edits = Vec::new();
+        let mut merged = self.clone();
+        while let Some(item) = items.next() {
+            merged = merged.try_merge(&item)?;
+            edits.push(Edit::Remove(item.into_either()));
+        }
+        if !edits.is_empty() {
+            edits.push(Edit::replace(self, merged));
+            Some(edits)
+        } else {
+            None
+        }
+    }
+    fn try_merge(&self, other: &Self) -> Option<Self>;
+    fn into_either(self) -> Either<ast::Use, ast::UseTree>;
+}
+
+impl Merge for ast::Use {
+    fn try_merge(&self, other: &Self) -> Option<Self> {
+        try_merge_imports(self, other, MergeBehavior::Crate)
+    }
+    fn into_either(self) -> Either<ast::Use, ast::UseTree> {
+        Either::Left(self)
+    }
+}
+
+impl Merge for ast::UseTree {
+    fn try_merge(&self, other: &Self) -> Option<Self> {
+        try_merge_trees(self, other, MergeBehavior::Crate)
+    }
+    fn into_either(self) -> Either<ast::Use, ast::UseTree> {
+        Either::Right(self)
+    }
+}
+
+enum Edit {
+    Remove(Either<ast::Use, ast::UseTree>),
+    Replace(SyntaxNode, SyntaxNode),
+}
+
+impl Edit {
+    fn replace(old: impl AstNode, new: impl AstNode) -> Self {
+        Edit::Replace(old.syntax().clone(), new.syntax().clone())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use crate::tests::{check_assist, check_assist_not_applicable};
@@ -454,4 +521,50 @@ use foo::{*, bar::Baz};
 ",
         );
     }
+
+    #[test]
+    fn merge_selection_uses() {
+        check_assist(
+            merge_imports,
+            r"
+use std::fmt::Error;
+$0use std::fmt::Display;
+use std::fmt::Debug;
+use std::fmt::Write;
+$0use std::fmt::Result;
+",
+            r"
+use std::fmt::Error;
+use std::fmt::{Display, Debug, Write};
+use std::fmt::Result;
+",
+        );
+    }
+
+    #[test]
+    fn merge_selection_use_trees() {
+        check_assist(
+            merge_imports,
+            r"
+use std::{
+    fmt::Error,
+    $0fmt::Display,
+    fmt::Debug,
+    fmt::Write,$0
+    fmt::Result,
+};",
+            r"
+use std::{
+    fmt::Error,
+    fmt::{Display, Debug, Write},
+    fmt::Result,
+};",
+        );
+        // FIXME: Remove redundant braces. See also unnecessary-braces diagnostic.
+        check_assist(
+            merge_imports,
+            r"use std::$0{fmt::Display, fmt::Debug}$0;",
+            r"use std::{fmt::{Display, Debug}};",
+        );
+    }
 }

crates/ide-db/src/imports/merge_imports.rs

@@ -30,6 +30,8 @@ impl MergeBehavior {
     }
 }
 
+/// Merge `rhs` into `lhs` keeping both intact.
+/// Returned AST is mutable.
 pub fn try_merge_imports(
     lhs: &ast::Use,
     rhs: &ast::Use,
@@ -51,6 +53,8 @@ pub fn try_merge_imports(
     Some(lhs)
 }
 
+/// Merge `rhs` into `lhs` keeping both intact.
+/// Returned AST is mutable.
 pub fn try_merge_trees(
     lhs: &ast::UseTree,
     rhs: &ast::UseTree,