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 <lukastw97@gmail.com>

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};
 ",
         )
     }

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;

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))
+    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))
 }
 
-/// 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()
-    })
-}
-
-// 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()?;
+    let lhs = lhs.split_prefix(&lhs_prefix);
+    let rhs = rhs.split_prefix(&rhs_prefix);
+    recursive_merge(&lhs, &rhs, merge).map(|(merged, _)| merged)
+}
 
-    // 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;
-    }
+/// 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;
+                        }
+                        _ => (),
+                    }
 
-    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());
+                    // 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);
+            }
+        }
     }
-    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))
+    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}};",
         )
     }
 
@@ -644,6 +739,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(
@@ -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);
+    }
 }

crates/syntax/src/ast/edit.rs

@@ -347,6 +347,7 @@ impl ast::UseTree {
         self.clone()
     }
 
+    /// Splits off the given prefix, making it the path component of the use tree, appending the rest of the path to all UseTreeList items.
     #[must_use]
     pub fn split_prefix(&self, prefix: &ast::Path) -> ast::UseTree {
         let suffix = if self.path().as_ref() == Some(prefix) && self.use_tree_list().is_none() {

crates/syntax/src/ast/make.rs

@@ -38,6 +38,10 @@ pub fn path_from_text(text: &str) -> ast::Path {
     ast_from_text(text)
 }
 
+pub fn glob_use_tree() -> ast::UseTree {
+    ast_from_text("use *;")
+}
+
 pub fn use_tree(
     path: ast::Path,
     use_tree_list: Option<ast::UseTreeList>,