normalize use trees when merging imports

5 files changed, 350 insertions, 26 deletions

diff --git a/crates/ide-assists/src/handlers/merge_imports.rs b/crates/ide-assists/src/handlers/merge_imports.rs
index 2beab26dce7..02dbde803e5 100644
--- a/crates/ide-assists/src/handlers/merge_imports.rs
+++ b/crates/ide-assists/src/handlers/merge_imports.rs
@@ -201,7 +201,7 @@ use std::fmt$0::{Display, Debug};
 use std::fmt::{Display, Debug};
 ",
             r"
-use std::fmt::{Display, Debug};
+use std::fmt::{Debug, Display};
 ",
         );
 
@@ -214,7 +214,7 @@ use std::fmt::{Display, Debug};
         check_assist_import_one_variations!(
             "std::fmt$0::{Display, Debug}",
             "std::fmt::{Display, Debug}",
-            "use {std::fmt::{Display, Debug}};"
+            "use {std::fmt::{Debug, Display}};"
         );
     }
 
@@ -419,13 +419,13 @@ use std$0::{fmt::{Write, Display}};
 use std::{fmt::{self, Debug}};
 ",
             r"
-use std::{fmt::{self, Debug, Display, Write}};
+use std::fmt::{self, Debug, Display, Write};
 ",
         );
         check_assist_import_one_variations!(
             "std$0::{fmt::{Write, Display}}",
             "std::{fmt::{self, Debug}}",
-            "use {std::{fmt::{self, Debug, Display, Write}}};"
+            "use {std::fmt::{self, Debug, Display, Write}};"
         );
     }
 
@@ -438,13 +438,13 @@ use std$0::{fmt::{self, Debug}};
 use std::{fmt::{Write, Display}};
 ",
             r"
-use std::{fmt::{self, Debug, Display, Write}};
+use std::fmt::{self, Debug, Display, Write};
 ",
         );
         check_assist_import_one_variations!(
             "std$0::{fmt::{self, Debug}}",
             "std::{fmt::{Write, Display}}",
-            "use {std::{fmt::{self, Debug, Display, Write}}};"
+            "use {std::fmt::{self, Debug, Display, Write}};"
         );
     }
 
@@ -470,13 +470,13 @@ use foo::$0{bar::{self}};
 use foo::{bar};
 ",
             r"
-use foo::{bar::{self}};
+use foo::bar;
 ",
         );
         check_assist_import_one_variations!(
             "foo::$0{bar::{self}}",
             "foo::{bar}",
-            "use {foo::{bar::{self}}};"
+            "use {foo::bar};"
         );
     }
 
@@ -489,13 +489,13 @@ use foo::$0{bar};
 use foo::{bar::{self}};
 ",
             r"
-use foo::{bar::{self}};
+use foo::bar;
 ",
         );
         check_assist_import_one_variations!(
             "foo::$0{bar}",
             "foo::{bar::{self}}",
-            "use {foo::{bar::{self}}};"
+            "use {foo::bar};"
         );
     }
 
@@ -508,13 +508,13 @@ use std$0::{fmt::*};
 use std::{fmt::{self, Display}};
 ",
             r"
-use std::{fmt::{self, Display, *}};
+use std::fmt::{self, Display, *};
 ",
         );
         check_assist_import_one_variations!(
             "std$0::{fmt::*}",
             "std::{fmt::{self, Display}}",
-            "use {std::{fmt::{self, Display, *}}};"
+            "use {std::fmt::{self, Display, *}};"
         );
     }
 
diff --git a/crates/ide-completion/src/tests/flyimport.rs b/crates/ide-completion/src/tests/flyimport.rs
index 1af16ef857d..eaa1bebc03c 100644
--- a/crates/ide-completion/src/tests/flyimport.rs
+++ b/crates/ide-completion/src/tests/flyimport.rs
@@ -106,7 +106,7 @@ fn main() {
 }
 "#,
         r#"
-use dep::{FirstStruct, some_module::{SecondStruct, ThirdStruct}};
+use dep::{some_module::{SecondStruct, ThirdStruct}, FirstStruct};
 
 fn main() {
     ThirdStruct
diff --git a/crates/ide-db/src/imports/insert_use.rs b/crates/ide-db/src/imports/insert_use.rs
index 09b4a1c1baa..5f5ec446870 100644
--- a/crates/ide-db/src/imports/insert_use.rs
+++ b/crates/ide-db/src/imports/insert_use.rs
@@ -17,6 +17,7 @@ use syntax::{
 use crate::{
     imports::merge_imports::{
         common_prefix, eq_attrs, eq_visibility, try_merge_imports, use_tree_cmp, MergeBehavior,
+        NormalizationStyle,
     },
     RootDatabase,
 };
@@ -40,6 +41,15 @@ pub enum ImportGranularity {
     One,
 }
 
+impl From<ImportGranularity> for NormalizationStyle {
+    fn from(granularity: ImportGranularity) -> Self {
+        match granularity {
+            ImportGranularity::One => NormalizationStyle::One,
+            _ => NormalizationStyle::Default,
+        }
+    }
+}
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
 pub struct InsertUseConfig {
     pub granularity: ImportGranularity,
diff --git a/crates/ide-db/src/imports/insert_use/tests.rs b/crates/ide-db/src/imports/insert_use/tests.rs
index 2ed60698871..6b0fecae267 100644
--- a/crates/ide-db/src/imports/insert_use/tests.rs
+++ b/crates/ide-db/src/imports/insert_use/tests.rs
@@ -635,7 +635,7 @@ use std::io;",
     check_one(
         "std::io",
         r"use {std::fmt::{Result, Display}};",
-        r"use {std::{fmt::{Result, Display}, io}};",
+        r"use {std::{fmt::{Display, Result}, io}};",
     );
 }
 
@@ -650,12 +650,12 @@ fn merge_groups_full() {
     check_crate(
         "std::io",
         r"use std::fmt::{Result, Display};",
-        r"use std::{fmt::{Result, Display}, io};",
+        r"use std::{fmt::{Display, Result}, io};",
     );
     check_one(
         "std::io",
         r"use {std::fmt::{Result, Display}};",
-        r"use {std::{fmt::{Result, Display}, io}};",
+        r"use {std::{fmt::{Display, Result}, io}};",
     );
 }
 
@@ -749,12 +749,12 @@ fn merge_groups_full_nested_deep() {
     check_crate(
         "std::foo::bar::quux::Baz",
         r"use std::foo::bar::{Qux, quux::{Fez, Fizz}};",
-        r"use std::foo::bar::{Qux, quux::{Baz, Fez, Fizz}};",
+        r"use std::foo::bar::{quux::{Baz, Fez, Fizz}, Qux};",
     );
     check_one(
         "std::foo::bar::quux::Baz",
         r"use {std::foo::bar::{Qux, quux::{Fez, Fizz}}};",
-        r"use {std::foo::bar::{Qux, quux::{Baz, Fez, Fizz}}};",
+        r"use {std::foo::bar::{quux::{Baz, Fez, Fizz}, Qux}};",
     );
 }
 
@@ -763,7 +763,7 @@ fn merge_groups_full_nested_long() {
     check_crate(
         "std::foo::bar::Baz",
         r"use std::{foo::bar::Qux};",
-        r"use std::{foo::bar::{Baz, Qux}};",
+        r"use std::foo::bar::{Baz, Qux};",
     );
 }
 
@@ -772,12 +772,12 @@ fn merge_groups_last_nested_long() {
     check_crate(
         "std::foo::bar::Baz",
         r"use std::{foo::bar::Qux};",
-        r"use std::{foo::bar::{Baz, Qux}};",
+        r"use std::foo::bar::{Baz, Qux};",
     );
     check_one(
         "std::foo::bar::Baz",
         r"use {std::{foo::bar::Qux}};",
-        r"use {std::{foo::bar::{Baz, Qux}}};",
+        r"use {std::foo::bar::{Baz, Qux}};",
     );
 }
 
@@ -898,7 +898,7 @@ fn merge_glob() {
         r"
 use syntax::{SyntaxKind::*};",
         r"
-use syntax::{SyntaxKind::{self, *}};",
+use syntax::SyntaxKind::{self, *};",
     )
 }
 
@@ -907,7 +907,7 @@ fn merge_glob_nested() {
     check_crate(
         "foo::bar::quux::Fez",
         r"use foo::bar::{Baz, quux::*};",
-        r"use foo::bar::{Baz, quux::{Fez, *}};",
+        r"use foo::bar::{quux::{Fez, *}, Baz};",
     )
 }
 
@@ -1211,7 +1211,7 @@ fn insert_with_renamed_import_complex_use() {
 use self::foo::{self, Foo as _, Bar};
 "#,
         r#"
-use self::foo::{self, Foo, Bar};
+use self::foo::{self, Bar, Foo};
 "#,
         &InsertUseConfig {
             granularity: ImportGranularity::Crate,
diff --git a/crates/ide-db/src/imports/merge_imports.rs b/crates/ide-db/src/imports/merge_imports.rs
index 7ec38c317df..5c97ffaf4b4 100644
--- a/crates/ide-db/src/imports/merge_imports.rs
+++ b/crates/ide-db/src/imports/merge_imports.rs
@@ -7,9 +7,12 @@ use parser::T;
 use stdx::is_upper_snake_case;
 use syntax::{
     algo,
-    ast::{self, make, AstNode, HasAttrs, HasName, HasVisibility, PathSegmentKind},
+    ast::{
+        self, edit_in_place::Removable, make, AstNode, HasAttrs, HasName, HasVisibility,
+        PathSegmentKind,
+    },
     ted::{self, Position},
-    Direction,
+    Direction, SyntaxElement,
 };
 
 use crate::syntax_helpers::node_ext::vis_eq;
@@ -58,6 +61,10 @@ pub fn try_merge_imports(
     let lhs_tree = lhs.use_tree()?;
     let rhs_tree = rhs.use_tree()?;
     try_merge_trees_mut(&lhs_tree, &rhs_tree, merge_behavior)?;
+
+    // Ignore `None` result because normalization should not affect the merge result.
+    try_normalize_use_tree_mut(&lhs_tree, merge_behavior.into());
+
     Some(lhs)
 }
 
@@ -71,6 +78,10 @@ pub fn try_merge_trees(
     let lhs = lhs.clone_subtree().clone_for_update();
     let rhs = rhs.clone_subtree().clone_for_update();
     try_merge_trees_mut(&lhs, &rhs, merge)?;
+
+    // Ignore `None` result because normalization should not affect the merge result.
+    try_normalize_use_tree_mut(&lhs, merge.into());
+
     Some(lhs)
 }
 
@@ -232,6 +243,309 @@ fn recursive_merge(lhs: &ast::UseTree, rhs: &ast::UseTree, merge: MergeBehavior)
     Some(())
 }
 
+/// Style to follow when normalizing a use tree.
+#[derive(Copy, Clone, Debug, PartialEq, Eq)]
+pub enum NormalizationStyle {
+    /// Merges all descendant use tree lists with only one child use tree into their parent use tree.
+    ///
+    /// Examples:
+    /// - `foo::{bar::{Qux}}` -> `foo::bar::Qux`
+    /// - `foo::{bar::{self}}` -> `foo::bar`
+    /// - `{foo::bar}` -> `foo::bar`
+    Default,
+    /// Same as default but wraps the root use tree in a use tree list.
+    ///
+    /// Examples:
+    /// - `foo::{bar::{Qux}}` -> `{foo::bar::Qux}`
+    /// - `foo::{bar::{self}}` -> `{foo::bar}`
+    /// - `{foo::bar}` -> `{foo::bar}`
+    One,
+}
+
+impl From<MergeBehavior> for NormalizationStyle {
+    fn from(mb: MergeBehavior) -> Self {
+        match mb {
+            MergeBehavior::One => NormalizationStyle::One,
+            _ => NormalizationStyle::Default,
+        }
+    }
+}
+
+/// Normalizes a use item by:
+/// - Ordering all use trees
+/// - Merging use trees with common prefixes
+/// - Removing redundant braces based on the specified normalization style
+///   (see [`NormalizationStyle`] doc)
+///
+/// Examples:
+///
+/// Using the "Default" normalization style
+///
+/// - `foo::{bar::Qux, bar::{self}}` -> `foo::bar::{self, Qux}`
+/// - `foo::bar::{self}` -> `foo::bar`
+/// - `{foo::bar}` -> `foo::bar`
+///
+/// Using the "One" normalization style
+///
+/// - `foo::{bar::Qux, bar::{self}}` -> `{foo::bar::{self, Qux}}`
+/// - `foo::bar::{self}` -> `{foo::bar}`
+/// - `foo::bar` -> `{foo::bar}`
+pub fn try_normalize_import(use_item: &ast::Use, style: NormalizationStyle) -> Option<ast::Use> {
+    let use_item = use_item.clone_subtree().clone_for_update();
+    try_normalize_use_tree_mut(&use_item.use_tree()?, style)?;
+    Some(use_item)
+}
+
+/// Normalizes a use tree (see [`try_normalize_import`] doc).
+pub fn try_normalize_use_tree(
+    use_tree: &ast::UseTree,
+    style: NormalizationStyle,
+) -> Option<ast::UseTree> {
+    let use_tree = use_tree.clone_subtree().clone_for_update();
+    try_normalize_use_tree_mut(&use_tree, style)?;
+    Some(use_tree)
+}
+
+macro_rules! call_and_track_result {
+    ($call:expr, $tracker: ident) => {
+        let result = $call;
+        if !$tracker && result.is_some() {
+            $tracker = true;
+        }
+    };
+}
+
+pub fn try_normalize_use_tree_mut(
+    use_tree: &ast::UseTree,
+    style: NormalizationStyle,
+) -> Option<()> {
+    if style == NormalizationStyle::One {
+        let mut modified = false;
+        call_and_track_result!(use_tree.wrap_in_tree_list(), modified);
+        call_and_track_result!(recursive_normalize(use_tree, style), modified);
+        if !modified {
+            // Either the use tree was already normalized or its semantically empty.
+            return None;
+        }
+    } else {
+        recursive_normalize(use_tree, NormalizationStyle::Default)?;
+    }
+    Some(())
+}
+
+/// Recursively normalizes a use tree and its subtrees (if any).
+fn recursive_normalize(use_tree: &ast::UseTree, style: NormalizationStyle) -> Option<()> {
+    let use_tree_list = use_tree.use_tree_list()?;
+    let merge_subtree_into_parent_tree = |single_subtree: &ast::UseTree| {
+        let merged_path = match (use_tree.path(), single_subtree.path()) {
+            (None, None) => None,
+            (Some(outer), None) => Some(outer),
+            (None, Some(inner)) if path_is_self(&inner) => None,
+            (None, Some(inner)) => Some(inner),
+            (Some(outer), Some(inner)) if path_is_self(&inner) => Some(outer),
+            (Some(outer), Some(inner)) => Some(make::path_concat(outer, inner).clone_for_update()),
+        };
+        if merged_path.is_some()
+            || single_subtree.use_tree_list().is_some()
+            || single_subtree.star_token().is_some()
+        {
+            ted::remove_all_iter(use_tree.syntax().children_with_tokens());
+            if let Some(path) = merged_path {
+                ted::insert_raw(Position::first_child_of(use_tree.syntax()), path.syntax());
+                if single_subtree.use_tree_list().is_some() || single_subtree.star_token().is_some()
+                {
+                    ted::insert_raw(
+                        Position::last_child_of(use_tree.syntax()),
+                        make::token(T![::]),
+                    );
+                }
+            }
+            if let Some(inner_use_tree_list) = single_subtree.use_tree_list() {
+                ted::insert_raw(
+                    Position::last_child_of(use_tree.syntax()),
+                    inner_use_tree_list.syntax(),
+                );
+            } else if single_subtree.star_token().is_some() {
+                ted::insert_raw(Position::last_child_of(use_tree.syntax()), make::token(T![*]));
+            } else if let Some(rename) = single_subtree.rename() {
+                ted::insert_raw(
+                    Position::last_child_of(use_tree.syntax()),
+                    make::tokens::single_space(),
+                );
+                ted::insert_raw(Position::last_child_of(use_tree.syntax()), rename.syntax());
+            }
+            Some(())
+        } else {
+            // Bail on semantically empty use trees.
+            None
+        }
+    };
+    let one_style_tree_list = |subtree: &ast::UseTree| match (
+        subtree.path().is_none() && subtree.star_token().is_none() && subtree.rename().is_none(),
+        subtree.use_tree_list(),
+    ) {
+        (true, tree_list) => tree_list,
+        _ => None,
+    };
+    let add_element_to_list = |elem: SyntaxElement, elements: &mut Vec<SyntaxElement>| {
+        if !elements.is_empty() {
+            elements.push(make::token(T![,]).into());
+            elements.push(make::tokens::single_space().into());
+        }
+        elements.push(elem);
+    };
+    if let Some((single_subtree,)) = use_tree_list.use_trees().collect_tuple() {
+        if style == NormalizationStyle::One {
+            // Only normalize descendant subtrees if the normalization style is "one".
+            recursive_normalize(&single_subtree, NormalizationStyle::Default)?;
+        } else {
+            // Otherwise, merge the single subtree into it's parent (if possible)
+            // and then normalize the result.
+            merge_subtree_into_parent_tree(&single_subtree)?;
+            recursive_normalize(use_tree, style);
+        }
+    } else {
+        // Tracks whether any changes have been made to the use tree.
+        let mut modified = false;
+
+        // Recursively un-nests (if necessary) and then normalizes each subtree in the tree list.
+        for subtree in use_tree_list.use_trees() {
+            if let Some(one_tree_list) = one_style_tree_list(&subtree) {
+                let mut elements = Vec::new();
+                let mut one_tree_list_iter = one_tree_list.use_trees();
+                let mut prev_skipped = Vec::new();
+                loop {
+                    let mut prev_skipped_iter = prev_skipped.into_iter();
+                    let mut curr_skipped = Vec::new();
+
+                    while let Some(sub_sub_tree) =
+                        one_tree_list_iter.next().or(prev_skipped_iter.next())
+                    {
+                        if let Some(sub_one_tree_list) = one_style_tree_list(&sub_sub_tree) {
+                            curr_skipped.extend(sub_one_tree_list.use_trees());
+                        } else {
+                            call_and_track_result!(
+                                recursive_normalize(&sub_sub_tree, NormalizationStyle::Default),
+                                modified
+                            );
+                            add_element_to_list(
+                                sub_sub_tree.syntax().clone().into(),
+                                &mut elements,
+                            );
+                        }
+                    }
+
+                    if curr_skipped.is_empty() {
+                        // Un-nesting is complete.
+                        break;
+                    }
+                    prev_skipped = curr_skipped;
+                }
+
+                // Either removes the subtree (if its semantically empty) or replaces it with
+                // the un-nested elements.
+                if elements.is_empty() {
+                    subtree.remove();
+                } else {
+                    ted::replace_with_many(subtree.syntax(), elements);
+                }
+                modified = true;
+            } else {
+                call_and_track_result!(
+                    recursive_normalize(&subtree, NormalizationStyle::Default),
+                    modified
+                );
+            }
+        }
+
+        // Merge all merge-able subtrees.
+        let mut tree_list_iter = use_tree_list.use_trees();
+        let mut anchor = tree_list_iter.next()?;
+        let mut prev_skipped = Vec::new();
+        loop {
+            let mut has_merged = false;
+            let mut prev_skipped_iter = prev_skipped.into_iter();
+            let mut next_anchor = None;
+            let mut curr_skipped = Vec::new();
+
+            while let Some(candidate) = tree_list_iter.next().or(prev_skipped_iter.next()) {
+                let result = try_merge_trees_mut(&anchor, &candidate, MergeBehavior::Crate);
+                if result.is_some() {
+                    // Remove merged subtree.
+                    candidate.remove();
+                    has_merged = true;
+                } else if next_anchor.is_none() {
+                    next_anchor = Some(candidate);
+                } else {
+                    curr_skipped.push(candidate);
+                }
+            }
+
+            if has_merged {
+                // Normalize the merge result.
+                recursive_normalize(&anchor, NormalizationStyle::Default);
+                modified = true;
+            }
+
+            let (Some(next_anchor), true) = (next_anchor, !curr_skipped.is_empty()) else {
+                // Merging is complete.
+                break;
+            };
+
+            // Try to merge the remaining subtrees in the next iteration.
+            anchor = next_anchor;
+            prev_skipped = curr_skipped;
+        }
+
+        let mut subtrees: Vec<_> = use_tree_list.use_trees().collect();
+        // Merge the remaining subtree into its parent, if its only one and
+        // the normalization style is not "one".
+        if subtrees.len() == 1 && style != NormalizationStyle::One {
+            call_and_track_result!(merge_subtree_into_parent_tree(&subtrees[0]), modified);
+        }
+        // Order the remaining subtrees (if necessary).
+        if subtrees.len() > 1 {
+            let mut did_sort = false;
+            subtrees.sort_unstable_by(|a, b| {
+                let order = use_tree_cmp_bin_search(a, b);
+                if !did_sort && order == Ordering::Less {
+                    did_sort = true;
+                }
+                order
+            });
+            if did_sort {
+                let start = use_tree_list
+                    .l_curly_token()
+                    .and_then(|l_curly| algo::non_trivia_sibling(l_curly.into(), Direction::Next))
+                    .filter(|it| it.kind() != T!['}']);
+                let end = use_tree_list
+                    .r_curly_token()
+                    .and_then(|r_curly| algo::non_trivia_sibling(r_curly.into(), Direction::Prev))
+                    .filter(|it| it.kind() != T!['{']);
+                if let Some((start, end)) = start.zip(end) {
+                    // Attempt to insert elements while preserving preceding and trailing trivia.
+                    let mut elements = Vec::new();
+                    for subtree in subtrees {
+                        add_element_to_list(subtree.syntax().clone().into(), &mut elements);
+                    }
+                    ted::replace_all(start..=end, elements);
+                } else {
+                    let new_use_tree_list =
+                        make::use_tree_list(subtrees.into_iter()).clone_for_update();
+                    ted::replace(use_tree_list.syntax(), new_use_tree_list.syntax());
+                }
+                modified = true;
+            }
+        }
+
+        if !modified {
+            // Either the use tree was already normalized or its semantically empty.
+            return None;
+        }
+    }
+    Some(())
+}
+
 /// Traverses both paths until they differ, returning the common prefix of both.
 pub fn common_prefix(lhs: &ast::Path, rhs: &ast::Path) -> Option<(ast::Path, ast::Path)> {
     let mut res = None;