Prepare ImportMap for supportin multiple import paths per item

Veykril · Veykril · commit 2339ba4440bc · 2023-11-11T14:04:24.000+01:00
diff --git a/crates/hir-def/src/find_path.rs b/crates/hir-def/src/find_path.rs
@@ -367,18 +367,18 @@ fn calculate_best_path(
         // too (unless we can't name it at all). It could *also* be (re)exported by the same crate
         // that wants to import it here, but we always prefer to use the external path here.
 
-        let crate_graph = db.crate_graph();
-        let extern_paths = crate_graph[from.krate].dependencies.iter().filter_map(|dep| {
+        for dep in &db.crate_graph()[from.krate].dependencies {
             let import_map = db.import_map(dep.crate_id);
-            import_map.import_info_for(item).and_then(|info| {
+            let Some(import_info_for) = import_map.import_info_for(item) else { continue };
+            for info in import_info_for {
                 if info.is_doc_hidden {
                     // the item or import is `#[doc(hidden)]`, so skip it as it is in an external crate
-                    return None;
+                    continue;
                 }
 
                 // Determine best path for containing module and append last segment from `info`.
                 // FIXME: we should guide this to look up the path locally, or from the same crate again?
-                let (mut path, path_stability) = find_path_for_module(
+                let Some((mut path, path_stability)) = find_path_for_module(
                     db,
                     def_map,
                     visited_modules,
@@ -388,22 +388,23 @@ fn calculate_best_path(
                     max_len - 1,
                     prefixed,
                     prefer_no_std,
-                )?;
+                ) else {
+                    continue;
+                };
                 cov_mark::hit!(partially_imported);
                 path.push_segment(info.name.clone());
-                Some((
+
+                let path_with_stab = (
                     path,
                     zip_stability(path_stability, if info.is_unstable { Unstable } else { Stable }),
-                ))
-            })
-        });
+                );
 
-        for path in extern_paths {
-            let new_path = match best_path.take() {
-                Some(best_path) => select_best_path(best_path, path, prefer_no_std),
-                None => path,
-            };
-            update_best_path(&mut best_path, new_path);
+                let new_path_with_stab = match best_path.take() {
+                    Some(best_path) => select_best_path(best_path, path_with_stab, prefer_no_std),
+                    None => path_with_stab,
+                };
+                update_best_path(&mut best_path, new_path_with_stab);
+            }
         }
     }
     if let Some(module) = item.module(db) {
@@ -593,7 +594,7 @@ mod tests {
 
         let found_path =
             find_path_inner(&db, ItemInNs::Types(resolved), module, prefix_kind, false);
-        assert_eq!(found_path, Some(mod_path), "{prefix_kind:?}");
+        assert_eq!(found_path, Some(mod_path), "on kind: {prefix_kind:?}");
     }
 
     fn check_found_path(
diff --git a/crates/hir-def/src/import_map.rs b/crates/hir-def/src/import_map.rs
@@ -3,11 +3,12 @@
 use std::{collections::hash_map::Entry, fmt, hash::BuildHasherDefault};
 
 use base_db::CrateId;
-use fst::{self, Streamer};
+use fst::{self, raw::IndexedValue, Streamer};
 use hir_expand::name::Name;
 use indexmap::IndexMap;
 use itertools::Itertools;
 use rustc_hash::{FxHashMap, FxHashSet, FxHasher};
+use smallvec::{smallvec, SmallVec};
 use triomphe::Arc;
 
 use crate::{
@@ -20,31 +21,28 @@ use crate::{
 
 type FxIndexMap<K, V> = IndexMap<K, V, BuildHasherDefault<FxHasher>>;
 
-// FIXME: Support aliases: an item may be exported under multiple names, so `ImportInfo` should
-// have `Vec<(Name, ModuleId)>` instead of `(Name, ModuleId)`.
 /// Item import details stored in the `ImportMap`.
 #[derive(Debug, Clone, Eq, PartialEq)]
 pub struct ImportInfo {
     /// A name that can be used to import the item, relative to the crate's root.
     pub name: Name,
     /// The module containing this item.
     pub container: ModuleId,
-    /// Whether the import is a trait associated item or not.
-    pub is_trait_assoc_item: bool,
     /// Whether this item is annotated with `#[doc(hidden)]`.
     pub is_doc_hidden: bool,
     /// Whether this item is annotated with `#[unstable(..)]`.
     pub is_unstable: bool,
 }
 
+type ImportMapIndex = FxIndexMap<ItemInNs, (SmallVec<[ImportInfo; 1]>, IsTraitAssocItem)>;
+
 /// A map from publicly exported items to its name.
 ///
 /// Reexports of items are taken into account, ie. if something is exported under multiple
 /// names, the one with the shortest import path will be used.
 #[derive(Default)]
 pub struct ImportMap {
-    map: FxIndexMap<ItemInNs, ImportInfo>,
-
+    map: ImportMapIndex,
     /// List of keys stored in `map`, sorted lexicographically by their `ModPath`. Indexed by the
     /// values returned by running `fst`.
     ///
@@ -55,6 +53,12 @@ pub struct ImportMap {
     fst: fst::Map<Vec<u8>>,
 }
 
+#[derive(Copy, Clone, PartialEq, Eq)]
+enum IsTraitAssocItem {
+    Yes,
+    No,
+}
+
 impl ImportMap {
     pub(crate) fn import_map_query(db: &dyn DefDatabase, krate: CrateId) -> Arc<Self> {
         let _p = profile::span("import_map_query");
@@ -64,9 +68,13 @@ impl ImportMap {
         let mut importables: Vec<_> = map
             .iter()
             // We've only collected items, whose name cannot be tuple field.
-            .map(|(&item, info)| (item, info.name.as_str().unwrap().to_ascii_lowercase()))
+            .flat_map(|(&item, (info, _))| {
+                info.iter()
+                    .map(move |info| (item, info.name.as_str().unwrap().to_ascii_lowercase()))
+            })
             .collect();
         importables.sort_by(|(_, lhs_name), (_, rhs_name)| lhs_name.cmp(rhs_name));
+        importables.dedup();
 
         // Build the FST, taking care not to insert duplicate values.
         let mut builder = fst::MapBuilder::memory();
@@ -82,12 +90,12 @@ impl ImportMap {
         })
     }
 
-    pub fn import_info_for(&self, item: ItemInNs) -> Option<&ImportInfo> {
-        self.map.get(&item)
+    pub fn import_info_for(&self, item: ItemInNs) -> Option<&[ImportInfo]> {
+        self.map.get(&item).map(|(info, _)| &**info)
     }
 }
 
-fn collect_import_map(db: &dyn DefDatabase, krate: CrateId) -> FxIndexMap<ItemInNs, ImportInfo> {
+fn collect_import_map(db: &dyn DefDatabase, krate: CrateId) -> ImportMapIndex {
     let _p = profile::span("collect_import_map");
 
     let def_map = db.crate_def_map(krate);
@@ -140,7 +148,6 @@ fn collect_import_map(db: &dyn DefDatabase, krate: CrateId) -> FxIndexMap<ItemIn
                 let import_info = ImportInfo {
                     name: name.clone(),
                     container: module,
-                    is_trait_assoc_item: false,
                     is_doc_hidden,
                     is_unstable,
                 };
@@ -180,6 +187,8 @@ fn collect_import_map(db: &dyn DefDatabase, krate: CrateId) -> FxIndexMap<ItemIn
                                 depth < occ_depth
                             }
                         };
+                        // FIXME: Remove the overwrite rules as we can now record exports and
+                        // aliases for the same item
                         if !overwrite {
                             continue;
                         }
@@ -197,7 +206,7 @@ fn collect_import_map(db: &dyn DefDatabase, krate: CrateId) -> FxIndexMap<ItemIn
                     );
                 }
 
-                map.insert(item, import_info);
+                map.insert(item, (smallvec![import_info], IsTraitAssocItem::No));
 
                 // If we've just added a module, descend into it. We might traverse modules
                 // multiple times, but only if the module depth is smaller (else we `continue`
@@ -214,7 +223,7 @@ fn collect_import_map(db: &dyn DefDatabase, krate: CrateId) -> FxIndexMap<ItemIn
 
 fn collect_trait_assoc_items(
     db: &dyn DefDatabase,
-    map: &mut FxIndexMap<ItemInNs, ImportInfo>,
+    map: &mut ImportMapIndex,
     tr: TraitId,
     is_type_in_ns: bool,
     trait_import_info: &ImportInfo,
@@ -241,11 +250,10 @@ fn collect_trait_assoc_items(
         let assoc_item_info = ImportInfo {
             container: trait_import_info.container,
             name: assoc_item_name.clone(),
-            is_trait_assoc_item: true,
             is_doc_hidden: attrs.has_doc_hidden(),
             is_unstable: attrs.is_unstable(),
         };
-        map.insert(assoc_item, assoc_item_info);
+        map.insert(assoc_item, (smallvec![assoc_item_info], IsTraitAssocItem::Yes));
     }
 }
 
@@ -349,19 +357,24 @@ impl Query {
         Self { case_sensitive: true, ..self }
     }
 
+    fn matches_assoc_mode(&self, is_trait_assoc_item: IsTraitAssocItem) -> bool {
+        match (is_trait_assoc_item, self.assoc_mode) {
+            (IsTraitAssocItem::Yes, AssocSearchMode::Exclude)
+            | (IsTraitAssocItem::No, AssocSearchMode::AssocItemsOnly) => false,
+            _ => true,
+        }
+    }
+
+    /// Checks whether the import map entry matches the query.
     fn import_matches(
         &self,
         db: &dyn DefDatabase,
         import: &ImportInfo,
         enforce_lowercase: bool,
     ) -> bool {
         let _p = profile::span("import_map::Query::import_matches");
-        match (import.is_trait_assoc_item, self.assoc_mode) {
-            (true, AssocSearchMode::Exclude) => return false,
-            (false, AssocSearchMode::AssocItemsOnly) => return false,
-            _ => {}
-        }
 
+        // FIXME: Can we get rid of the alloc here?
         let mut input = import.name.display(db.upcast()).to_string();
         let case_insensitive = enforce_lowercase || !self.case_sensitive;
         if case_insensitive {
@@ -411,32 +424,55 @@ pub fn search_dependencies(
 
     let mut res = FxHashSet::default();
     while let Some((_, indexed_values)) = stream.next() {
-        for indexed_value in indexed_values {
-            let import_map = &import_maps[indexed_value.index];
-            let importables = &import_map.importables[indexed_value.value as usize..];
+        for &IndexedValue { index, value } in indexed_values {
+            let import_map = &import_maps[index];
+            let [importable, importables @ ..] = &import_map.importables[value as usize..] else {
+                continue;
+            };
 
-            let common_importable_data = &import_map.map[&importables[0]];
-            if !query.import_matches(db, common_importable_data, true) {
+            let &(ref importable_data, is_trait_assoc_item) = &import_map.map[importable];
+            if !query.matches_assoc_mode(is_trait_assoc_item) {
                 continue;
             }
 
-            // Name shared by the importable items in this group.
-            let common_importable_name =
-                common_importable_data.name.to_smol_str().to_ascii_lowercase();
-            // Add the items from this name group. Those are all subsequent items in
-            // `importables` whose name match `common_importable_name`.
-            let iter = importables
-                .iter()
-                .copied()
-                .take_while(|item| {
-                    common_importable_name
-                        == import_map.map[item].name.to_smol_str().to_ascii_lowercase()
-                })
-                .filter(|item| {
-                    !query.case_sensitive // we've already checked the common importables name case-insensitively
-                        || query.import_matches(db, &import_map.map[item], false)
-                });
-            res.extend(iter);
+            // FIXME: We probably need to account for other possible matches in this alias group?
+            let Some(common_importable_data) =
+                importable_data.iter().find(|&info| query.import_matches(db, info, true))
+            else {
+                continue;
+            };
+            res.insert(*importable);
+
+            if !importables.is_empty() {
+                // FIXME: so many allocs...
+                // Name shared by the importable items in this group.
+                let common_importable_name =
+                    common_importable_data.name.to_smol_str().to_ascii_lowercase();
+                // Add the items from this name group. Those are all subsequent items in
+                // `importables` whose name match `common_importable_name`.
+                let iter = importables
+                    .iter()
+                    .copied()
+                    .take_while(|item| {
+                        let &(ref import_infos, assoc_mode) = &import_map.map[item];
+                        query.matches_assoc_mode(assoc_mode)
+                            && import_infos.iter().any(|info| {
+                                info.name.to_smol_str().to_ascii_lowercase()
+                                    == common_importable_name
+                            })
+                    })
+                    .filter(|item| {
+                        !query.case_sensitive || {
+                            // we've already checked the common importables name case-insensitively
+                            let &(ref import_infos, assoc_mode) = &import_map.map[item];
+                            query.matches_assoc_mode(assoc_mode)
+                                && import_infos
+                                    .iter()
+                                    .any(|info| query.import_matches(db, info, false))
+                        }
+                    });
+                res.extend(iter);
+            }
 
             if res.len() >= query.limit {
                 return res;
@@ -461,6 +497,7 @@ mod tests {
             let mut importable_paths: Vec<_> = self
                 .map
                 .iter()
+                .flat_map(|(item, (info, _))| info.iter().map(move |info| (item, info)))
                 .map(|(item, info)| {
                     let path = render_path(db, info);
                     let ns = match item {
@@ -499,7 +536,7 @@ mod tests {
                 let (path, mark) = match assoc_item_path(&db, &dependency_imports, dependency) {
                     Some(assoc_item_path) => (assoc_item_path, "a"),
                     None => (
-                        render_path(&db, dependency_imports.import_info_for(dependency)?),
+                        render_path(&db, &dependency_imports.import_info_for(dependency)?[0]),
                         match dependency {
                             ItemInNs::Types(ModuleDefId::FunctionId(_))
                             | ItemInNs::Values(ModuleDefId::FunctionId(_)) => "f",
@@ -547,7 +584,12 @@ mod tests {
             .items
             .iter()
             .find(|(_, assoc_item_id)| &dependency_assoc_item_id == assoc_item_id)?;
-        Some(format!("{}::{}", render_path(db, trait_info), assoc_item_name.display(db.upcast())))
+        // FIXME: This should check all import infos, not just the first
+        Some(format!(
+            "{}::{}",
+            render_path(db, &trait_info[0]),
+            assoc_item_name.display(db.upcast())
+        ))
     }
 
     fn check(ra_fixture: &str, expect: Expect) {
