5 files changed, 67 insertions, 36 deletions

diff --git a/crates/hir_def/src/body/tests.rs b/crates/hir_def/src/body/tests.rs
index 404603360..a92134ba7 100644
--- a/crates/hir_def/src/body/tests.rs
+++ b/crates/hir_def/src/body/tests.rs
@@ -34,7 +34,7 @@ fn check_diagnostics(ra_fixture: &str) {
     db.check_diagnostics();
 }
 
-fn block_def_map_at(ra_fixture: &str) -> Arc<DefMap> {
+fn block_def_map_at(ra_fixture: &str) -> String {
     let (db, position) = crate::test_db::TestDB::with_position(ra_fixture);
 
     let krate = db.crate_graph().iter().next().unwrap();
@@ -51,7 +51,7 @@ fn block_def_map_at(ra_fixture: &str) -> Arc<DefMap> {
                 block = new_block;
             }
             None => {
-                return def_map;
+                return def_map.dump(&db);
             }
         }
     }
@@ -138,8 +138,7 @@ fn block_at_pos(db: &dyn DefDatabase, def_map: &DefMap, position: FilePosition)
 }
 
 fn check_at(ra_fixture: &str, expect: Expect) {
-    let def_map = block_def_map_at(ra_fixture);
-    let actual = def_map.dump();
+    let actual = block_def_map_at(ra_fixture);
     expect.assert_eq(&actual);
 }
 
diff --git a/crates/hir_def/src/nameres.rs b/crates/hir_def/src/nameres.rs
index 68998c121..ad2e9bcac 100644
--- a/crates/hir_def/src/nameres.rs
+++ b/crates/hir_def/src/nameres.rs
@@ -100,11 +100,12 @@ pub struct DefMap {
 }
 
 /// For `DefMap`s computed for a block expression, this stores its location in the parent map.
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, PartialEq, Eq, Clone, Copy)]
 struct BlockInfo {
+    /// The `BlockId` this `DefMap` was created from.
     block: BlockId,
-    parent: Arc<DefMap>,
-    parent_module: LocalModuleId,
+    /// The containing module.
+    parent: ModuleId,
 }
 
 impl std::ops::Index<LocalModuleId> for DefMap {
@@ -211,17 +212,16 @@ impl DefMap {
         block_id: BlockId,
     ) -> Option<Arc<DefMap>> {
         let block: BlockLoc = db.lookup_intern_block(block_id);
-        let parent = block.module.def_map(db);
 
         let item_tree = db.item_tree(block.ast_id.file_id);
         if item_tree.inner_items_of_block(block.ast_id.value).is_empty() {
             return None;
         }
 
-        let block_info =
-            BlockInfo { block: block_id, parent, parent_module: block.module.local_id };
+        let block_info = BlockInfo { block: block_id, parent: block.module };
 
-        let mut def_map = DefMap::empty(block.module.krate, block_info.parent.edition);
+        let parent_map = block.module.def_map(db);
+        let mut def_map = DefMap::empty(block.module.krate, parent_map.edition);
         def_map.block = Some(block_info);
 
         let def_map = collector::collect_defs(db, def_map, Some(block.ast_id));
@@ -289,9 +289,15 @@ impl DefMap {
         ModuleId { krate: self.krate, local_id, block }
     }
 
-    pub(crate) fn crate_root(&self) -> ModuleId {
-        let (root_map, _) = self.ancestor_maps(self.root).last().unwrap();
-        root_map.module_id(root_map.root)
+    pub(crate) fn crate_root(&self, db: &dyn DefDatabase) -> ModuleId {
+        self.with_ancestor_maps(db, self.root, &mut |def_map, _module| {
+            if def_map.block.is_none() {
+                Some(def_map.module_id(def_map.root))
+            } else {
+                None
+            }
+        })
+        .expect("DefMap chain without root")
     }
 
     pub(crate) fn resolve_path(
@@ -306,26 +312,42 @@ impl DefMap {
         (res.resolved_def, res.segment_index)
     }
 
-    /// Iterates over the containing `DefMap`s, if `self` is a `DefMap` corresponding to a block
-    /// expression.
-    fn ancestor_maps(
+    /// Ascends the `DefMap` hierarchy and calls `f` with every `DefMap` and containing module.
+    ///
+    /// If `f` returns `Some(val)`, iteration is stopped and `Some(val)` is returned. If `f` returns
+    /// `None`, iteration continues.
+    fn with_ancestor_maps<T>(
         &self,
+        db: &dyn DefDatabase,
         local_mod: LocalModuleId,
-    ) -> impl Iterator<Item = (&DefMap, LocalModuleId)> {
-        std::iter::successors(Some((self, local_mod)), |(map, _)| {
-            map.block.as_ref().map(|block| (&*block.parent, block.parent_module))
-        })
+        f: &mut dyn FnMut(&DefMap, LocalModuleId) -> Option<T>,
+    ) -> Option<T> {
+        if let Some(it) = f(self, local_mod) {
+            return Some(it);
+        }
+        let mut block = self.block;
+        while let Some(block_info) = block {
+            let parent = block_info.parent.def_map(db);
+            if let Some(it) = f(&parent, block_info.parent.local_id) {
+                return Some(it);
+            }
+            block = parent.block;
+        }
+
+        None
     }
 
     // FIXME: this can use some more human-readable format (ideally, an IR
     // even), as this should be a great debugging aid.
-    pub fn dump(&self) -> String {
+    pub fn dump(&self, db: &dyn DefDatabase) -> String {
         let mut buf = String::new();
+        let mut arc;
         let mut current_map = self;
         while let Some(block) = &current_map.block {
             go(&mut buf, current_map, "block scope", current_map.root);
             buf.push('\n');
-            current_map = &*block.parent;
+            arc = block.parent.def_map(db);
+            current_map = &*arc;
         }
         go(&mut buf, current_map, "crate", current_map.root);
         return buf;
diff --git a/crates/hir_def/src/nameres/collector.rs b/crates/hir_def/src/nameres/collector.rs
index 6e86cc4a7..f904a97de 100644
--- a/crates/hir_def/src/nameres/collector.rs
+++ b/crates/hir_def/src/nameres/collector.rs
@@ -1449,10 +1449,11 @@ impl ModCollector<'_, '_> {
         if let Some(macro_call_id) =
             ast_id.as_call_id(self.def_collector.db, self.def_collector.def_map.krate, |path| {
                 path.as_ident().and_then(|name| {
-                    self.def_collector
-                        .def_map
-                        .ancestor_maps(self.module_id)
-                        .find_map(|(map, module)| map[module].scope.get_legacy_macro(&name))
+                    self.def_collector.def_map.with_ancestor_maps(
+                        self.def_collector.db,
+                        self.module_id,
+                        &mut |map, module| map[module].scope.get_legacy_macro(&name),
+                    )
                 })
             })
         {
diff --git a/crates/hir_def/src/nameres/path_resolution.rs b/crates/hir_def/src/nameres/path_resolution.rs
index ecf75c777..2a3ac5d7b 100644
--- a/crates/hir_def/src/nameres/path_resolution.rs
+++ b/crates/hir_def/src/nameres/path_resolution.rs
@@ -110,6 +110,7 @@ impl DefMap {
         let mut result = ResolvePathResult::empty(ReachedFixedPoint::No);
         result.segment_index = Some(usize::max_value());
 
+        let mut arc;
         let mut current_map = self;
         loop {
             let new = current_map.resolve_path_fp_with_macro_single(
@@ -131,8 +132,9 @@ impl DefMap {
 
             match &current_map.block {
                 Some(block) => {
-                    current_map = &block.parent;
-                    original_module = block.parent_module;
+                    original_module = block.parent.local_id;
+                    arc = block.parent.def_map(db);
+                    current_map = &*arc;
                 }
                 None => return result,
             }
@@ -152,7 +154,7 @@ impl DefMap {
             PathKind::DollarCrate(krate) => {
                 if krate == self.krate {
                     mark::hit!(macro_dollar_crate_self);
-                    PerNs::types(self.crate_root().into(), Visibility::Public)
+                    PerNs::types(self.crate_root(db).into(), Visibility::Public)
                 } else {
                     let def_map = db.crate_def_map(krate);
                     let module = def_map.module_id(def_map.root);
@@ -160,7 +162,7 @@ impl DefMap {
                     PerNs::types(module.into(), Visibility::Public)
                 }
             }
-            PathKind::Crate => PerNs::types(self.crate_root().into(), Visibility::Public),
+            PathKind::Crate => PerNs::types(self.crate_root(db).into(), Visibility::Public),
             // plain import or absolute path in 2015: crate-relative with
             // fallback to extern prelude (with the simplification in
             // rust-lang/rust#57745)
@@ -206,10 +208,10 @@ impl DefMap {
                                     segments: path.segments.clone(),
                                 };
                                 log::debug!("`super` path: {} -> {} in parent map", path, new_path);
-                                return block.parent.resolve_path_fp_with_macro(
+                                return block.parent.def_map(db).resolve_path_fp_with_macro(
                                     db,
                                     mode,
-                                    block.parent_module,
+                                    block.parent.local_id,
                                     &new_path,
                                     shadow,
                                 );
diff --git a/crates/hir_def/src/nameres/tests.rs b/crates/hir_def/src/nameres/tests.rs
index 723481c36..bd3e2701b 100644
--- a/crates/hir_def/src/nameres/tests.rs
+++ b/crates/hir_def/src/nameres/tests.rs
@@ -11,7 +11,9 @@ use base_db::{fixture::WithFixture, SourceDatabase};
 use expect_test::{expect, Expect};
 use test_utils::mark;
 
-use crate::{db::DefDatabase, nameres::*, test_db::TestDB};
+use crate::{db::DefDatabase, test_db::TestDB};
+
+use super::DefMap;
 
 fn compute_crate_def_map(ra_fixture: &str) -> Arc<DefMap> {
     let db = TestDB::with_files(ra_fixture);
@@ -19,9 +21,14 @@ fn compute_crate_def_map(ra_fixture: &str) -> Arc<DefMap> {
     db.crate_def_map(krate)
 }
 
+fn render_crate_def_map(ra_fixture: &str) -> String {
+    let db = TestDB::with_files(ra_fixture);
+    let krate = db.crate_graph().iter().next().unwrap();
+    db.crate_def_map(krate).dump(&db)
+}
+
 fn check(ra_fixture: &str, expect: Expect) {
-    let def_map = compute_crate_def_map(ra_fixture);
-    let actual = def_map.dump();
+    let actual = render_crate_def_map(ra_fixture);
     expect.assert_eq(&actual);
 }