Merge #5989

5989: Rewrite import merging r=jonas-schievink a=Veykril

Rewrites how import merging is being handled. It is now a recursive function to properly handle merging of intermediate levels in the import trees. With this ordering the imports is also now possible tho it doesn't quite order it the same way as `rustfmt` does yet, namely it orders lowercase identifiers after uppercase identifiers as that is the standard character order that rust uses. This also fixes a few weird behaviors that were visible in some of the `replace_qualified_name_with_use.rs` tests.

This really took longer than I was hoping for, fighting with import trees is quite the exhausting task 😅 

Co-authored-by: Lukas Wirth <[email protected]>

3 files changed, 287 insertions, 103 deletions

diff --git a/crates/assists/src/handlers/merge_imports.rs b/crates/assists/src/handlers/merge_imports.rs
index 0bd679260..fe33cee53 100644
--- a/crates/assists/src/handlers/merge_imports.rs
+++ b/crates/assists/src/handlers/merge_imports.rs
@@ -95,7 +95,7 @@ use std::fmt::Debug;
 use std::fmt<|>::Display;
 ",
             r"
-use std::fmt::{Display, Debug};
+use std::fmt::{Debug, Display};
 ",
         );
     }
@@ -122,7 +122,7 @@ use std::fmt::{self, Display};
 use std::{fmt, <|>fmt::Display};
 ",
             r"
-use std::{fmt::{Display, self}};
+use std::{fmt::{self, Display}};
 ",
         );
     }
@@ -210,13 +210,17 @@ use std::{fmt<|>::Debug, fmt::Display};
 use std::{fmt::{Debug, Display}};
 ",
         );
+    }
+
+    #[test]
+    fn test_merge_nested2() {
         check_assist(
             merge_imports,
             r"
 use std::{fmt::Debug, fmt<|>::Display};
 ",
             r"
-use std::{fmt::{Display, Debug}};
+use std::{fmt::{Debug, Display}};
 ",
         );
     }
@@ -310,9 +314,7 @@ use foo::<|>{
 };
 ",
             r"
-use foo::{
-    FooBar,
-bar::baz};
+use foo::{FooBar, bar::baz};
 ",
         )
     }
diff --git a/crates/assists/src/handlers/replace_qualified_name_with_use.rs b/crates/assists/src/handlers/replace_qualified_name_with_use.rs
index e48407fcc..8ac907707 100644
--- a/crates/assists/src/handlers/replace_qualified_name_with_use.rs
+++ b/crates/assists/src/handlers/replace_qualified_name_with_use.rs
@@ -312,7 +312,7 @@ impl std::fmt<|> for Foo {
 }
     ",
             r"
-use std::fmt::{Debug, self};
+use std::fmt::{self, Debug};
 
 impl fmt for Foo {
 }
@@ -330,9 +330,8 @@ use std::fmt::{Debug, nested::{Display}};
 impl std::fmt::nested<|> for Foo {
 }
 ",
-            // FIXME(veykril): should be nested::{self, Display} here
             r"
-use std::fmt::{Debug, nested::{Display}, nested};
+use std::fmt::{Debug, nested::{self, Display}};
 
 impl nested for Foo {
 }
@@ -350,9 +349,8 @@ use std::fmt::{Debug, nested::{self, Display}};
 impl std::fmt::nested<|> for Foo {
 }
 ",
-            // FIXME(veykril): nested is duplicated now
             r"
-use std::fmt::{Debug, nested::{self, Display}, nested};
+use std::fmt::{Debug, nested::{self, Display}};
 
 impl nested for Foo {
 }
@@ -371,7 +369,7 @@ impl std::fmt::nested::Debug<|> for Foo {
 }
 ",
             r"
-use std::fmt::{Debug, nested::{Display}, nested::Debug};
+use std::fmt::{Debug, nested::{Debug, Display}};
 
 impl Debug for Foo {
 }
@@ -409,7 +407,7 @@ impl std::fmt::Display<|> for Foo {
 }
 ",
             r"
-use std::fmt::{nested::Debug, Display};
+use std::fmt::{Display, nested::Debug};
 
 impl Display for Foo {
 }
@@ -429,12 +427,8 @@ use crate::{
 
 fn foo() { crate::ty::lower<|>::trait_env() }
 ",
-            // FIXME(veykril): formatting broke here
             r"
-use crate::{
-    ty::{Substs, Ty},
-    AssocItem,
-ty::lower};
+use crate::{AssocItem, ty::{Substs, Ty, lower}};
 
 fn foo() { lower::trait_env() }
 ",
@@ -633,7 +627,7 @@ fn main() {
 }
     ",
             r"
-use std::fmt::{Display, self};
+use std::fmt::{self, Display};
 
 fn main() {
     fmt;
diff --git a/crates/assists/src/utils/insert_use.rs b/crates/assists/src/utils/insert_use.rs
index 6d110aaaf..09f4a2224 100644
--- a/crates/assists/src/utils/insert_use.rs
+++ b/crates/assists/src/utils/insert_use.rs
@@ -1,7 +1,9 @@
 //! Handle syntactic aspects of inserting a new `use`.
-use std::iter::{self, successors};
+use std::{
+    cmp::Ordering,
+    iter::{self, successors},
+};
 
-use algo::skip_trivia_token;
 use ast::{
     edit::{AstNodeEdit, IndentLevel},
     PathSegmentKind, VisibilityOwner,
@@ -9,9 +11,8 @@ use ast::{
 use syntax::{
     algo,
     ast::{self, make, AstNode},
-    Direction, InsertPosition, SyntaxElement, SyntaxNode, T,
+    InsertPosition, SyntaxElement, SyntaxNode,
 };
-use test_utils::mark;
 
 #[derive(Debug)]
 pub enum ImportScope {
@@ -119,7 +120,6 @@ pub(crate) fn insert_use(
         }
 
         if let ident_level @ 1..=usize::MAX = scope.indent_level().0 as usize {
-            // FIXME: this alone doesnt properly re-align all cases
             buf.push(make::tokens::whitespace(&" ".repeat(4 * ident_level)).into());
         }
         buf.push(use_item.syntax().clone().into());
@@ -149,66 +149,123 @@ fn eq_visibility(vis0: Option<ast::Visibility>, vis1: Option<ast::Visibility>) -
 }
 
 pub(crate) fn try_merge_imports(
-    old: &ast::Use,
-    new: &ast::Use,
+    lhs: &ast::Use,
+    rhs: &ast::Use,
     merge_behaviour: MergeBehaviour,
 ) -> Option<ast::Use> {
     // don't merge imports with different visibilities
-    if !eq_visibility(old.visibility(), new.visibility()) {
+    if !eq_visibility(lhs.visibility(), rhs.visibility()) {
         return None;
     }
-    let old_tree = old.use_tree()?;
-    let new_tree = new.use_tree()?;
-    let merged = try_merge_trees(&old_tree, &new_tree, merge_behaviour)?;
-    Some(old.with_use_tree(merged))
-}
-
-/// Simple function that checks if a UseTreeList is deeper than one level
-fn use_tree_list_is_nested(tl: &ast::UseTreeList) -> bool {
-    tl.use_trees().any(|use_tree| {
-        use_tree.use_tree_list().is_some() || use_tree.path().and_then(|p| p.qualifier()).is_some()
-    })
+    let lhs_tree = lhs.use_tree()?;
+    let rhs_tree = rhs.use_tree()?;
+    let merged = try_merge_trees(&lhs_tree, &rhs_tree, merge_behaviour)?;
+    Some(lhs.with_use_tree(merged))
 }
 
-// FIXME: currently this merely prepends the new tree into old, ideally it would insert the items in a sorted fashion
 pub(crate) fn try_merge_trees(
-    old: &ast::UseTree,
-    new: &ast::UseTree,
-    merge_behaviour: MergeBehaviour,
+    lhs: &ast::UseTree,
+    rhs: &ast::UseTree,
+    merge: MergeBehaviour,
 ) -> Option<ast::UseTree> {
-    let lhs_path = old.path()?;
-    let rhs_path = new.path()?;
+    let lhs_path = lhs.path()?;
+    let rhs_path = rhs.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 lhs_tl = lhs.use_tree_list()?;
-    let rhs_tl = rhs.use_tree_list()?;
-
-    // if we are only allowed to merge the last level check if the split off paths are only one level deep
-    if merge_behaviour == MergeBehaviour::Last
-        && (use_tree_list_is_nested(&lhs_tl) || use_tree_list_is_nested(&rhs_tl))
-    {
-        mark::hit!(test_last_merge_too_long);
-        return None;
-    }
+    let lhs = lhs.split_prefix(&lhs_prefix);
+    let rhs = rhs.split_prefix(&rhs_prefix);
+    recursive_merge(&lhs, &rhs, merge).map(|(merged, _)| merged)
+}
 
-    let should_insert_comma = lhs_tl
-        .r_curly_token()
-        .and_then(|it| skip_trivia_token(it.prev_token()?, Direction::Prev))
-        .map(|it| it.kind())
-        != Some(T![,]);
-    let mut to_insert: Vec<SyntaxElement> = Vec::new();
-    if should_insert_comma {
-        to_insert.push(make::token(T![,]).into());
-        to_insert.push(make::tokens::single_space().into());
-    }
-    to_insert.extend(
-        rhs_tl.syntax().children_with_tokens().filter(|it| !matches!(it.kind(), T!['{'] | T!['}'])),
-    );
-    let pos = InsertPosition::Before(lhs_tl.r_curly_token()?.into());
-    let use_tree_list = lhs_tl.insert_children(pos, to_insert);
-    Some(lhs.with_use_tree_list(use_tree_list))
+/// Recursively "zips" together lhs and rhs.
+fn recursive_merge(
+    lhs: &ast::UseTree,
+    rhs: &ast::UseTree,
+    merge: MergeBehaviour,
+) -> Option<(ast::UseTree, bool)> {
+    let mut use_trees = lhs
+        .use_tree_list()
+        .into_iter()
+        .flat_map(|list| list.use_trees())
+        // check if any of the use trees are nested, if they are and the behaviour is `last` we are not allowed to merge this
+        // so early exit the iterator by using Option's Intoiterator impl
+        .map(|tree| match merge == MergeBehaviour::Last && tree.use_tree_list().is_some() {
+            true => None,
+            false => Some(tree),
+        })
+        .collect::<Option<Vec<_>>>()?;
+    use_trees.sort_unstable_by(|a, b| path_cmp_opt(a.path(), b.path()));
+    for rhs_t in rhs.use_tree_list().into_iter().flat_map(|list| list.use_trees()) {
+        let rhs_path = rhs_t.path();
+        match use_trees.binary_search_by(|p| path_cmp_opt(p.path(), rhs_path.clone())) {
+            Ok(idx) => {
+                let lhs_t = &mut use_trees[idx];
+                let lhs_path = lhs_t.path()?;
+                let rhs_path = rhs_path?;
+                let (lhs_prefix, rhs_prefix) = common_prefix(&lhs_path, &rhs_path)?;
+                if lhs_prefix == lhs_path && rhs_prefix == rhs_path {
+                    let tree_is_self = |tree: ast::UseTree| {
+                        tree.path().as_ref().map(path_is_self).unwrap_or(false)
+                    };
+                    // check if only one of the two trees has a tree list, and whether that then contains `self` or not.
+                    // If this is the case we can skip this iteration since the path without the list is already included in the other one via `self`
+                    let tree_contains_self = |tree: &ast::UseTree| {
+                        tree.use_tree_list()
+                            .map(|tree_list| tree_list.use_trees().any(tree_is_self))
+                            .unwrap_or(false)
+                    };
+                    match (tree_contains_self(&lhs_t), tree_contains_self(&rhs_t)) {
+                        (true, false) => continue,
+                        (false, true) => {
+                            *lhs_t = rhs_t;
+                            continue;
+                        }
+                        _ => (),
+                    }
+
+                    // glob imports arent part of the use-tree lists so we need to special handle them here as well
+                    // this special handling is only required for when we merge a module import into a glob import of said module
+                    // see the `merge_self_glob` or `merge_mod_into_glob` tests
+                    if lhs_t.star_token().is_some() || rhs_t.star_token().is_some() {
+                        *lhs_t = make::use_tree(
+                            make::path_unqualified(make::path_segment_self()),
+                            None,
+                            None,
+                            false,
+                        );
+                        use_trees.insert(idx, make::glob_use_tree());
+                        continue;
+                    }
+                }
+                let lhs = lhs_t.split_prefix(&lhs_prefix);
+                let rhs = rhs_t.split_prefix(&rhs_prefix);
+                let this_has_children = use_trees.len() > 0;
+                match recursive_merge(&lhs, &rhs, merge) {
+                    Some((_, has_multiple_children))
+                        if merge == MergeBehaviour::Last
+                            && this_has_children
+                            && has_multiple_children =>
+                    {
+                        return None
+                    }
+                    Some((use_tree, _)) => use_trees[idx] = use_tree,
+                    None => use_trees.insert(idx, rhs_t),
+                }
+            }
+            Err(_)
+                if merge == MergeBehaviour::Last
+                    && use_trees.len() > 0
+                    && rhs_t.use_tree_list().is_some() =>
+            {
+                return None
+            }
+            Err(idx) => {
+                use_trees.insert(idx, rhs_t);
+            }
+        }
+    }
+    let has_multiple_children = use_trees.len() > 1;
+    Some((lhs.with_use_tree_list(make::use_tree_list(use_trees)), has_multiple_children))
 }
 
 /// Traverses both paths until they differ, returning the common prefix of both.
@@ -219,7 +276,7 @@ fn common_prefix(lhs: &ast::Path, rhs: &ast::Path) -> Option<(ast::Path, ast::Pa
     loop {
         match (lhs_curr.segment(), rhs_curr.segment()) {
             (Some(lhs), Some(rhs)) if lhs.syntax().text() == rhs.syntax().text() => (),
-            _ => break,
+            _ => break res,
         }
         res = Some((lhs_curr.clone(), rhs_curr.clone()));
 
@@ -228,11 +285,62 @@ fn common_prefix(lhs: &ast::Path, rhs: &ast::Path) -> Option<(ast::Path, ast::Pa
                 lhs_curr = lhs;
                 rhs_curr = rhs;
             }
-            _ => break,
+            _ => break res,
         }
     }
+}
 
-    res
+fn path_is_self(path: &ast::Path) -> bool {
+    path.segment().and_then(|seg| seg.self_token()).is_some() && path.qualifier().is_none()
+}
+
+#[inline]
+fn first_segment(path: &ast::Path) -> Option<ast::PathSegment> {
+    first_path(path).segment()
+}
+
+fn first_path(path: &ast::Path) -> ast::Path {
+    successors(Some(path.clone()), ast::Path::qualifier).last().unwrap()
+}
+
+fn segment_iter(path: &ast::Path) -> impl Iterator<Item = ast::PathSegment> + Clone {
+    // cant make use of SyntaxNode::siblings, because the returned Iterator is not clone
+    successors(first_segment(path), |p| p.parent_path().parent_path().and_then(|p| p.segment()))
+}
+
+/// Orders paths in the following way:
+/// the sole self token comes first, after that come uppercase identifiers, then lowercase identifiers
+// FIXME: rustfmt sort lowercase idents before uppercase, in general we want to have the same ordering rustfmt has
+// which is `self` and `super` first, then identifier imports with lowercase ones first, then glob imports and at last list imports.
+// Example foo::{self, foo, baz, Baz, Qux, *, {Bar}}
+fn path_cmp(a: &ast::Path, b: &ast::Path) -> Ordering {
+    match (path_is_self(a), path_is_self(b)) {
+        (true, true) => Ordering::Equal,
+        (true, false) => Ordering::Less,
+        (false, true) => Ordering::Greater,
+        (false, false) => {
+            let a = segment_iter(a);
+            let b = segment_iter(b);
+            // cmp_by would be useful for us here but that is currently unstable
+            // cmp doesnt work due the lifetimes on text's return type
+            a.zip(b)
+                .flat_map(|(seg, seg2)| seg.name_ref().zip(seg2.name_ref()))
+                .find_map(|(a, b)| match a.text().cmp(b.text()) {
+                    ord @ Ordering::Greater | ord @ Ordering::Less => Some(ord),
+                    Ordering::Equal => None,
+                })
+                .unwrap_or(Ordering::Equal)
+        }
+    }
+}
+
+fn path_cmp_opt(a: Option<ast::Path>, b: Option<ast::Path>) -> Ordering {
+    match (a, b) {
+        (None, None) => Ordering::Equal,
+        (None, Some(_)) => Ordering::Less,
+        (Some(_), None) => Ordering::Greater,
+        (Some(a), Some(b)) => path_cmp(&a, &b),
+    }
 }
 
 /// What type of merges are allowed.
@@ -279,19 +387,6 @@ impl ImportGroup {
     }
 }
 
-fn first_segment(path: &ast::Path) -> Option<ast::PathSegment> {
-    first_path(path).segment()
-}
-
-fn first_path(path: &ast::Path) -> ast::Path {
-    successors(Some(path.clone()), ast::Path::qualifier).last().unwrap()
-}
-
-fn segment_iter(path: &ast::Path) -> impl Iterator<Item = ast::PathSegment> + Clone {
-    // cant make use of SyntaxNode::siblings, because the returned Iterator is not clone
-    successors(first_segment(path), |p| p.parent_path().parent_path().and_then(|p| p.segment()))
-}
-
 #[derive(PartialEq, Eq)]
 enum AddBlankLine {
     Before,
@@ -594,7 +689,7 @@ use std::io;",
         check_full(
             "std::foo::bar::Baz",
             r"use std::foo::bar::Qux;",
-            r"use std::foo::bar::{Qux, Baz};",
+            r"use std::foo::bar::{Baz, Qux};",
         )
     }
 
@@ -603,7 +698,7 @@ use std::io;",
         check_last(
             "std::foo::bar::Baz",
             r"use std::foo::bar::Qux;",
-            r"use std::foo::bar::{Qux, Baz};",
+            r"use std::foo::bar::{Baz, Qux};",
         )
     }
 
@@ -612,7 +707,7 @@ use std::io;",
         check_full(
             "std::foo::bar::Baz",
             r"use std::foo::bar::{Qux, Quux};",
-            r"use std::foo::bar::{Qux, Quux, Baz};",
+            r"use std::foo::bar::{Baz, Quux, Qux};",
         )
     }
 
@@ -621,7 +716,7 @@ use std::io;",
         check_last(
             "std::foo::bar::Baz",
             r"use std::foo::bar::{Qux, Quux};",
-            r"use std::foo::bar::{Qux, Quux, Baz};",
+            r"use std::foo::bar::{Baz, Quux, Qux};",
         )
     }
 
@@ -630,7 +725,7 @@ use std::io;",
         check_full(
             "std::foo::bar::Baz",
             r"use std::foo::bar::{Qux, quux::{Fez, Fizz}};",
-            r"use std::foo::bar::{Qux, quux::{Fez, Fizz}, Baz};",
+            r"use std::foo::bar::{Baz, Qux, quux::{Fez, Fizz}};",
         )
     }
 
@@ -645,6 +740,15 @@ use std::foo::bar::{Qux, quux::{Fez, Fizz}};",
     }
 
     #[test]
+    fn merge_groups_full_nested_deep() {
+        check_full(
+            "std::foo::bar::quux::Baz",
+            r"use std::foo::bar::{Qux, quux::{Fez, Fizz}};",
+            r"use std::foo::bar::{Qux, quux::{Baz, Fez, Fizz}};",
+        )
+    }
+
+    #[test]
     fn merge_groups_skip_pub() {
         check_full(
             "std::io",
@@ -670,34 +774,63 @@ use std::io;",
         check_last(
             "std::fmt::Result",
             r"use std::{fmt, io};",
-            r"use std::{self, fmt::Result};
+            r"use std::fmt::{self, Result};
 use std::io;",
         )
     }
 
     #[test]
+    fn merge_into_module_import() {
+        check_full(
+            "std::fmt::Result",
+            r"use std::{fmt, io};",
+            r"use std::{fmt::{self, Result}, io};",
+        )
+    }
+
+    #[test]
     fn merge_groups_self() {
         check_full("std::fmt::Debug", r"use std::fmt;", r"use std::fmt::{self, Debug};")
     }
 
     #[test]
-    fn merge_self_glob() {
+    fn merge_mod_into_glob() {
         check_full(
             "token::TokenKind",
             r"use token::TokenKind::*;",
-            r"use token::TokenKind::{self::*, self};",
+            r"use token::TokenKind::{*, self};",
+        )
+        // FIXME: have it emit `use token::TokenKind::{self, *}`?
+    }
+
+    #[test]
+    fn merge_self_glob() {
+        check_full("self", r"use self::*;", r"use self::{*, self};")
+        // FIXME: have it emit `use {self, *}`?
+    }
+
+    #[test]
+    #[ignore] // FIXME: Support this
+    fn merge_partial_path() {
+        check_full(
+            "ast::Foo",
+            r"use syntax::{ast, algo};",
+            r"use syntax::{ast::{self, Foo}, algo};",
+        )
+    }
+
+    #[test]
+    fn merge_glob_nested() {
+        check_full(
+            "foo::bar::quux::Fez",
+            r"use foo::bar::{Baz, quux::*};",
+            r"use foo::bar::{Baz, quux::{self::*, Fez}};",
         )
     }
 
     #[test]
     fn merge_last_too_long() {
-        mark::check!(test_last_merge_too_long);
-        check_last(
-            "foo::bar",
-            r"use foo::bar::baz::Qux;",
-            r"use foo::bar;
-use foo::bar::baz::Qux;",
-        );
+        check_last("foo::bar", r"use foo::bar::baz::Qux;", r"use foo::bar::{self, baz::Qux};");
     }
 
     #[test]
@@ -710,6 +843,42 @@ use foo::bar::baz::Qux;",
         );
     }
 
+    #[test]
+    fn merge_last_fail() {
+        check_merge_only_fail(
+            r"use foo::bar::{baz::{Qux, Fez}};",
+            r"use foo::bar::{baaz::{Quux, Feez}};",
+            MergeBehaviour::Last,
+        );
+    }
+
+    #[test]
+    fn merge_last_fail1() {
+        check_merge_only_fail(
+            r"use foo::bar::{baz::{Qux, Fez}};",
+            r"use foo::bar::baaz::{Quux, Feez};",
+            MergeBehaviour::Last,
+        );
+    }
+
+    #[test]
+    fn merge_last_fail2() {
+        check_merge_only_fail(
+            r"use foo::bar::baz::{Qux, Fez};",
+            r"use foo::bar::{baaz::{Quux, Feez}};",
+            MergeBehaviour::Last,
+        );
+    }
+
+    #[test]
+    fn merge_last_fail3() {
+        check_merge_only_fail(
+            r"use foo::bar::baz::{Qux, Fez};",
+            r"use foo::bar::baaz::{Quux, Feez};",
+            MergeBehaviour::Last,
+        );
+    }
+
     fn check(
         path: &str,
         ra_fixture_before: &str,
@@ -742,4 +911,23 @@ use foo::bar::baz::Qux;",
     fn check_none(path: &str, ra_fixture_before: &str, ra_fixture_after: &str) {
         check(path, ra_fixture_before, ra_fixture_after, None)
     }
+
+    fn check_merge_only_fail(ra_fixture0: &str, ra_fixture1: &str, mb: MergeBehaviour) {
+        let use0 = ast::SourceFile::parse(ra_fixture0)
+            .tree()
+            .syntax()
+            .descendants()
+            .find_map(ast::Use::cast)
+            .unwrap();
+
+        let use1 = ast::SourceFile::parse(ra_fixture1)
+            .tree()
+            .syntax()
+            .descendants()
+            .find_map(ast::Use::cast)
+            .unwrap();
+
+        let result = try_merge_imports(&use0, &use1, mb);
+        assert_eq!(result.map(|u| u.to_string()), None);
+    }
 }