3 files changed, 150 insertions, 61 deletions

diff --git a/crates/ra_assists/src/assists/auto_import.rs b/crates/ra_assists/src/assists/auto_import.rs
index d196f6c5c..295fdf2e2 100644
--- a/crates/ra_assists/src/assists/auto_import.rs
+++ b/crates/ra_assists/src/assists/auto_import.rs
@@ -100,88 +100,122 @@ mod tests {
     use super::*;
     use crate::helpers::{
         check_assist_with_imports_locator, check_assist_with_imports_locator_not_applicable,
+        TestImportsLocator,
     };
-    use hir::Name;
 
-    #[derive(Clone)]
-    struct TestImportsLocator<'a> {
-        import_path: &'a [Name],
-    }
+    #[test]
+    fn applicable_when_found_an_import() {
+        check_assist_with_imports_locator(
+            auto_import,
+            TestImportsLocator::new,
+            r"
+            PubStruct<|>
 
-    impl<'a> TestImportsLocator<'a> {
-        fn new(import_path: &'a [Name]) -> Self {
-            TestImportsLocator { import_path }
-        }
-    }
+            pub mod PubMod {
+                pub struct PubStruct;
+            }
+            ",
+            r"
+            use PubMod::PubStruct;
 
-    impl<'a> ImportsLocator for TestImportsLocator<'a> {
-        fn find_imports(
-            &mut self,
-            _: hir::InFile<&ast::NameRef>,
-            _: hir::Module,
-        ) -> Option<Vec<hir::ModPath>> {
-            if self.import_path.is_empty() {
-                None
-            } else {
-                Some(vec![hir::ModPath {
-                    kind: hir::PathKind::Plain,
-                    segments: self.import_path.to_owned(),
-                }])
+            PubStruct<|>
+
+            pub mod PubMod {
+                pub struct PubStruct;
             }
-        }
+            ",
+        );
     }
 
     #[test]
-    fn applicable_when_found_an_import() {
-        let import_path = &[hir::name::known::std, hir::name::known::ops, hir::name::known::Debug];
-        let mut imports_locator = TestImportsLocator::new(import_path);
+    fn applicable_when_found_multiple_imports() {
         check_assist_with_imports_locator(
             auto_import,
-            &mut imports_locator,
-            "
-            fn main() {
+            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<|>
 
-            Debug<|>",
-            &format!(
-                "
-            use {};
-
-            fn main() {{
-            }}
-
-            Debug<|>",
-                import_path
-                    .into_iter()
-                    .map(|name| name.to_string())
-                    .collect::<Vec<String>>()
-                    .join("::")
-            ),
+            pub mod PubMod {
+                pub struct PubStruct;
+            }
+            ",
         );
     }
 
     #[test]
-    fn not_applicable_when_no_imports_found() {
-        let mut imports_locator = TestImportsLocator::new(&[]);
+    fn not_applicable_for_types_with_private_paths() {
         check_assist_with_imports_locator_not_applicable(
             auto_import,
-            &mut imports_locator,
-            "
-            fn main() {
+            TestImportsLocator::new,
+            r"
+            PrivateStruct<|>
+
+            pub mod PubMod {
+                struct PrivateStruct;
             }
+            ",
+        );
+    }
 
-            Debug<|>",
+    #[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() {
-        let import_path = &[hir::name::known::std, hir::name::known::ops, hir::name::known::Debug];
-        let mut imports_locator = TestImportsLocator::new(import_path);
         check_assist_with_imports_locator_not_applicable(
             auto_import,
-            &mut imports_locator,
-            "use Debug<|>;",
+            TestImportsLocator::new,
+            r"
+            use PubStruct<|>;
+
+            pub mod PubMod {
+                pub struct PubStruct;
+            }",
         );
     }
 }
diff --git a/crates/ra_assists/src/lib.rs b/crates/ra_assists/src/lib.rs
index 9e4ebec47..724bce191 100644
--- a/crates/ra_assists/src/lib.rs
+++ b/crates/ra_assists/src/lib.rs
@@ -226,11 +226,59 @@ mod assists {
 
 #[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, ImportsLocator};
+    use std::sync::Arc;
+
+    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![];
+
+            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![];
+                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>,
@@ -262,16 +310,19 @@ mod helpers {
 
     pub(crate) fn check_assist_with_imports_locator<F: ImportsLocator>(
         assist: fn(AssistCtx<TestDB>, &mut F) -> Option<Assist>,
-        imports_locator: &mut F,
+        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, frange, true, |ctx| assist(ctx, imports_locator))
-            .expect("code action is not applicable");
+        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(),
@@ -364,14 +415,17 @@ mod helpers {
 
     pub(crate) fn check_assist_with_imports_locator_not_applicable<F: ImportsLocator>(
         assist: fn(AssistCtx<TestDB>, &mut F) -> Option<Assist>,
-        imports_locator: &mut F,
+        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, frange, true, |ctx| assist(ctx, imports_locator));
+        let assist =
+            AssistCtx::with_ctx(db.as_ref(), frange, true, |ctx| assist(ctx, &mut imports_locator));
         assert!(assist.is_none());
     }
 
diff --git a/crates/ra_hir/src/lib.rs b/crates/ra_hir/src/lib.rs
index 21b0553be..2cfeaba72 100644
--- a/crates/ra_hir/src/lib.rs
+++ b/crates/ra_hir/src/lib.rs
@@ -56,6 +56,7 @@ pub use hir_def::{
     nameres::ModuleSource,
     path::{ModPath, Path, PathKind},
     type_ref::Mutability,
+    ModuleDefId,
 };
 pub use hir_expand::{
     name::{AsName, Name},