Draft the qualifier import resolution

4 files changed, 224 insertions, 43 deletions

diff --git a/crates/ide_assists/src/handlers/qualify_path.rs b/crates/ide_assists/src/handlers/qualify_path.rs
index 1fb4931cd..c0463311e 100644
--- a/crates/ide_assists/src/handlers/qualify_path.rs
+++ b/crates/ide_assists/src/handlers/qualify_path.rs
@@ -1,7 +1,10 @@
 use std::iter;
 
 use hir::AsAssocItem;
-use ide_db::helpers::{import_assets::ImportCandidate, mod_path_to_ast};
+use ide_db::helpers::{
+    import_assets::{ImportCandidate, Qualifier},
+    mod_path_to_ast,
+};
 use ide_db::RootDatabase;
 use syntax::{
     ast,
@@ -45,7 +48,7 @@ pub(crate) fn qualify_path(acc: &mut Assists, ctx: &AssistContext) -> Option<()>
 
     let qualify_candidate = match candidate {
         ImportCandidate::Path(candidate) => {
-            if candidate.unresolved_qualifier.is_some() {
+            if !matches!(candidate.qualifier, Qualifier::Absent) {
                 cov_mark::hit!(qualify_path_qualifier_start);
                 let path = ast::Path::cast(syntax_under_caret)?;
                 let (prev_segment, segment) = (path.qualifier()?.segment()?, path.segment()?);
@@ -192,7 +195,7 @@ fn group_label(candidate: &ImportCandidate) -> GroupLabel {
 fn label(candidate: &ImportCandidate, import: &hir::ModPath) -> String {
     match candidate {
         ImportCandidate::Path(candidate) => {
-            if candidate.unresolved_qualifier.is_some() {
+            if !matches!(candidate.qualifier, Qualifier::Absent) {
                 format!("Qualify with `{}`", &import)
             } else {
                 format!("Qualify as `{}`", &import)
diff --git a/crates/ide_completion/src/completions/flyimport.rs b/crates/ide_completion/src/completions/flyimport.rs
index 16384551c..64b60bbdd 100644
--- a/crates/ide_completion/src/completions/flyimport.rs
+++ b/crates/ide_completion/src/completions/flyimport.rs
@@ -775,18 +775,92 @@ fn main() {
     }
 
     #[test]
-    fn unresolved_qualifiers() {
+    fn unresolved_qualifier() {
+        check_edit(
+            "Item",
+            r#"
+mod foo {
+    pub mod bar {
+        pub mod baz {
+            pub struct Item;
+        }
+    }
+}
+
+fn main() {
+    bar::baz::Ite$0
+}
+"#,
+            r#"
+use foo::bar;
+
+mod foo {
+    pub mod bar {
+        pub mod baz {
+            pub struct Item;
+        }
+    }
+}
+
+fn main() {
+    bar::baz::Item
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn unresolved_assoc_item_container() {
+        check_edit(
+            "Item",
+            r#"
+mod foo {
+    pub struct Item;
+
+    impl Item {
+        pub const TEST_ASSOC: usize = 3;
+    }
+}
+
+fn main() {
+    Item::TEST_A$0;
+}
+"#,
+            r#"
+use foo::Item;
+
+mod foo {
+    pub struct Item;
+
+    impl Item {
+        pub const TEST_ASSOC: usize = 3;
+    }
+}
+
+fn main() {
+    Item::TEST_ASSOC
+}
+"#,
+        );
+    }
+
+    #[test]
+    fn unresolved_assoc_item_container_with_path() {
         check_edit(
             "Item",
             r#"
 mod foo {
     pub mod bar {
         pub struct Item;
+
+        impl Item {
+            pub const TEST_ASSOC: usize = 3;
+        }
     }
 }
 
 fn main() {
-    bar::Ite$0
+    bar::Item::TEST_A$0;
 }
 "#,
             r#"
@@ -795,11 +869,15 @@ use foo::bar;
 mod foo {
     pub mod bar {
         pub struct Item;
+
+        impl Item {
+            pub const TEST_ASSOC: usize = 3;
+        }
     }
 }
 
 fn main() {
-    bar::Item
+    bar::Item::TEST_ASSOC
 }
 "#,
         );
diff --git a/crates/ide_db/src/helpers/import_assets.rs b/crates/ide_db/src/helpers/import_assets.rs
index 976dc92fb..dc3b92a64 100644
--- a/crates/ide_db/src/helpers/import_assets.rs
+++ b/crates/ide_db/src/helpers/import_assets.rs
@@ -1,8 +1,8 @@
 //! Look up accessible paths for items.
 use either::Either;
 use hir::{
-    AsAssocItem, AssocItem, Crate, ItemInNs, MacroDef, ModPath, Module, ModuleDef, PrefixKind,
-    Semantics,
+    AsAssocItem, AssocItem, Crate, ItemInNs, MacroDef, ModPath, Module, ModuleDef, Name,
+    PrefixKind, Semantics,
 };
 use rustc_hash::FxHashSet;
 use syntax::{ast, AstNode};
@@ -34,11 +34,17 @@ pub struct TraitImportCandidate {
 
 #[derive(Debug)]
 pub struct PathImportCandidate {
-    pub unresolved_qualifier: Option<ast::Path>,
+    pub qualifier: Qualifier,
     pub name: NameToImport,
 }
 
 #[derive(Debug)]
+pub enum Qualifier {
+    Absent,
+    FirstSegmentUnresolved(ast::PathSegment, ast::Path),
+}
+
+#[derive(Debug)]
 pub enum NameToImport {
     Exact(String),
     Fuzzy(String),
@@ -162,8 +168,9 @@ impl ImportAssets {
                 let (assoc_item_search, limit) = if self.import_candidate.is_trait_candidate() {
                     (AssocItemSearch::AssocItemsOnly, None)
                 } else {
-                    (AssocItemSearch::Exclude, Some(DEFAULT_QUERY_SEARCH_LIMIT))
+                    (AssocItemSearch::Include, Some(DEFAULT_QUERY_SEARCH_LIMIT))
                 };
+
                 imports_locator::find_similar_imports(
                     sema,
                     current_crate,
@@ -192,17 +199,16 @@ impl ImportAssets {
         let db = sema.db;
 
         match &self.import_candidate {
-            ImportCandidate::Path(path_candidate) => Box::new(path_applicable_items(
-                sema,
-                path_candidate,
-                unfiltered_defs
-                    .into_iter()
-                    .map(|def| def.either(ItemInNs::from, ItemInNs::from))
-                    .filter_map(move |item_to_search| {
-                        get_mod_path(db, item_to_search, &self.module_with_candidate, prefixed)
-                            .zip(Some(item_to_search))
-                    }),
-            )),
+            ImportCandidate::Path(path_candidate) => Box::new(
+                path_applicable_items(
+                    db,
+                    path_candidate,
+                    &self.module_with_candidate,
+                    prefixed,
+                    unfiltered_defs,
+                )
+                .into_iter(),
+            ),
             ImportCandidate::TraitAssocItem(trait_candidate) => Box::new(
                 trait_applicable_defs(db, current_crate, trait_candidate, true, unfiltered_defs)
                     .into_iter()
@@ -224,27 +230,110 @@ impl ImportAssets {
 }
 
 fn path_applicable_items<'a>(
-    sema: &'a Semantics<RootDatabase>,
-    path_candidate: &PathImportCandidate,
-    unfiltered_defs: impl Iterator<Item = (ModPath, ItemInNs)> + 'a,
-) -> Box<dyn Iterator<Item = (ModPath, ItemInNs)> + 'a> {
-    let unresolved_qualifier = match &path_candidate.unresolved_qualifier {
-        Some(qualifier) => qualifier,
-        None => {
-            return Box::new(unfiltered_defs);
+    db: &'a RootDatabase,
+    path_candidate: &'a PathImportCandidate,
+    module_with_candidate: &hir::Module,
+    prefixed: Option<hir::PrefixKind>,
+    unfiltered_defs: impl Iterator<Item = Either<ModuleDef, MacroDef>> + 'a,
+) -> FxHashSet<(ModPath, ItemInNs)> {
+    let applicable_items = unfiltered_defs
+        .filter_map(|def| {
+            let (assoc_original, candidate) = match def {
+                Either::Left(module_def) => match module_def.as_assoc_item(db) {
+                    Some(assoc_item) => match assoc_item.container(db) {
+                        hir::AssocItemContainer::Trait(trait_) => {
+                            (Some(module_def), ItemInNs::from(ModuleDef::from(trait_)))
+                        }
+                        hir::AssocItemContainer::Impl(impl_) => (
+                            Some(module_def),
+                            ItemInNs::from(ModuleDef::from(impl_.target_ty(db).as_adt()?)),
+                        ),
+                    },
+                    None => (None, ItemInNs::from(module_def)),
+                },
+                Either::Right(macro_def) => (None, ItemInNs::from(macro_def)),
+            };
+            Some((assoc_original, candidate))
+        })
+        .filter_map(|(assoc_original, candidate)| {
+            get_mod_path(db, candidate, module_with_candidate, prefixed)
+                .zip(Some((assoc_original, candidate)))
+        });
+
+    let (unresolved_first_segment, unresolved_qualifier) = match &path_candidate.qualifier {
+        Qualifier::Absent => {
+            return applicable_items
+                .map(|(candidate_path, (_, candidate))| (candidate_path, candidate))
+                .collect();
         }
+        Qualifier::FirstSegmentUnresolved(first_segment, qualifier) => (first_segment, qualifier),
     };
 
-    let qualifier_string = unresolved_qualifier.to_string();
-    Box::new(unfiltered_defs.filter(move |(candidate_path, _)| {
-        let mut candidate_qualifier = candidate_path.clone();
-        candidate_qualifier.pop_segment();
+    // TODO kb need to remove turbofish from the qualifier, maybe use the segments instead?
+    let unresolved_qualifier_string = unresolved_qualifier.to_string();
+    let unresolved_first_segment_string = unresolved_first_segment.to_string();
+
+    applicable_items
+        .filter(|(candidate_path, _)| {
+            let candidate_path_string = candidate_path.to_string();
+            candidate_path_string.contains(&unresolved_qualifier_string)
+                && candidate_path_string.contains(&unresolved_first_segment_string)
+        })
+        // TODO kb need to adjust the return type: I get the results rendered rather badly
+        .filter_map(|(candidate_path, (assoc_original, candidate))| {
+            if let Some(assoc_original) = assoc_original {
+                if item_name(db, candidate)?.to_string() == unresolved_first_segment_string {
+                    return Some((candidate_path, ItemInNs::from(assoc_original)));
+                }
+            }
+
+            let matching_module =
+                module_with_matching_name(db, &unresolved_first_segment_string, candidate)?;
+            let path = get_mod_path(
+                db,
+                ItemInNs::from(ModuleDef::from(matching_module)),
+                module_with_candidate,
+                prefixed,
+            )?;
+            Some((path, candidate))
+        })
+        .collect()
+}
+
+fn item_name(db: &RootDatabase, item: ItemInNs) -> Option<Name> {
+    match item {
+        ItemInNs::Types(module_def_id) => ModuleDef::from(module_def_id).name(db),
+        ItemInNs::Values(module_def_id) => ModuleDef::from(module_def_id).name(db),
+        ItemInNs::Macros(macro_def_id) => MacroDef::from(macro_def_id).name(db),
+    }
+}
 
-        // TODO kb
-        // * take 1st segment of `unresolved_qualifier` and return it instead of the original `ItemInNs`
-        // * Update `ModPath`: pop until 1st segment of `unresolved_qualifier` reached (do not rely on name comparison, nested mod names can repeat)
-        candidate_qualifier.to_string().ends_with(&qualifier_string)
-    }))
+fn item_module(db: &RootDatabase, item: ItemInNs) -> Option<Module> {
+    match item {
+        ItemInNs::Types(module_def_id) => ModuleDef::from(module_def_id).module(db),
+        ItemInNs::Values(module_def_id) => ModuleDef::from(module_def_id).module(db),
+        ItemInNs::Macros(macro_def_id) => MacroDef::from(macro_def_id).module(db),
+    }
+}
+
+fn module_with_matching_name(
+    db: &RootDatabase,
+    unresolved_first_segment_string: &str,
+    candidate: ItemInNs,
+) -> Option<Module> {
+    let mut current_module = item_module(db, candidate);
+    while let Some(module) = current_module {
+        match module.name(db) {
+            Some(module_name) => {
+                if module_name.to_string().as_str() == unresolved_first_segment_string {
+                    return Some(module);
+                }
+            }
+            None => {}
+        }
+        current_module = module.parent(db);
+    }
+    None
 }
 
 fn trait_applicable_defs<'a>(
@@ -367,10 +456,20 @@ fn path_import_candidate(
 ) -> Option<ImportCandidate> {
     Some(match qualifier {
         Some(qualifier) => match sema.resolve_path(&qualifier) {
-            None => ImportCandidate::Path(PathImportCandidate {
-                unresolved_qualifier: Some(qualifier),
-                name,
-            }),
+            None => {
+                let qualifier_start =
+                    qualifier.syntax().descendants().find_map(ast::PathSegment::cast)?;
+                let qualifier_start_path =
+                    qualifier_start.syntax().ancestors().find_map(ast::Path::cast)?;
+                if sema.resolve_path(&qualifier_start_path).is_none() {
+                    ImportCandidate::Path(PathImportCandidate {
+                        qualifier: Qualifier::FirstSegmentUnresolved(qualifier_start, qualifier),
+                        name,
+                    })
+                } else {
+                    return None;
+                }
+            }
             Some(hir::PathResolution::Def(hir::ModuleDef::Adt(assoc_item_path))) => {
                 ImportCandidate::TraitAssocItem(TraitImportCandidate {
                     receiver_ty: assoc_item_path.ty(sema.db),
@@ -379,6 +478,6 @@ fn path_import_candidate(
             }
             Some(_) => return None,
         },
-        None => ImportCandidate::Path(PathImportCandidate { unresolved_qualifier: None, name }),
+        None => ImportCandidate::Path(PathImportCandidate { qualifier: Qualifier::Absent, name }),
     })
 }
diff --git a/crates/ide_db/src/imports_locator.rs b/crates/ide_db/src/imports_locator.rs
index 502e8281a..480cbf1ea 100644
--- a/crates/ide_db/src/imports_locator.rs
+++ b/crates/ide_db/src/imports_locator.rs
@@ -40,6 +40,7 @@ pub fn find_exact_imports<'a>(
     ))
 }
 
+#[derive(Debug)]
 pub enum AssocItemSearch {
     Include,
     Exclude,