Fix `self` and `super` path resolution in block modules

3 files changed, 160 insertions, 47 deletions

diff --git a/crates/hir-def/src/body/tests/block.rs b/crates/hir-def/src/body/tests/block.rs
index 6e77744f215..4e015a7fbbb 100644
--- a/crates/hir-def/src/body/tests/block.rs
+++ b/crates/hir-def/src/body/tests/block.rs
@@ -134,6 +134,47 @@ struct Struct {}
 }
 
 #[test]
+fn super_imports_2() {
+    check_at(
+        r#"
+fn outer() {
+    mod m {
+        struct ResolveMe {}
+        fn middle() {
+            mod m2 {
+                fn inner() {
+                    use super::ResolveMe;
+                    $0
+                }
+            }
+        }
+    }
+}
+"#,
+        expect![[r#"
+            block scope
+            ResolveMe: t
+
+            block scope
+            m2: t
+
+            block scope::m2
+            inner: v
+
+            block scope
+            m: t
+
+            block scope::m
+            ResolveMe: t
+            middle: v
+
+            crate
+            outer: v
+        "#]],
+    );
+}
+
+#[test]
 fn nested_module_scoping() {
     check_block_scopes_at(
         r#"
@@ -156,6 +197,42 @@ fn f() {
 }
 
 #[test]
+fn self_imports() {
+    check_at(
+        r#"
+fn f() {
+    mod m {
+        struct ResolveMe {}
+        fn g() {
+            fn h() {
+                use self::ResolveMe;
+                $0
+            }
+        }
+    }
+}
+"#,
+        expect![[r#"
+            block scope
+            ResolveMe: t
+
+            block scope
+            h: v
+
+            block scope
+            m: t
+
+            block scope::m
+            ResolveMe: t
+            g: v
+
+            crate
+            f: v
+        "#]],
+    );
+}
+
+#[test]
 fn legacy_macro_items() {
     // Checks that legacy-scoped `macro_rules!` from parent namespaces are resolved and expanded
     // correctly.
diff --git a/crates/hir-def/src/nameres.rs b/crates/hir-def/src/nameres.rs
index e7a4355d258..86818ce26dd 100644
--- a/crates/hir-def/src/nameres.rs
+++ b/crates/hir-def/src/nameres.rs
@@ -196,6 +196,10 @@ impl BlockRelativeModuleId {
     fn into_module(self, krate: CrateId) -> ModuleId {
         ModuleId { krate, block: self.block, local_id: self.local_id }
     }
+
+    fn is_block_module(self) -> bool {
+        self.block.is_some() && self.local_id == DefMap::ROOT
+    }
 }
 
 impl std::ops::Index<LocalModuleId> for DefMap {
@@ -278,7 +282,9 @@ pub struct ModuleData {
     pub origin: ModuleOrigin,
     /// Declared visibility of this module.
     pub visibility: Visibility,
-    /// Always [`None`] for block modules
+    /// Parent module in the same `DefMap`.
+    ///
+    /// [`None`] for block modules because they are always its `DefMap`'s root.
     pub parent: Option<LocalModuleId>,
     pub children: FxHashMap<Name, LocalModuleId>,
     pub scope: ItemScope,
diff --git a/crates/hir-def/src/nameres/path_resolution.rs b/crates/hir-def/src/nameres/path_resolution.rs
index 5f6163175a7..38edd89879d 100644
--- a/crates/hir-def/src/nameres/path_resolution.rs
+++ b/crates/hir-def/src/nameres/path_resolution.rs
@@ -12,11 +12,12 @@
 
 use base_db::Edition;
 use hir_expand::name::Name;
+use triomphe::Arc;
 
 use crate::{
     db::DefDatabase,
     item_scope::BUILTIN_SCOPE,
-    nameres::{sub_namespace_match, BuiltinShadowMode, DefMap, MacroSubNs},
+    nameres::{sub_namespace_match, BlockInfo, BuiltinShadowMode, DefMap, MacroSubNs},
     path::{ModPath, PathKind},
     per_ns::PerNs,
     visibility::{RawVisibility, Visibility},
@@ -159,13 +160,15 @@ impl DefMap {
                 (None, new) => new,
             };
 
-            match &current_map.block {
-                Some(block) => {
+            match current_map.block {
+                Some(block) if original_module == Self::ROOT => {
+                    // Block modules "inherit" names from its parent module.
                     original_module = block.parent.local_id;
                     arc = block.parent.def_map(db, current_map.krate);
-                    current_map = &*arc;
+                    current_map = &arc;
                 }
-                None => return result,
+                // Proper (non-block) modules, including those in block `DefMap`s, don't.
+                _ => return result,
             }
         }
     }
@@ -241,51 +244,54 @@ impl DefMap {
                 )
             }
             PathKind::Super(lvl) => {
-                let mut module = original_module;
-                for i in 0..lvl {
-                    match self.modules[module].parent {
-                        Some(it) => module = it,
-                        None => match &self.block {
-                            Some(block) => {
-                                // Look up remaining path in parent `DefMap`
-                                let new_path = ModPath::from_segments(
-                                    PathKind::Super(lvl - i),
-                                    path.segments().to_vec(),
-                                );
-                                tracing::debug!(
-                                    "`super` path: {} -> {} in parent map",
-                                    path.display(db.upcast()),
-                                    new_path.display(db.upcast())
-                                );
-                                return block
-                                    .parent
-                                    .def_map(db, self.krate)
-                                    .resolve_path_fp_with_macro(
-                                        db,
-                                        mode,
-                                        block.parent.local_id,
-                                        &new_path,
-                                        shadow,
-                                        expected_macro_subns,
-                                    );
-                            }
-                            None => {
-                                tracing::debug!("super path in root module");
-                                return ResolvePathResult::empty(ReachedFixedPoint::Yes);
-                            }
-                        },
-                    }
+                let mut local_id = original_module;
+                let mut ext;
+                let mut def_map = self;
+
+                // Adjust `local_id` to `self`, i.e. the nearest non-block module.
+                if def_map.module_id(local_id).is_block_module() {
+                    (ext, local_id) = adjust_to_nearest_non_block_module(db, def_map, local_id);
+                    def_map = &ext;
                 }
 
-                // Resolve `self` to the containing crate-rooted module if we're a block
-                self.with_ancestor_maps(db, module, &mut |def_map, module| {
-                    if def_map.block.is_some() {
-                        None // keep ascending
+                // Go up the module tree but skip block modules as `super` always refers to the
+                // nearest non-block module.
+                for _ in 0..lvl {
+                    // Loop invariant: at the beginning of each loop, `local_id` must refer to a
+                    // non-block module.
+                    if let Some(parent) = def_map.modules[local_id].parent {
+                        local_id = parent;
+                        if def_map.module_id(local_id).is_block_module() {
+                            (ext, local_id) =
+                                adjust_to_nearest_non_block_module(db, def_map, local_id);
+                            def_map = &ext;
+                        }
                     } else {
-                        Some(PerNs::types(def_map.module_id(module).into(), Visibility::Public))
+                        stdx::always!(def_map.block.is_none());
+                        tracing::debug!("super path in root module");
+                        return ResolvePathResult::empty(ReachedFixedPoint::Yes);
                     }
-                })
-                .expect("block DefMap not rooted in crate DefMap")
+                }
+
+                let module = def_map.module_id(local_id);
+                stdx::never!(module.is_block_module());
+
+                if self.block != def_map.block {
+                    // If we have a different `DefMap` from `self` (the orignal `DefMap` we started
+                    // with), resolve the remaining path segments in that `DefMap`.
+                    let path =
+                        ModPath::from_segments(PathKind::Super(0), path.segments().iter().cloned());
+                    return def_map.resolve_path_fp_with_macro(
+                        db,
+                        mode,
+                        local_id,
+                        &path,
+                        shadow,
+                        expected_macro_subns,
+                    );
+                }
+
+                PerNs::types(module.into(), Visibility::Public)
             }
             PathKind::Abs => {
                 // 2018-style absolute path -- only extern prelude
@@ -508,3 +514,27 @@ impl DefMap {
         }
     }
 }
+
+/// Given a block module, returns its nearest non-block module and the `DefMap` it blongs to.
+fn adjust_to_nearest_non_block_module(
+    db: &dyn DefDatabase,
+    def_map: &DefMap,
+    mut local_id: LocalModuleId,
+) -> (Arc<DefMap>, LocalModuleId) {
+    // INVARIANT: `local_id` in `def_map` must be a block module.
+    stdx::always!(def_map.module_id(local_id).is_block_module());
+
+    let mut ext;
+    // This needs to be a local variable due to our mighty lifetime.
+    let mut def_map = def_map;
+    loop {
+        let BlockInfo { parent, .. } = def_map.block.expect("block module without parent module");
+
+        ext = parent.def_map(db, def_map.krate);
+        def_map = &ext;
+        local_id = parent.local_id;
+        if !parent.is_block_module() {
+            return (ext, local_id);
+        }
+    }
+}