5 files changed, 409 insertions, 15 deletions

diff --git a/crates/ra_assists/src/assist_ctx.rs b/crates/ra_assists/src/assist_ctx.rs
index 43f0d664b..2ab65ab99 100644
--- a/crates/ra_assists/src/assist_ctx.rs
+++ b/crates/ra_assists/src/assist_ctx.rs
@@ -101,7 +101,6 @@ impl<'a, DB: HirDatabase> AssistCtx<'a, DB> {
         Some(assist)
     }
 
-    #[allow(dead_code)] // will be used for auto import assist with multiple actions
     pub(crate) fn add_assist_group(
         self,
         id: AssistId,
@@ -168,7 +167,6 @@ pub(crate) struct ActionBuilder {
 }
 
 impl ActionBuilder {
-    #[allow(dead_code)]
     /// Adds a custom label to the action, if it needs to be different from the assist label
     pub(crate) fn label(&mut self, label: impl Into<String>) {
         self.label = Some(label.into())
diff --git a/crates/ra_assists/src/assists/auto_import.rs b/crates/ra_assists/src/assists/auto_import.rs
new file mode 100644
index 000000000..9163cc662
--- /dev/null
+++ b/crates/ra_assists/src/assists/auto_import.rs
@@ -0,0 +1,222 @@
+use hir::db::HirDatabase;
+use ra_syntax::{
+    ast::{self, AstNode},
+    SmolStr, SyntaxElement,
+    SyntaxKind::{NAME_REF, USE_ITEM},
+    SyntaxNode,
+};
+
+use crate::{
+    assist_ctx::{ActionBuilder, Assist, AssistCtx},
+    auto_import_text_edit, AssistId, ImportsLocator,
+};
+
+// Assist: auto_import
+//
+// If the name is unresolved, provides all possible imports for it.
+//
+// ```
+// fn main() {
+//     let map = HashMap<|>::new();
+// }
+// ```
+// ->
+// ```
+// use std::collections::HashMap;
+//
+// fn main() {
+//     let map = HashMap<|>::new();
+// }
+// ```
+pub(crate) fn auto_import<F: ImportsLocator>(
+    ctx: AssistCtx<impl HirDatabase>,
+    imports_locator: &mut F,
+) -> Option<Assist> {
+    let path: ast::Path = ctx.find_node_at_offset()?;
+    let module = path.syntax().ancestors().find_map(ast::Module::cast);
+    let position = match module.and_then(|it| it.item_list()) {
+        Some(item_list) => item_list.syntax().clone(),
+        None => {
+            let current_file = path.syntax().ancestors().find_map(ast::SourceFile::cast)?;
+            current_file.syntax().clone()
+        }
+    };
+    let source_analyzer = ctx.source_analyzer(&position, None);
+    let module_with_name_to_import = source_analyzer.module()?;
+    let path_to_import = ctx.covering_element().ancestors().find_map(ast::Path::cast)?;
+    if source_analyzer.resolve_path(ctx.db, &path_to_import).is_some() {
+        return None;
+    }
+
+    let name_to_import = &find_applicable_name_ref(ctx.covering_element())?.syntax().to_string();
+    let proposed_imports = imports_locator
+        .find_imports(&name_to_import.to_string())
+        .into_iter()
+        .filter_map(|module_def| module_with_name_to_import.find_use_path(ctx.db, module_def))
+        .filter(|use_path| !use_path.segments.is_empty())
+        .take(20)
+        .map(|import| import.to_string())
+        .collect::<std::collections::BTreeSet<_>>();
+    if proposed_imports.is_empty() {
+        return None;
+    }
+
+    ctx.add_assist_group(AssistId("auto_import"), "auto import", || {
+        proposed_imports
+            .into_iter()
+            .map(|import| import_to_action(import, &position, &path_to_import.syntax()))
+            .collect()
+    })
+}
+
+fn find_applicable_name_ref(element: SyntaxElement) -> Option<ast::NameRef> {
+    if element.ancestors().find(|ancestor| ancestor.kind() == USE_ITEM).is_some() {
+        None
+    } else if element.kind() == NAME_REF {
+        Some(element.as_node().cloned().and_then(ast::NameRef::cast)?)
+    } else {
+        let parent = element.parent()?;
+        if parent.kind() == NAME_REF {
+            Some(ast::NameRef::cast(parent)?)
+        } else {
+            None
+        }
+    }
+}
+
+fn import_to_action(import: String, position: &SyntaxNode, anchor: &SyntaxNode) -> ActionBuilder {
+    let mut action_builder = ActionBuilder::default();
+    action_builder.label(format!("Import `{}`", &import));
+    auto_import_text_edit(
+        position,
+        anchor,
+        &[SmolStr::new(import)],
+        action_builder.text_edit_builder(),
+    );
+    action_builder
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::helpers::{
+        check_assist_with_imports_locator, check_assist_with_imports_locator_not_applicable,
+        TestImportsLocator,
+    };
+
+    #[test]
+    fn applicable_when_found_an_import() {
+        check_assist_with_imports_locator(
+            auto_import,
+            TestImportsLocator::new,
+            r"
+            PubStruct<|>
+
+            pub mod PubMod {
+                pub struct PubStruct;
+            }
+            ",
+            r"
+            use PubMod::PubStruct;
+
+            PubStruct<|>
+
+            pub mod PubMod {
+                pub struct PubStruct;
+            }
+            ",
+        );
+    }
+
+    #[test]
+    fn applicable_when_found_multiple_imports() {
+        check_assist_with_imports_locator(
+            auto_import,
+            TestImportsLocator::new,
+            r"
+            PubStruct<|>
+
+            pub mod PubMod1 {
+                pub struct PubStruct;
+            }
+            pub mod PubMod2 {
+                pub struct PubStruct;
+            }
+            pub mod PubMod3 {
+                pub struct PubStruct;
+            }
+            ",
+            r"
+            use PubMod1::PubStruct;
+
+            PubStruct<|>
+
+            pub mod PubMod1 {
+                pub struct PubStruct;
+            }
+            pub mod PubMod2 {
+                pub struct PubStruct;
+            }
+            pub mod PubMod3 {
+                pub struct PubStruct;
+            }
+            ",
+        );
+    }
+
+    #[test]
+    fn not_applicable_for_already_imported_types() {
+        check_assist_with_imports_locator_not_applicable(
+            auto_import,
+            TestImportsLocator::new,
+            r"
+            use PubMod::PubStruct;
+
+            PubStruct<|>
+
+            pub mod PubMod {
+                pub struct PubStruct;
+            }
+            ",
+        );
+    }
+
+    #[test]
+    fn not_applicable_for_types_with_private_paths() {
+        check_assist_with_imports_locator_not_applicable(
+            auto_import,
+            TestImportsLocator::new,
+            r"
+            PrivateStruct<|>
+
+            pub mod PubMod {
+                struct PrivateStruct;
+            }
+            ",
+        );
+    }
+
+    #[test]
+    fn not_applicable_when_no_imports_found() {
+        check_assist_with_imports_locator_not_applicable(
+            auto_import,
+            TestImportsLocator::new,
+            "
+            PubStruct<|>",
+        );
+    }
+
+    #[test]
+    fn not_applicable_in_import_statements() {
+        check_assist_with_imports_locator_not_applicable(
+            auto_import,
+            TestImportsLocator::new,
+            r"
+            use PubStruct<|>;
+
+            pub mod PubMod {
+                pub struct PubStruct;
+            }",
+        );
+    }
+}
diff --git a/crates/ra_assists/src/doc_tests.rs b/crates/ra_assists/src/doc_tests.rs
index 5dc1ee233..65d51428b 100644
--- a/crates/ra_assists/src/doc_tests.rs
+++ b/crates/ra_assists/src/doc_tests.rs
@@ -11,6 +11,10 @@ use test_utils::{assert_eq_text, extract_range_or_offset};
 use crate::test_db::TestDB;
 
 fn check(assist_id: &str, before: &str, after: &str) {
+    // FIXME we cannot get the imports search functionality here yet, but still need to generate a test and a doc for an assist
+    if assist_id == "auto_import" {
+        return;
+    }
     let (selection, before) = extract_range_or_offset(before);
     let (db, file_id) = TestDB::with_single_file(&before);
     let frange = FileRange { file_id, range: selection.into() };
diff --git a/crates/ra_assists/src/doc_tests/generated.rs b/crates/ra_assists/src/doc_tests/generated.rs
index 7d84dc8fb..ec4587ce7 100644
--- a/crates/ra_assists/src/doc_tests/generated.rs
+++ b/crates/ra_assists/src/doc_tests/generated.rs
@@ -215,6 +215,25 @@ fn main() {
 }
 
 #[test]
+fn doctest_auto_import() {
+    check(
+        "auto_import",
+        r#####"
+fn main() {
+    let map = HashMap<|>::new();
+}
+"#####,
+        r#####"
+use std::collections::HashMap;
+
+fn main() {
+    let map = HashMap<|>::new();
+}
+"#####,
+    )
+}
+
+#[test]
 fn doctest_change_visibility() {
     check(
         "change_visibility",
diff --git a/crates/ra_assists/src/lib.rs b/crates/ra_assists/src/lib.rs
index 3337805a5..625ebc4a2 100644
--- a/crates/ra_assists/src/lib.rs
+++ b/crates/ra_assists/src/lib.rs
@@ -14,7 +14,7 @@ mod test_db;
 pub mod ast_transform;
 
 use either::Either;
-use hir::db::HirDatabase;
+use hir::{db::HirDatabase, ModuleDef};
 use ra_db::FileRange;
 use ra_syntax::{TextRange, TextUnit};
 use ra_text_edit::TextEdit;
@@ -77,6 +77,51 @@ where
     })
 }
 
+/// A functionality for locating imports for the given name.
+///
+/// Currently has to be a trait with the real implementation provided by the ra_ide_api crate,
+/// due to the search functionality located there.
+/// Later, this trait should be removed completely and the search functionality moved to a separate crate,
+/// accessible from the ra_assists crate.
+pub trait ImportsLocator {
+    /// Finds all imports for the given name and the module that contains this name.
+    fn find_imports(&mut self, name_to_import: &str) -> Vec<ModuleDef>;
+}
+
+/// Return all the assists applicable at the given position
+/// and additional assists that need the imports locator functionality to work.
+///
+/// Assists are returned in the "resolved" state, that is with edit fully
+/// computed.
+pub fn assists_with_imports_locator<H, F>(
+    db: &H,
+    range: FileRange,
+    mut imports_locator: F,
+) -> Vec<ResolvedAssist>
+where
+    H: HirDatabase + 'static,
+    F: ImportsLocator,
+{
+    AssistCtx::with_ctx(db, range, true, |ctx| {
+        let mut assists = assists::all()
+            .iter()
+            .map(|f| f(ctx.clone()))
+            .chain(
+                assists::all_with_imports_locator()
+                    .iter()
+                    .map(|f| f(ctx.clone(), &mut imports_locator)),
+            )
+            .filter_map(std::convert::identity)
+            .map(|a| match a {
+                Assist::Resolved { assist } => assist,
+                Assist::Unresolved { .. } => unreachable!(),
+            })
+            .collect();
+        sort_assists(&mut assists);
+        assists
+    })
+}
+
 /// Return all the assists applicable at the given position.
 ///
 /// Assists are returned in the "resolved" state, that is with edit fully
@@ -85,8 +130,6 @@ pub fn assists<H>(db: &H, range: FileRange) -> Vec<ResolvedAssist>
 where
     H: HirDatabase + 'static,
 {
-    use std::cmp::Ordering;
-
     AssistCtx::with_ctx(db, range, true, |ctx| {
         let mut a = assists::all()
             .iter()
@@ -95,19 +138,24 @@ where
                 Assist::Resolved { assist } => assist,
                 Assist::Unresolved { .. } => unreachable!(),
             })
-            .collect::<Vec<_>>();
-        a.sort_by(|a, b| match (a.get_first_action().target, b.get_first_action().target) {
-            (Some(a), Some(b)) => a.len().cmp(&b.len()),
-            (Some(_), None) => Ordering::Less,
-            (None, Some(_)) => Ordering::Greater,
-            (None, None) => Ordering::Equal,
-        });
+            .collect();
+        sort_assists(&mut a);
         a
     })
 }
 
+fn sort_assists(assists: &mut Vec<ResolvedAssist>) {
+    use std::cmp::Ordering;
+    assists.sort_by(|a, b| match (a.get_first_action().target, b.get_first_action().target) {
+        (Some(a), Some(b)) => a.len().cmp(&b.len()),
+        (Some(_), None) => Ordering::Less,
+        (None, Some(_)) => Ordering::Greater,
+        (None, None) => Ordering::Equal,
+    });
+}
+
 mod assists {
-    use crate::{Assist, AssistCtx};
+    use crate::{Assist, AssistCtx, ImportsLocator};
     use hir::db::HirDatabase;
 
     mod add_derive;
@@ -116,6 +164,7 @@ mod assists {
     mod add_custom_impl;
     mod add_new;
     mod apply_demorgan;
+    mod auto_import;
     mod invert_if;
     mod flip_comma;
     mod flip_binexpr;
@@ -168,15 +217,69 @@ mod assists {
             early_return::convert_to_guarded_return,
         ]
     }
+
+    pub(crate) fn all_with_imports_locator<'a, DB: HirDatabase, F: ImportsLocator>(
+    ) -> &'a [fn(AssistCtx<DB>, &mut F) -> Option<Assist>] {
+        &[auto_import::auto_import]
+    }
 }
 
 #[cfg(test)]
 mod helpers {
-    use ra_db::{fixture::WithFixture, FileRange};
+    use hir::db::DefDatabase;
+    use ra_db::{fixture::WithFixture, FileId, FileRange};
     use ra_syntax::TextRange;
     use test_utils::{add_cursor, assert_eq_text, extract_offset, extract_range};
 
-    use crate::{test_db::TestDB, Assist, AssistCtx};
+    use crate::{test_db::TestDB, Assist, AssistCtx, ImportsLocator};
+    use std::sync::Arc;
+
+    // FIXME remove the `ModuleDefId` reexport from `ra_hir` when this gets removed.
+    pub(crate) struct TestImportsLocator {
+        db: Arc<TestDB>,
+        test_file_id: FileId,
+    }
+
+    impl TestImportsLocator {
+        pub(crate) fn new(db: Arc<TestDB>, test_file_id: FileId) -> Self {
+            TestImportsLocator { db, test_file_id }
+        }
+    }
+
+    impl ImportsLocator for TestImportsLocator {
+        fn find_imports(&mut self, name_to_import: &str) -> Vec<hir::ModuleDef> {
+            let crate_def_map = self.db.crate_def_map(self.db.test_crate());
+            let mut findings = Vec::new();
+
+            let mut module_ids_to_process =
+                crate_def_map.modules_for_file(self.test_file_id).collect::<Vec<_>>();
+
+            while !module_ids_to_process.is_empty() {
+                let mut more_ids_to_process = Vec::new();
+                for local_module_id in module_ids_to_process.drain(..) {
+                    for (name, namespace_data) in
+                        crate_def_map[local_module_id].scope.entries_without_primitives()
+                    {
+                        let found_a_match = &name.to_string() == name_to_import;
+                        vec![namespace_data.types, namespace_data.values]
+                            .into_iter()
+                            .filter_map(std::convert::identity)
+                            .for_each(|(module_def_id, _)| {
+                                if found_a_match {
+                                    findings.push(module_def_id.into());
+                                }
+                                if let hir::ModuleDefId::ModuleId(module_id) = module_def_id {
+                                    more_ids_to_process.push(module_id.local_id);
+                                }
+                            });
+                    }
+                }
+                module_ids_to_process = more_ids_to_process;
+            }
+
+            findings
+        }
+    }
 
     pub(crate) fn check_assist(
         assist: fn(AssistCtx<TestDB>) -> Option<Assist>,
@@ -206,6 +309,38 @@ mod helpers {
         assert_eq_text!(after, &actual);
     }
 
+    pub(crate) fn check_assist_with_imports_locator<F: ImportsLocator>(
+        assist: fn(AssistCtx<TestDB>, &mut F) -> Option<Assist>,
+        imports_locator_provider: fn(db: Arc<TestDB>, file_id: FileId) -> F,
+        before: &str,
+        after: &str,
+    ) {
+        let (before_cursor_pos, before) = extract_offset(before);
+        let (db, file_id) = TestDB::with_single_file(&before);
+        let db = Arc::new(db);
+        let mut imports_locator = imports_locator_provider(Arc::clone(&db), file_id);
+        let frange =
+            FileRange { file_id, range: TextRange::offset_len(before_cursor_pos, 0.into()) };
+        let assist =
+            AssistCtx::with_ctx(db.as_ref(), frange, true, |ctx| assist(ctx, &mut imports_locator))
+                .expect("code action is not applicable");
+        let action = match assist {
+            Assist::Unresolved { .. } => unreachable!(),
+            Assist::Resolved { assist } => assist.get_first_action(),
+        };
+
+        let actual = action.edit.apply(&before);
+        let actual_cursor_pos = match action.cursor_position {
+            None => action
+                .edit
+                .apply_to_offset(before_cursor_pos)
+                .expect("cursor position is affected by the edit"),
+            Some(off) => off,
+        };
+        let actual = add_cursor(&actual, actual_cursor_pos);
+        assert_eq_text!(after, &actual);
+    }
+
     pub(crate) fn check_assist_range(
         assist: fn(AssistCtx<TestDB>) -> Option<Assist>,
         before: &str,
@@ -279,6 +414,22 @@ mod helpers {
         assert!(assist.is_none());
     }
 
+    pub(crate) fn check_assist_with_imports_locator_not_applicable<F: ImportsLocator>(
+        assist: fn(AssistCtx<TestDB>, &mut F) -> Option<Assist>,
+        imports_locator_provider: fn(db: Arc<TestDB>, file_id: FileId) -> F,
+        before: &str,
+    ) {
+        let (before_cursor_pos, before) = extract_offset(before);
+        let (db, file_id) = TestDB::with_single_file(&before);
+        let db = Arc::new(db);
+        let mut imports_locator = imports_locator_provider(Arc::clone(&db), file_id);
+        let frange =
+            FileRange { file_id, range: TextRange::offset_len(before_cursor_pos, 0.into()) };
+        let assist =
+            AssistCtx::with_ctx(db.as_ref(), frange, true, |ctx| assist(ctx, &mut imports_locator));
+        assert!(assist.is_none());
+    }
+
     pub(crate) fn check_assist_range_not_applicable(
         assist: fn(AssistCtx<TestDB>) -> Option<Assist>,
         before: &str,