feat: auto import all missing items

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

@@ -2,13 +2,16 @@ use std::cmp::Reverse;
 
 use hir::{Module, db::HirDatabase};
 use ide_db::{
+    FxHashSet,
     helpers::mod_path_to_ast,
     imports::{
         import_assets::{ImportAssets, ImportCandidate, LocatedImport},
-        insert_use::{ImportScope, insert_use, insert_use_as_alias},
+        insert_use::{
+            ImportScope, insert_multiple_use_with_alias_option, insert_use, insert_use_as_alias,
+        },
     },
 };
-use syntax::{AstNode, Edition, NodeOrToken, SyntaxElement, ast};
+use syntax::{AstNode, Edition, NodeOrToken, SyntaxElement, TextRange, WalkEvent, ast};
 
 use crate::{AssistContext, AssistId, Assists, GroupLabel};
 
@@ -201,6 +204,170 @@ pub(super) fn find_importable_node(
     }
 }
 
+/// Represents an import action to be performed
+#[derive(Clone)]
+struct ImportAction {
+    import: LocatedImport,
+    scope: ImportScope,
+    range: TextRange,
+    edition: Edition,
+}
+
+/// Feature: Auto Import All
+/// auto import all missing imports in the current file. Share the same configuration as `auto_import`.
+/// and basically is just calling `auto_import` multiple times. But if multiple proposed imports are provided for one importable node,
+/// then this assist will automatically choose the most relevant one(as api forbid us from proposing multiple choices **multiple times** to user), also for trait import it default to import as alias `_`
+///
+/// Another limitation is that if a `Foo::bar()` where `bar()` is a trait method and both struct `Foo` and trait `Bar` are not in scope, two pass of `auto_import_all` will be required, as in the first pass, there is no way to know whether `bar()` is a trait method or not.
+pub(crate) fn auto_import_all(acc: &mut Assists, ctx: &AssistContext<'_>) -> Option<()> {
+    let cfg = ctx.config.import_path_config();
+
+    let importable_nodes = find_all_importable_nodes(ctx);
+    // dedup all proposed_imports first, so we wouldn't propose the same import multiple times
+    let mut all_proposed_import_paths = FxHashSet::default();
+    let mut all_insert_use_actions = Vec::with_capacity(importable_nodes.len());
+    for (import_assets, syntax_under_caret) in &importable_nodes {
+        let mut proposed_imports: Vec<_> = import_assets
+            .search_for_imports(&ctx.sema, cfg, ctx.config.insert_use.prefix_kind)
+            .collect();
+        // remove duplicates
+        proposed_imports.retain(|import| !all_proposed_import_paths.contains(&import.import_path));
+        if proposed_imports.is_empty() {
+            continue;
+        }
+        all_proposed_import_paths
+            .extend(proposed_imports.iter().map(|import| import.import_path.clone()));
+
+        let range = match &syntax_under_caret {
+            NodeOrToken::Node(node) => ctx.sema.original_range(node).range,
+            NodeOrToken::Token(token) => token.text_range(),
+        };
+        let scope = ImportScope::find_insert_use_container(
+            &match syntax_under_caret {
+                NodeOrToken::Node(it) => it.clone(),
+                NodeOrToken::Token(it) => it.parent()?,
+            },
+            &ctx.sema,
+        )?;
+
+        // we aren't interested in different namespaces
+        proposed_imports.sort_by(|a, b| a.import_path.cmp(&b.import_path));
+        proposed_imports.dedup_by(|a, b| a.import_path == b.import_path);
+
+        let current_module = ctx.sema.scope(scope.as_syntax_node()).map(|scope| scope.module());
+        // prioritize more relevant imports
+        proposed_imports
+            .sort_by_key(|import| Reverse(relevance_score(ctx, import, current_module.as_ref())));
+        let edition =
+            current_module.map(|it| it.krate().edition(ctx.db())).unwrap_or(Edition::CURRENT);
+
+        let chosen_import =
+            proposed_imports.first().expect("should have at least one available import");
+
+        all_insert_use_actions.push(ImportAction {
+            import: chosen_import.clone(),
+            scope,
+            range,
+            edition,
+        });
+    }
+
+    let group_label = GroupLabel("Import all missing items".to_owned());
+    let assist_id = AssistId::quick_fix("auto_import_all");
+    let label = "Import all missing items";
+
+    // add the same import all action in all the places where we need import
+    for action in &all_insert_use_actions {
+        acc.add_group(&group_label, assist_id, label, action.range, |builder| {
+            let path_alias = gen_insert_args(
+                &action.scope,
+                importable_nodes.iter().map(|(import_assets, _)| import_assets),
+                &all_insert_use_actions,
+            );
+
+            let scope = match action.scope.clone() {
+                ImportScope::File(it) => ImportScope::File(builder.make_mut(it)),
+                ImportScope::Module(it) => ImportScope::Module(builder.make_mut(it)),
+                ImportScope::Block(it) => ImportScope::Block(builder.make_mut(it)),
+            };
+
+            insert_multiple_use_with_alias_option(&scope, &path_alias, &ctx.config.insert_use);
+        });
+    }
+    Some(())
+}
+
+/// Find imports in the given scope
+fn gen_insert_args<'a>(
+    given_scope: &ImportScope,
+    import_assets: impl Iterator<Item = &'a ImportAssets>,
+    insert_uses: &[ImportAction],
+) -> Vec<(ast::Path, Option<ast::Rename>)> {
+    let mut path_alias = Vec::new();
+    for (import_assets, action) in import_assets.zip(insert_uses.iter()) {
+        if action.scope.as_syntax_node() != given_scope.as_syntax_node() {
+            continue;
+        }
+
+        match import_assets.import_candidate() {
+            ImportCandidate::TraitAssocItem(_name) | ImportCandidate::TraitMethod(_name) => {
+                // get a ast node `Rename` from `use foo as _`(the same code from `insert_use_as_alias`)
+                let text: &str = "use foo as _";
+                let parse = syntax::SourceFile::parse(text, action.edition);
+                let node = parse
+                    .tree()
+                    .syntax()
+                    .descendants()
+                    .find_map(ast::UseTree::cast)
+                    .expect("Failed to make ast node `Rename`");
+                let alias = node.rename();
+                let ast_path = mod_path_to_ast(&action.import.import_path, action.edition);
+                path_alias.push((ast_path, alias));
+            }
+            _ => {
+                path_alias
+                    .push((mod_path_to_ast(&action.import.import_path, action.edition), None));
+            }
+        }
+    }
+
+    path_alias
+}
+
+pub(super) fn find_all_importable_nodes(
+    ctx: &AssistContext<'_>,
+) -> Vec<(ImportAssets, SyntaxElement)> {
+    // walk from root to find all importable nodes
+    let mut result: Vec<(ImportAssets, SyntaxElement)> = Vec::new();
+
+    for syntax_node in ctx.source_file().syntax().preorder() {
+        let WalkEvent::Enter(syntax_node) = syntax_node else { continue };
+        let node = if ast::Path::can_cast(syntax_node.kind()) {
+            let path_under_caret =
+                ast::Path::cast(syntax_node).expect("Should be able to cast to Path");
+            ImportAssets::for_exact_path(&path_under_caret, &ctx.sema)
+                .zip(Some(path_under_caret.syntax().clone().into()))
+        } else if ast::MethodCallExpr::can_cast(syntax_node.kind()) {
+            let method_under_caret = ast::MethodCallExpr::cast(syntax_node)
+                .expect("Should be able to cast to MethodCallExpr");
+            ImportAssets::for_method_call(&method_under_caret, &ctx.sema)
+                .zip(Some(method_under_caret.syntax().clone().into()))
+        } else if ast::Param::can_cast(syntax_node.kind()) {
+            None
+        } else if let Some(pat) =
+            ast::IdentPat::cast(syntax_node).filter(ast::IdentPat::is_simple_ident)
+        {
+            ImportAssets::for_ident_pat(&ctx.sema, &pat).zip(Some(pat.syntax().clone().into()))
+        } else {
+            None
+        };
+        if let Some(node) = node {
+            result.push(node);
+        }
+    }
+    result
+}
+
 fn group_label(import_candidate: &ImportCandidate) -> GroupLabel {
     let name = match import_candidate {
         ImportCandidate::Path(candidate) => format!("Import {}", candidate.name.text()),
@@ -1722,4 +1889,189 @@ mod foo {
             ",
         );
     }
+
+    #[test]
+    fn basic_auto_import_all() {
+        check_assist(
+            auto_import_all,
+            r"
+mod foo { pub struct Foo; }
+mod bar { pub struct Bar; }
+
+fn main() {
+    Foo$0;
+    Bar;
+}
+",
+            r"
+use {foo::Foo, bar::Bar};
+
+mod foo { pub struct Foo; }
+mod bar { pub struct Bar; }
+
+fn main() {
+    Foo;
+    Bar;
+}
+",
+        )
+    }
+
+    /// only import in current scope
+    #[test]
+    fn auto_import_all_scope() {
+        check_assist_by_label(
+            auto_import_all,
+            r"
+mod foo { pub struct Foo; }
+mod bar {
+    pub struct Bar;
+    fn use_foo(){
+        Foo$0;
+    }
+}
+
+fn main() {
+    Foo;
+    Bar;
+}
+",
+            r"
+mod foo { pub struct Foo; }
+mod bar {
+    use {crate::foo::Foo};
+
+    pub struct Bar;
+    fn use_foo(){
+        Foo;
+    }
+}
+
+fn main() {
+    Foo;
+    Bar;
+}
+",
+            "Import all missing items",
+        )
+    }
+
+    /// notice that only after struct is imported,
+    /// the `ImportAssets` can search and found
+    /// the trait.
+    #[test]
+    fn auto_import_all_two_step_struct_trait() {
+        check_assist(
+            auto_import_all,
+            r"
+mod foo { pub struct Foo; }
+mod bar { pub trait Bar{fn foo() -> bool{true}} }
+
+impl bar::Bar for foo::Foo{}
+
+fn main() {
+Foo::foo()$0;
+}
+",
+            r"
+use {foo::Foo};
+
+mod foo { pub struct Foo; }
+mod bar { pub trait Bar{fn foo() -> bool{true}} }
+
+impl bar::Bar for foo::Foo{}
+
+fn main() {
+Foo::foo();
+}
+",
+        );
+        check_assist(
+            auto_import_all,
+            r"
+use foo::Foo;
+mod foo { pub struct Foo; }
+mod bar { pub trait Bar{fn foo() -> bool{true}} }
+
+impl bar::Bar for foo::Foo{}
+
+fn main() {
+    Foo::foo()$0;
+}
+",
+            r"
+use foo::Foo;
+
+use {bar::Bar as _};
+mod foo { pub struct Foo; }
+mod bar { pub trait Bar{fn foo() -> bool{true}} }
+
+impl bar::Bar for foo::Foo{}
+
+fn main() {
+    Foo::foo();
+}
+",
+        );
+    }
+
+    #[test]
+    fn auto_import_all_multiple_scopes() {
+        // This test verifies that the ImportAction struct correctly handles
+        // multiple scopes and different import types
+        check_assist_by_label(
+            auto_import_all,
+            r"
+mod foo {
+    pub struct Foo;
+    pub trait FooTrait { fn foo_method() {} }
+}
+mod bar {
+    pub struct Bar;
+    pub mod nested { pub struct Nested; }
+}
+
+fn main() {
+    // Multiple imports in the same scope
+    Foo$0;
+    Bar;
+
+    // Nested import
+    nested::Nested;
+
+    // Block scope with trait method
+    {
+        Foo::foo_method();
+    }
+}
+",
+            r"
+use {foo::Foo, bar::Bar, bar::nested};
+
+mod foo {
+    pub struct Foo;
+    pub trait FooTrait { fn foo_method() {} }
+}
+mod bar {
+    pub struct Bar;
+    pub mod nested { pub struct Nested; }
+}
+
+fn main() {
+    // Multiple imports in the same scope
+    Foo;
+    Bar;
+
+    // Nested import
+    nested::Nested;
+
+    // Block scope with trait method
+    {
+        Foo::foo_method();
+    }
+}
+",
+            "Import all missing items",
+        );
+    }
 }

crates/ide-assists/src/lib.rs

@@ -245,6 +245,7 @@ mod handlers {
             apply_demorgan::apply_demorgan_iterator,
             apply_demorgan::apply_demorgan,
             auto_import::auto_import,
+            auto_import::auto_import_all,
             bind_unused_param::bind_unused_param,
             change_visibility::change_visibility,
             convert_bool_then::convert_bool_then_to_if,