internal: Improve `find_path` and extern prelude handling

6 files changed, 103 insertions, 93 deletions

diff --git a/crates/hir_def/src/find_path.rs b/crates/hir_def/src/find_path.rs
index e1e5ded52a6..bb89b8cff44 100644
--- a/crates/hir_def/src/find_path.rs
+++ b/crates/hir_def/src/find_path.rs
@@ -18,8 +18,7 @@ use crate::{
 /// *from where* you're referring to the item, hence the `from` parameter.
 pub fn find_path(db: &dyn DefDatabase, item: ItemInNs, from: ModuleId) -> Option<ModPath> {
     let _p = profile::span("find_path");
-    let mut visited_modules = FxHashSet::default();
-    find_path_inner(db, item, from, MAX_PATH_LEN, None, &mut visited_modules)
+    find_path_inner(db, item, from, None)
 }
 
 pub fn find_path_prefixed(
@@ -29,8 +28,7 @@ pub fn find_path_prefixed(
     prefix_kind: PrefixKind,
 ) -> Option<ModPath> {
     let _p = profile::span("find_path_prefixed");
-    let mut visited_modules = FxHashSet::default();
-    find_path_inner(db, item, from, MAX_PATH_LEN, Some(prefix_kind), &mut visited_modules)
+    find_path_inner(db, item, from, Some(prefix_kind))
 }
 
 const MAX_PATH_LEN: usize = 15;
@@ -56,12 +54,12 @@ impl ModPathExt for ModPath {
 fn check_self_super(def_map: &DefMap, item: ItemInNs, from: ModuleId) -> Option<ModPath> {
     if item == ItemInNs::Types(from.into()) {
         // - if the item is the module we're in, use `self`
-        Some(ModPath::from_segments(PathKind::Super(0), Vec::new()))
+        Some(ModPath::from_segments(PathKind::Super(0), None))
     } else if let Some(parent_id) = def_map[from.local_id].parent {
         // - if the item is the parent module, use `super` (this is not used recursively, since `super::super` is ugly)
         let parent_id = def_map.module_id(parent_id);
         if item == ItemInNs::Types(ModuleDefId::ModuleId(parent_id)) {
-            Some(ModPath::from_segments(PathKind::Super(1), Vec::new()))
+            Some(ModPath::from_segments(PathKind::Super(1), None))
         } else {
             None
         }
@@ -97,12 +95,25 @@ impl PrefixKind {
         self == &PrefixKind::ByCrate
     }
 }
-
 /// Attempts to find a path to refer to the given `item` visible from the `from` ModuleId
 fn find_path_inner(
     db: &dyn DefDatabase,
     item: ItemInNs,
     from: ModuleId,
+    prefixed: Option<PrefixKind>,
+) -> Option<ModPath> {
+    // FIXME: Do fast path for std/core libs?
+
+    let mut visited_modules = FxHashSet::default();
+    let def_map = from.def_map(db);
+    find_path_inner_(db, &def_map, from, item, MAX_PATH_LEN, prefixed, &mut visited_modules)
+}
+
+fn find_path_inner_(
+    db: &dyn DefDatabase,
+    def_map: &DefMap,
+    from: ModuleId,
+    item: ItemInNs,
     max_len: usize,
     mut prefixed: Option<PrefixKind>,
     visited_modules: &mut FxHashSet<ModuleId>,
@@ -114,19 +125,24 @@ fn find_path_inner(
     // Base cases:
 
     // - if the item is already in scope, return the name under which it is
-    let def_map = from.def_map(db);
     let scope_name = def_map.with_ancestor_maps(db, from.local_id, &mut |def_map, local_id| {
         def_map[local_id].scope.name_of(item).map(|(name, _)| name.clone())
     });
-    if prefixed.is_none() && scope_name.is_some() {
-        return scope_name
-            .map(|scope_name| ModPath::from_segments(PathKind::Plain, vec![scope_name]));
+    if prefixed.is_none() {
+        if let Some(scope_name) = scope_name {
+            return Some(ModPath::from_segments(PathKind::Plain, Some(scope_name)));
+        }
+    }
+
+    // - if the item is a builtin, it's in scope
+    if let ItemInNs::Types(ModuleDefId::BuiltinType(builtin)) = item {
+        return Some(ModPath::from_segments(PathKind::Plain, Some(builtin.as_name())));
     }
 
     // - if the item is the crate root, return `crate`
-    let root = def_map.crate_root(db);
-    if item == ItemInNs::Types(ModuleDefId::ModuleId(root)) {
-        return Some(ModPath::from_segments(PathKind::Crate, Vec::new()));
+    let crate_root = def_map.crate_root(db);
+    if item == ItemInNs::Types(ModuleDefId::ModuleId(crate_root)) {
+        return Some(ModPath::from_segments(PathKind::Crate, None));
     }
 
     if prefixed.filter(PrefixKind::is_absolute).is_none() {
@@ -136,25 +152,28 @@ fn find_path_inner(
     }
 
     // - if the item is the crate root of a dependency crate, return the name from the extern prelude
-    let root_def_map = root.def_map(db);
-    for (name, def_id) in root_def_map.extern_prelude() {
-        if item == ItemInNs::Types(*def_id) {
-            let name = scope_name.unwrap_or_else(|| name.clone());
-
-            let name_already_occupied_in_type_ns = def_map
-                .with_ancestor_maps(db, from.local_id, &mut |def_map, local_id| {
-                    def_map[local_id].scope.get(&name).take_types().filter(|&id| id != *def_id)
-                })
-                .is_some();
-            return Some(ModPath::from_segments(
-                if name_already_occupied_in_type_ns {
+    let root_def_map = crate_root.def_map(db);
+    if let ItemInNs::Types(ModuleDefId::ModuleId(item)) = item {
+        for (name, &def_id) in root_def_map.extern_prelude() {
+            if item == def_id {
+                let name = scope_name.unwrap_or_else(|| name.clone());
+
+                let name_already_occupied_in_type_ns = def_map
+                    .with_ancestor_maps(db, from.local_id, &mut |def_map, local_id| {
+                        def_map[local_id]
+                            .scope
+                            .type_(&name)
+                            .filter(|&(id, _)| id != ModuleDefId::ModuleId(def_id))
+                    })
+                    .is_some();
+                let kind = if name_already_occupied_in_type_ns {
                     cov_mark::hit!(ambiguous_crate_start);
                     PathKind::Abs
                 } else {
                     PathKind::Plain
-                },
-                vec![name],
-            ));
+                };
+                return Some(ModPath::from_segments(kind, Some(name)));
+            }
         }
     }
 
@@ -162,20 +181,14 @@ fn find_path_inner(
     if let Some(prelude_module) = root_def_map.prelude() {
         // Preludes in block DefMaps are ignored, only the crate DefMap is searched
         let prelude_def_map = prelude_module.def_map(db);
-        let prelude_scope: &crate::item_scope::ItemScope =
-            &prelude_def_map[prelude_module.local_id].scope;
+        let prelude_scope = &prelude_def_map[prelude_module.local_id].scope;
         if let Some((name, vis)) = prelude_scope.name_of(item) {
             if vis.is_visible_from(db, from) {
-                return Some(ModPath::from_segments(PathKind::Plain, vec![name.clone()]));
+                return Some(ModPath::from_segments(PathKind::Plain, Some(name.clone())));
             }
         }
     }
 
-    // - if the item is a builtin, it's in scope
-    if let ItemInNs::Types(ModuleDefId::BuiltinType(builtin)) = item {
-        return Some(ModPath::from_segments(PathKind::Plain, vec![builtin.as_name()]));
-    }
-
     // Recursive case:
     // - if the item is an enum variant, refer to it via the enum
     if let Some(ModuleDefId::EnumVariantId(variant)) = item.as_module_def_id() {
@@ -190,25 +203,24 @@ fn find_path_inner(
     }
 
     // - otherwise, look for modules containing (reexporting) it and import it from one of those
-
-    let crate_root = def_map.crate_root(db);
-    let crate_attrs = db.attrs(crate_root.into());
-    let prefer_no_std = crate_attrs.by_key("no_std").exists();
+    let prefer_no_std = db.attrs(crate_root.into()).by_key("no_std").exists();
     let mut best_path = None;
     let mut best_path_len = max_len;
 
     if item.krate(db) == Some(from.krate) {
         // Item was defined in the same crate that wants to import it. It cannot be found in any
         // dependency in this case.
+        // FIXME: this should have a fast path that doesn't look through the prelude again?
         for (module_id, name) in find_local_import_locations(db, item, from) {
             if !visited_modules.insert(module_id) {
                 cov_mark::hit!(recursive_imports);
                 continue;
             }
-            if let Some(mut path) = find_path_inner(
+            if let Some(mut path) = find_path_inner_(
                 db,
-                ItemInNs::Types(ModuleDefId::ModuleId(module_id)),
+                def_map,
                 from,
+                ItemInNs::Types(ModuleDefId::ModuleId(module_id)),
                 best_path_len - 1,
                 prefixed,
                 visited_modules,
@@ -233,16 +245,18 @@ fn find_path_inner(
             let import_map = db.import_map(dep.crate_id);
             import_map.import_info_for(item).and_then(|info| {
                 // Determine best path for containing module and append last segment from `info`.
-                let mut path = find_path_inner(
+                // FIXME: we should guide this to look up the path locally, or from the same crate again?
+                let mut path = find_path_inner_(
                     db,
-                    ItemInNs::Types(ModuleDefId::ModuleId(info.container)),
+                    def_map,
                     from,
+                    ItemInNs::Types(ModuleDefId::ModuleId(info.container)),
                     best_path_len - 1,
                     prefixed,
                     visited_modules,
                 )?;
                 cov_mark::hit!(partially_imported);
-                path.push_segment(info.path.segments.last().unwrap().clone());
+                path.push_segment(info.path.segments.last()?.clone());
                 Some(path)
             })
         });
@@ -267,7 +281,7 @@ fn find_path_inner(
 
     match prefixed.map(PrefixKind::prefix) {
         Some(prefix) => best_path.or_else(|| {
-            scope_name.map(|scope_name| ModPath::from_segments(prefix, vec![scope_name]))
+            scope_name.map(|scope_name| ModPath::from_segments(prefix, Some(scope_name)))
         }),
         None => best_path,
     }
@@ -370,8 +384,8 @@ fn find_local_import_locations(
         }
 
         // Descend into all modules visible from `from`.
-        for (_, per_ns) in data.scope.entries() {
-            if let Some((ModuleDefId::ModuleId(module), vis)) = per_ns.take_types_vis() {
+        for (ty, vis) in data.scope.types() {
+            if let ModuleDefId::ModuleId(module) = ty {
                 if vis.is_visible_from(db, from) {
                     worklist.push(module);
                 }
@@ -415,15 +429,7 @@ mod tests {
             .take_types()
             .unwrap();
 
-        let mut visited_modules = FxHashSet::default();
-        let found_path = find_path_inner(
-            &db,
-            ItemInNs::Types(resolved),
-            module,
-            MAX_PATH_LEN,
-            prefix_kind,
-            &mut visited_modules,
-        );
+        let found_path = find_path_inner(&db, ItemInNs::Types(resolved), module, prefix_kind);
         assert_eq!(found_path, Some(mod_path), "{:?}", prefix_kind);
     }
 
diff --git a/crates/hir_def/src/item_scope.rs b/crates/hir_def/src/item_scope.rs
index f1330b34d1f..0096df2288d 100644
--- a/crates/hir_def/src/item_scope.rs
+++ b/crates/hir_def/src/item_scope.rs
@@ -119,6 +119,12 @@ impl ItemScope {
         self.values.values().copied()
     }
 
+    pub fn types(
+        &self,
+    ) -> impl Iterator<Item = (ModuleDefId, Visibility)> + ExactSizeIterator + '_ {
+        self.types.values().copied()
+    }
+
     pub fn unnamed_consts(&self) -> impl Iterator<Item = ConstId> + '_ {
         self.unnamed_consts.iter().copied()
     }
@@ -142,6 +148,10 @@ impl ItemScope {
         }
     }
 
+    pub(crate) fn type_(&self, name: &Name) -> Option<(ModuleDefId, Visibility)> {
+        self.types.get(name).copied()
+    }
+
     /// XXX: this is O(N) rather than O(1), try to not introduce new usages.
     pub(crate) fn name_of(&self, item: ItemInNs) -> Option<(&Name, Visibility)> {
         let (def, mut iter) = match item {
diff --git a/crates/hir_def/src/nameres.rs b/crates/hir_def/src/nameres.rs
index e9d3d976f9e..e923d883ecf 100644
--- a/crates/hir_def/src/nameres.rs
+++ b/crates/hir_def/src/nameres.rs
@@ -75,7 +75,7 @@ use crate::{
     path::ModPath,
     per_ns::PerNs,
     visibility::Visibility,
-    AstId, BlockId, BlockLoc, FunctionId, LocalModuleId, ModuleDefId, ModuleId, ProcMacroId,
+    AstId, BlockId, BlockLoc, FunctionId, LocalModuleId, ModuleId, ProcMacroId,
 };
 
 /// Contains the results of (early) name resolution.
@@ -98,7 +98,7 @@ pub struct DefMap {
     /// marked with the `prelude_import` attribute, or (in the normal case) from
     /// a dependency (`std` or `core`).
     prelude: Option<ModuleId>,
-    extern_prelude: FxHashMap<Name, ModuleDefId>,
+    extern_prelude: FxHashMap<Name, ModuleId>,
 
     /// Side table for resolving derive helpers.
     exported_derives: FxHashMap<MacroDefId, Box<[Name]>>,
@@ -318,7 +318,7 @@ impl DefMap {
         self.prelude
     }
 
-    pub(crate) fn extern_prelude(&self) -> impl Iterator<Item = (&Name, &ModuleDefId)> + '_ {
+    pub(crate) fn extern_prelude(&self) -> impl Iterator<Item = (&Name, &ModuleId)> + '_ {
         self.extern_prelude.iter()
     }
 
diff --git a/crates/hir_def/src/nameres/collector.rs b/crates/hir_def/src/nameres/collector.rs
index 8e6b5283f19..ea29f804373 100644
--- a/crates/hir_def/src/nameres/collector.rs
+++ b/crates/hir_def/src/nameres/collector.rs
@@ -70,7 +70,7 @@ pub(super) fn collect_defs(db: &dyn DefDatabase, mut def_map: DefMap, tree_id: T
         deps.insert(dep.as_name(), dep_root.into());
 
         if dep.is_prelude() && !tree_id.is_block() {
-            def_map.extern_prelude.insert(dep.as_name(), dep_root.into());
+            def_map.extern_prelude.insert(dep.as_name(), dep_root);
         }
     }
 
@@ -234,7 +234,7 @@ enum MacroDirectiveKind {
 struct DefCollector<'a> {
     db: &'a dyn DefDatabase,
     def_map: DefMap,
-    deps: FxHashMap<Name, ModuleDefId>,
+    deps: FxHashMap<Name, ModuleId>,
     glob_imports: FxHashMap<LocalModuleId, Vec<(LocalModuleId, Visibility)>>,
     unresolved_imports: Vec<ImportDirective>,
     resolved_imports: Vec<ImportDirective>,
@@ -687,9 +687,7 @@ impl DefCollector<'_> {
             self.def_map.edition,
         );
 
-        let res = self.resolve_extern_crate(&extern_crate.name);
-
-        if let Some(ModuleDefId::ModuleId(m)) = res.take_types() {
+        if let Some(m) = self.resolve_extern_crate(&extern_crate.name) {
             if m == self.def_map.module_id(current_module_id) {
                 cov_mark::hit!(ignore_macro_use_extern_crate_self);
                 return;
@@ -757,10 +755,11 @@ impl DefCollector<'_> {
 
             let res = self.resolve_extern_crate(name);
 
-            if res.is_none() {
-                PartialResolvedImport::Unresolved
-            } else {
-                PartialResolvedImport::Resolved(res)
+            match res {
+                Some(res) => {
+                    PartialResolvedImport::Resolved(PerNs::types(res.into(), Visibility::Public))
+                }
+                None => PartialResolvedImport::Unresolved,
             }
         } else {
             let res = self.def_map.resolve_path_fp_with_macro(
@@ -796,7 +795,7 @@ impl DefCollector<'_> {
         }
     }
 
-    fn resolve_extern_crate(&self, name: &Name) -> PerNs {
+    fn resolve_extern_crate(&self, name: &Name) -> Option<ModuleId> {
         if *name == name!(self) {
             cov_mark::hit!(extern_crate_self_as);
             let root = match self.def_map.block {
@@ -806,9 +805,9 @@ impl DefCollector<'_> {
                 }
                 None => self.def_map.module_id(self.def_map.root()),
             };
-            PerNs::types(root.into(), Visibility::Public)
+            Some(root)
         } else {
-            self.deps.get(name).map_or(PerNs::none(), |&it| PerNs::types(it, Visibility::Public))
+            self.deps.get(name).copied()
         }
     }
 
@@ -846,7 +845,8 @@ impl DefCollector<'_> {
 
                 // extern crates in the crate root are special-cased to insert entries into the extern prelude: rust-lang/rust#54658
                 if import.is_extern_crate && module_id == self.def_map.root {
-                    if let (Some(def), Some(name)) = (def.take_types(), name) {
+                    if let (Some(ModuleDefId::ModuleId(def)), Some(name)) = (def.take_types(), name)
+                    {
                         self.def_map.extern_prelude.insert(name.clone(), def);
                     }
                 }
diff --git a/crates/hir_def/src/nameres/path_resolution.rs b/crates/hir_def/src/nameres/path_resolution.rs
index 809c6e0289c..cc5fc0a52a1 100644
--- a/crates/hir_def/src/nameres/path_resolution.rs
+++ b/crates/hir_def/src/nameres/path_resolution.rs
@@ -20,7 +20,7 @@ use crate::{
     path::{ModPath, PathKind},
     per_ns::PerNs,
     visibility::{RawVisibility, Visibility},
-    AdtId, CrateId, EnumVariantId, LocalModuleId, ModuleDefId,
+    AdtId, CrateId, EnumVariantId, LocalModuleId, ModuleDefId, ModuleId,
 };
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -63,19 +63,11 @@ impl DefMap {
         &self,
         db: &dyn DefDatabase,
         name: &Name,
-    ) -> PerNs {
-        let arc;
-        let root = match self.block {
-            Some(_) => {
-                arc = self.crate_root(db).def_map(db);
-                &*arc
-            }
-            None => self,
-        };
-
-        root.extern_prelude
-            .get(name)
-            .map_or(PerNs::none(), |&it| PerNs::types(it, Visibility::Public))
+    ) -> Option<ModuleId> {
+        match self.block {
+            Some(_) => self.crate_root(db).def_map(db).extern_prelude.get(name).copied(),
+            None => self.extern_prelude.get(name).copied(),
+        }
     }
 
     pub(crate) fn resolve_visibility(
@@ -273,9 +265,9 @@ impl DefMap {
                     Some((_, segment)) => segment,
                     None => return ResolvePathResult::empty(ReachedFixedPoint::Yes),
                 };
-                if let Some(def) = self.extern_prelude.get(segment) {
+                if let Some(&def) = self.extern_prelude.get(segment) {
                     tracing::debug!("absolute path {:?} resolved to crate {:?}", path, def);
-                    PerNs::types(*def, Visibility::Public)
+                    PerNs::types(def.into(), Visibility::Public)
                 } else {
                     return ResolvePathResult::empty(ReachedFixedPoint::No); // extern crate declarations can add to the extern prelude
                 }
@@ -408,7 +400,7 @@ impl DefMap {
         let from_extern_prelude = self
             .extern_prelude
             .get(name)
-            .map_or(PerNs::none(), |&it| PerNs::types(it, Visibility::Public));
+            .map_or(PerNs::none(), |&it| PerNs::types(it.into(), Visibility::Public));
 
         let from_prelude = self.resolve_in_prelude(db, name);
 
@@ -429,7 +421,9 @@ impl DefMap {
             None => self,
         };
         let from_crate_root = crate_def_map[crate_def_map.root].scope.get(name);
-        let from_extern_prelude = self.resolve_name_in_extern_prelude(db, name);
+        let from_extern_prelude = self
+            .resolve_name_in_extern_prelude(db, name)
+            .map_or(PerNs::none(), |it| PerNs::types(it.into(), Visibility::Public));
 
         from_crate_root.or(from_extern_prelude)
     }
diff --git a/crates/hir_def/src/resolver.rs b/crates/hir_def/src/resolver.rs
index b8c88282d67..ab99627d182 100644
--- a/crates/hir_def/src/resolver.rs
+++ b/crates/hir_def/src/resolver.rs
@@ -498,7 +498,7 @@ impl Scope {
                     acc.add(name, ScopeDef::ModuleDef(ModuleDefId::MacroId(MacroId::from(mac))));
                 });
                 m.def_map.extern_prelude().for_each(|(name, &def)| {
-                    acc.add(name, ScopeDef::ModuleDef(def));
+                    acc.add(name, ScopeDef::ModuleDef(ModuleDefId::ModuleId(def)));
                 });
                 BUILTIN_SCOPE.iter().for_each(|(name, &def)| {
                     acc.add_per_ns(name, def);