Intern block locations and handle them in ModuleId

5 files changed, 104 insertions, 65 deletions

diff --git a/crates/hir_def/src/db.rs b/crates/hir_def/src/db.rs
index a87c80b8a..aef7e1f6c 100644
--- a/crates/hir_def/src/db.rs
+++ b/crates/hir_def/src/db.rs
@@ -2,9 +2,9 @@
 use std::sync::Arc;
 
 use base_db::{salsa, CrateId, SourceDatabase, Upcast};
-use hir_expand::{db::AstDatabase, AstId, HirFileId};
+use hir_expand::{db::AstDatabase, HirFileId};
 use la_arena::ArenaMap;
-use syntax::{ast, SmolStr};
+use syntax::SmolStr;
 
 use crate::{
     adt::{EnumData, StructData},
@@ -16,9 +16,10 @@ use crate::{
     item_tree::ItemTree,
     lang_item::{LangItemTarget, LangItems},
     nameres::DefMap,
-    AttrDefId, ConstId, ConstLoc, DefWithBodyId, EnumId, EnumLoc, FunctionId, FunctionLoc,
-    GenericDefId, ImplId, ImplLoc, LocalEnumVariantId, LocalFieldId, StaticId, StaticLoc, StructId,
-    StructLoc, TraitId, TraitLoc, TypeAliasId, TypeAliasLoc, UnionId, UnionLoc, VariantId,
+    AttrDefId, BlockId, BlockLoc, ConstId, ConstLoc, DefWithBodyId, EnumId, EnumLoc, FunctionId,
+    FunctionLoc, GenericDefId, ImplId, ImplLoc, LocalEnumVariantId, LocalFieldId, StaticId,
+    StaticLoc, StructId, StructLoc, TraitId, TraitLoc, TypeAliasId, TypeAliasLoc, UnionId,
+    UnionLoc, VariantId,
 };
 
 #[salsa::query_group(InternDatabaseStorage)]
@@ -41,6 +42,8 @@ pub trait InternDatabase: SourceDatabase {
     fn intern_type_alias(&self, loc: TypeAliasLoc) -> TypeAliasId;
     #[salsa::interned]
     fn intern_impl(&self, loc: ImplLoc) -> ImplId;
+    #[salsa::interned]
+    fn intern_block(&self, loc: BlockLoc) -> BlockId;
 }
 
 #[salsa::query_group(DefDatabaseStorage)]
@@ -56,7 +59,7 @@ pub trait DefDatabase: InternDatabase + AstDatabase + Upcast<dyn AstDatabase> {
     fn crate_def_map_query(&self, krate: CrateId) -> Arc<DefMap>;
 
     #[salsa::invoke(DefMap::block_def_map_query)]
-    fn block_def_map(&self, krate: CrateId, block: AstId<ast::BlockExpr>) -> Arc<DefMap>;
+    fn block_def_map(&self, block: BlockId) -> Arc<DefMap>;
 
     #[salsa::invoke(StructData::struct_data_query)]
     fn struct_data(&self, id: StructId) -> Arc<StructData>;
diff --git a/crates/hir_def/src/lib.rs b/crates/hir_def/src/lib.rs
index cf09ebd3f..42b50b5b7 100644
--- a/crates/hir_def/src/lib.rs
+++ b/crates/hir_def/src/lib.rs
@@ -74,12 +74,16 @@ use stdx::impl_from;
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct ModuleId {
     krate: CrateId,
+    block: Option<BlockId>,
     pub local_id: LocalModuleId,
 }
 
 impl ModuleId {
     pub fn def_map(&self, db: &dyn db::DefDatabase) -> Arc<DefMap> {
-        db.crate_def_map(self.krate)
+        match self.block {
+            Some(block) => db.block_def_map(block),
+            None => db.crate_def_map(self.krate),
+        }
     }
 
     pub fn krate(&self) -> CrateId {
@@ -230,6 +234,15 @@ pub struct ImplId(salsa::InternId);
 type ImplLoc = ItemLoc<Impl>;
 impl_intern!(ImplId, ImplLoc, intern_impl, lookup_intern_impl);
 
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Ord, PartialOrd)]
+pub struct BlockId(salsa::InternId);
+#[derive(Debug, Hash, PartialEq, Eq, Clone)]
+pub struct BlockLoc {
+    ast_id: AstId<ast::BlockExpr>,
+    module: ModuleId,
+}
+impl_intern!(BlockId, BlockLoc, intern_block, lookup_intern_block);
+
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
 pub struct TypeParamId {
     pub parent: GenericDefId,
diff --git a/crates/hir_def/src/nameres.rs b/crates/hir_def/src/nameres.rs
index 4fbbecb38..199771e9a 100644
--- a/crates/hir_def/src/nameres.rs
+++ b/crates/hir_def/src/nameres.rs
@@ -62,7 +62,7 @@ use la_arena::Arena;
 use profile::Count;
 use rustc_hash::FxHashMap;
 use stdx::format_to;
-use syntax::{ast, AstNode};
+use syntax::ast;
 
 use crate::{
     db::DefDatabase,
@@ -70,14 +70,14 @@ use crate::{
     nameres::{diagnostics::DefDiagnostic, path_resolution::ResolveMode},
     path::ModPath,
     per_ns::PerNs,
-    AstId, LocalModuleId, ModuleDefId, ModuleId,
+    AstId, BlockId, BlockLoc, LocalModuleId, ModuleDefId, ModuleId,
 };
 
 /// Contains all top-level defs from a macro-expanded crate
 #[derive(Debug, PartialEq, Eq)]
 pub struct DefMap {
     _c: Count<Self>,
-    parent: Option<Arc<DefMap>>,
+    block: Option<BlockInfo>,
     root: LocalModuleId,
     modules: Arena<ModuleData>,
     krate: CrateId,
@@ -91,6 +91,13 @@ pub struct DefMap {
     diagnostics: Vec<DefDiagnostic>,
 }
 
+#[derive(Debug, PartialEq, Eq)]
+struct BlockInfo {
+    block: BlockId,
+    parent: Arc<DefMap>,
+    parent_module: LocalModuleId,
+}
+
 impl std::ops::Index<LocalModuleId> for DefMap {
     type Output = ModuleData;
     fn index(&self, id: LocalModuleId) -> &ModuleData {
@@ -190,15 +197,12 @@ impl DefMap {
         Arc::new(def_map)
     }
 
-    pub(crate) fn block_def_map_query(
-        db: &dyn DefDatabase,
-        krate: CrateId,
-        block: AstId<ast::BlockExpr>,
-    ) -> Arc<DefMap> {
-        let item_tree = db.item_tree(block.file_id);
-        let block_items = item_tree.inner_items_of_block(block.value);
+    pub(crate) fn block_def_map_query(db: &dyn DefDatabase, block_id: BlockId) -> Arc<DefMap> {
+        let block: BlockLoc = db.lookup_intern_block(block_id);
+        let item_tree = db.item_tree(block.ast_id.file_id);
+        let block_items = item_tree.inner_items_of_block(block.ast_id.value);
 
-        let parent = parent_def_map(db, krate, block);
+        let parent = block.module.def_map(db);
 
         if block_items.is_empty() {
             // If there are no inner items, nothing new is brought into scope, so we can just return
@@ -206,10 +210,13 @@ impl DefMap {
             return parent;
         }
 
-        let mut def_map = DefMap::empty(krate, parent.edition);
-        def_map.parent = Some(parent);
+        let block_info =
+            BlockInfo { block: block_id, parent, parent_module: block.module.local_id };
+
+        let mut def_map = DefMap::empty(block.module.krate, block_info.parent.edition);
+        def_map.block = Some(block_info);
 
-        let def_map = collector::collect_defs(db, def_map, Some(block.value));
+        let def_map = collector::collect_defs(db, def_map, Some(block.ast_id.value));
         Arc::new(def_map)
     }
 
@@ -218,7 +225,7 @@ impl DefMap {
         let root = modules.alloc(ModuleData::default());
         DefMap {
             _c: Count::new(),
-            parent: None,
+            block: None,
             krate,
             edition,
             extern_prelude: FxHashMap::default(),
@@ -266,7 +273,8 @@ impl DefMap {
     }
 
     pub fn module_id(&self, local_id: LocalModuleId) -> ModuleId {
-        ModuleId { krate: self.krate, local_id }
+        let block = self.block.as_ref().map(|b| b.block);
+        ModuleId { krate: self.krate, local_id, block }
     }
 
     pub(crate) fn resolve_path(
@@ -286,9 +294,9 @@ impl DefMap {
     pub fn dump(&self) -> String {
         let mut buf = String::new();
         let mut current_map = self;
-        while let Some(parent) = &current_map.parent {
+        while let Some(block) = &current_map.block {
             go(&mut buf, current_map, "block scope", current_map.root);
-            current_map = &**parent;
+            current_map = &*block.parent;
         }
         go(&mut buf, current_map, "crate", current_map.root);
         return buf;
@@ -342,35 +350,6 @@ impl ModuleData {
     }
 }
 
-fn parent_def_map(
-    db: &dyn DefDatabase,
-    krate: CrateId,
-    block: AstId<ast::BlockExpr>,
-) -> Arc<DefMap> {
-    // FIXME: store this info in the item tree instead of reparsing here
-    let ast_id_map = db.ast_id_map(block.file_id);
-    let block_ptr = ast_id_map.get(block.value);
-    let root = match db.parse_or_expand(block.file_id) {
-        Some(it) => it,
-        None => {
-            return Arc::new(DefMap::empty(krate, Edition::Edition2018));
-        }
-    };
-    let ast = block_ptr.to_node(&root);
-
-    for ancestor in ast.syntax().ancestors().skip(1) {
-        if let Some(block_expr) = ast::BlockExpr::cast(ancestor) {
-            let ancestor_id = ast_id_map.ast_id(&block_expr);
-            let ast_id = InFile::new(block.file_id, ancestor_id);
-            let parent_map = db.block_def_map(krate, ast_id);
-            return parent_map;
-        }
-    }
-
-    // No enclosing block scope, so the parent is the crate-level DefMap.
-    db.crate_def_map(krate)
-}
-
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum ModuleSource {
     SourceFile(ast::SourceFile),
diff --git a/crates/hir_def/src/nameres/path_resolution.rs b/crates/hir_def/src/nameres/path_resolution.rs
index c1eded5f2..8ce127dde 100644
--- a/crates/hir_def/src/nameres/path_resolution.rs
+++ b/crates/hir_def/src/nameres/path_resolution.rs
@@ -131,8 +131,8 @@ impl DefMap {
             result.krate = result.krate.or(new.krate);
             result.segment_index = result.segment_index.min(new.segment_index);
 
-            match &current_map.parent {
-                Some(map) => current_map = map,
+            match &current_map.block {
+                Some(block) => current_map = &block.parent,
                 None => return result,
             }
         }
diff --git a/crates/hir_def/src/nameres/tests.rs b/crates/hir_def/src/nameres/tests.rs
index 73e3a4702..b36d0b59b 100644
--- a/crates/hir_def/src/nameres/tests.rs
+++ b/crates/hir_def/src/nameres/tests.rs
@@ -8,12 +8,12 @@ mod block;
 
 use std::sync::Arc;
 
-use base_db::{fixture::WithFixture, SourceDatabase};
+use base_db::{fixture::WithFixture, FilePosition, SourceDatabase};
 use expect_test::{expect, Expect};
-use hir_expand::db::AstDatabase;
+use syntax::AstNode;
 use test_utils::mark;
 
-use crate::{db::DefDatabase, nameres::*, test_db::TestDB};
+use crate::{db::DefDatabase, nameres::*, test_db::TestDB, Lookup};
 
 fn compute_crate_def_map(ra_fixture: &str) -> Arc<DefMap> {
     let db = TestDB::with_files(ra_fixture);
@@ -23,14 +23,58 @@ fn compute_crate_def_map(ra_fixture: &str) -> Arc<DefMap> {
 
 fn compute_block_def_map(ra_fixture: &str) -> Arc<DefMap> {
     let (db, position) = TestDB::with_position(ra_fixture);
+
+    // FIXME: perhaps we should make this use body lowering tests instead?
+
     let module = db.module_for_file(position.file_id);
-    let ast_map = db.ast_id_map(position.file_id.into());
-    let ast = db.parse(position.file_id);
-    let block: ast::BlockExpr =
-        syntax::algo::find_node_at_offset(&ast.syntax_node(), position.offset).unwrap();
-    let block_id = ast_map.ast_id(&block);
+    let mut def_map = db.crate_def_map(module.krate);
+    while let Some(new_def_map) = descend_def_map_at_position(&db, position, def_map.clone()) {
+        def_map = new_def_map;
+    }
+
+    // FIXME: select the right module, not the root
+
+    def_map
+}
+
+fn descend_def_map_at_position(
+    db: &dyn DefDatabase,
+    position: FilePosition,
+    def_map: Arc<DefMap>,
+) -> Option<Arc<DefMap>> {
+    for (local_id, module_data) in def_map.modules() {
+        let mod_def = module_data.origin.definition_source(db);
+        let ast_map = db.ast_id_map(mod_def.file_id);
+        let item_tree = db.item_tree(mod_def.file_id);
+        let root = db.parse_or_expand(mod_def.file_id).unwrap();
+        for item in module_data.scope.declarations() {
+            match item {
+                ModuleDefId::FunctionId(it) => {
+                    // Technically blocks can be inside any type (due to arrays and const generics),
+                    // and also in const/static initializers. For tests we only really care about
+                    // functions though.
+
+                    let ast = ast_map.get(item_tree[it.lookup(db).id.value].ast_id).to_node(&root);
+
+                    if ast.syntax().text_range().contains(position.offset) {
+                        // Cursor inside function, descend into its body's DefMap.
+                        // Note that we don't handle block *expressions* inside function bodies.
+                        let ast_map = db.ast_id_map(position.file_id.into());
+                        let ast_id = ast_map.ast_id(&ast.body().unwrap());
+                        let block = BlockLoc {
+                            ast_id: InFile::new(position.file_id.into(), ast_id),
+                            module: def_map.module_id(local_id),
+                        };
+                        let block_id = db.intern_block(block);
+                        return Some(db.block_def_map(block_id));
+                    }
+                }
+                _ => continue,
+            }
+        }
+    }
 
-    db.block_def_map(module.krate, InFile::new(position.file_id.into(), block_id))
+    None
 }
 
 fn check(ra_fixture: &str, expect: Expect) {