Merge nested use trees

1 files changed, 52 insertions, 20 deletions

diff --git a/crates/ra_assists/src/handlers/merge_imports.rs b/crates/ra_assists/src/handlers/merge_imports.rs
index e9c49b7d0..89bc975bd 100644
--- a/crates/ra_assists/src/handlers/merge_imports.rs
+++ b/crates/ra_assists/src/handlers/merge_imports.rs
@@ -3,7 +3,7 @@ use std::iter::successors;
 use ra_syntax::{
     algo::neighbor,
     ast::{self, edit::AstNodeEdit, make},
-    AstNode, AstToken, Direction, InsertPosition, SyntaxElement, T,
+    AstNode, AstToken, Direction, InsertPosition, SyntaxElement, TextRange, T,
 };
 
 use crate::{Assist, AssistCtx, AssistId};
@@ -22,18 +22,15 @@ use crate::{Assist, AssistCtx, AssistId};
 // ```
 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| neighbor(&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 (new_tree, to_delete) = if let Some(use_item) =
+        tree.syntax().parent().and_then(ast::UseItem::cast)
+    {
+        let (merged, to_delete) = next_prev()
+            .filter_map(|dir| neighbor(&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 range = to_delete.syntax().text_range();
         let next_ws = to_delete
@@ -44,14 +41,41 @@ pub(crate) fn merge_imports(ctx: AssistCtx) -> Option<Assist> {
         if let Some(ws) = next_ws {
             range = range.extend_to(&ws.syntax().text_range())
         }
-        edit.delete(range);
-        if range.end() <= offset {
-            offset -= range.len();
+        (merged, range)
+    } else {
+        let (merged, to_delete) = next_prev()
+            .filter_map(|dir| neighbor(&tree, dir))
+            .find_map(|use_tree| Some((try_merge_trees(&tree, &use_tree)?, use_tree.clone())))?;
+
+        let mut range = to_delete.syntax().text_range();
+        if let Some((dir, nb)) = next_prev().find_map(|dir| Some((dir, neighbor(&to_delete, dir)?)))
+        {
+            let nb_range = nb.syntax().text_range();
+            if dir == Direction::Prev {
+                range = TextRange::from_to(nb_range.end(), range.end());
+            } else {
+                range = TextRange::from_to(range.start(), nb_range.start());
+            }
+        }
+        (merged, range)
+    };
+
+    let mut offset = ctx.frange.range.start();
+    ctx.add_assist(AssistId("merge_imports"), "Merge imports", |edit| {
+        edit.replace_ast(tree, new_tree);
+        edit.delete(to_delete);
+
+        if to_delete.end() <= offset {
+            offset -= to_delete.len();
         }
         edit.set_cursor(offset);
     })
 }
 
+fn next_prev() -> impl Iterator<Item = Direction> {
+    [Direction::Next, Direction::Prev].iter().copied()
+}
+
 fn try_merge_trees(old: &ast::UseTree, new: &ast::UseTree) -> Option<ast::UseTree> {
     let lhs_path = old.path()?;
     let rhs_path = new.path()?;
@@ -134,11 +158,10 @@ use std::fmt<|>::Display;
             r"
 use std::fmt<|>::{Display, Debug};
 ",
-        )
+        );
     }
 
     #[test]
-    #[ignore]
     fn test_merge_nested() {
         check_assist(
             merge_imports,
@@ -146,8 +169,17 @@ use std::fmt<|>::{Display, Debug};
 use std::{fmt<|>::Debug, fmt::Display};
 ",
             r"
-use std::{fmt::{Debug, Display}};
+use std::{fmt<|>::{Debug, Display}};
 ",
-        )
+        );
+        check_assist(
+            merge_imports,
+            r"
+use std::{fmt::Debug, fmt<|>::Display};
+",
+            r"
+use std::{fmt<|>::{Display, Debug}};
+",
+        );
     }
 }