Use ItemTree for modules in attrs_query

3 files changed, 32 insertions, 26 deletions

diff --git a/crates/hir-def/src/attr.rs b/crates/hir-def/src/attr.rs
index 8222cc97f0e..8a6b6f3effd 100644
--- a/crates/hir-def/src/attr.rs
+++ b/crates/hir-def/src/attr.rs
@@ -374,30 +374,23 @@ impl AttrsWithOwner {
                 let mod_data = &def_map[module.local_id];
 
                 match mod_data.origin {
-                    // FIXME: We should be able to leverage the item tree instead of parsing declaration here
-                    // but we don't save the module's item tree id anywhere
-                    ModuleOrigin::File { definition, declaration, .. } => {
-                        let value = declaration.to_node(db.upcast());
-                        let it = InFile { file_id: declaration.file_id, value };
-                        let raw_attrs = RawAttrs::from_attrs_owner(
-                            db,
-                            it.as_ref().map(|it| it as &dyn ast::HasAttrs),
-                        );
+                    ModuleOrigin::File { definition, declaration_tree_id, .. } => {
+                        let decl_attrs = declaration_tree_id
+                            .item_tree(db)
+                            .raw_attrs(AttrOwner::ModItem(declaration_tree_id.value.into()))
+                            .clone();
                         let tree = db.file_item_tree(definition.into());
-                        raw_attrs.merge(tree.raw_attrs(AttrOwner::TopLevel).clone())
+                        let def_attrs = tree.raw_attrs(AttrOwner::TopLevel).clone();
+                        decl_attrs.merge(def_attrs)
                     }
                     ModuleOrigin::CrateRoot { definition } => {
                         let tree = db.file_item_tree(definition.into());
                         tree.raw_attrs(AttrOwner::TopLevel).clone()
                     }
-                    // FIXME: We should be able to leverage the item tree instead of parsing here
-                    // but we don't save the module's item tree id anywhere
-                    ModuleOrigin::Inline { definition } => RawAttrs::from_attrs_owner(
-                        db,
-                        InFile::new(definition.file_id, definition.to_node(db.upcast()))
-                            .as_ref()
-                            .map(|it| it as &dyn ast::HasAttrs),
-                    ),
+                    ModuleOrigin::Inline { definition_tree_id, .. } => definition_tree_id
+                        .item_tree(db)
+                        .raw_attrs(AttrOwner::ModItem(definition_tree_id.value.into()))
+                        .clone(),
                     ModuleOrigin::BlockExpr { block } => RawAttrs::from_attrs_owner(
                         db,
                         InFile::new(block.file_id, block.to_node(db.upcast()))
diff --git a/crates/hir-def/src/nameres.rs b/crates/hir-def/src/nameres.rs
index a32513cf65b..c67046dfdab 100644
--- a/crates/hir-def/src/nameres.rs
+++ b/crates/hir-def/src/nameres.rs
@@ -70,7 +70,7 @@ use syntax::{ast, SmolStr};
 use crate::{
     db::DefDatabase,
     item_scope::{BuiltinShadowMode, ItemScope},
-    item_tree::TreeId,
+    item_tree::{ItemTreeId, Mod, TreeId},
     nameres::{diagnostics::DefDiagnostic, path_resolution::ResolveMode},
     path::ModPath,
     per_ns::PerNs,
@@ -141,9 +141,11 @@ pub enum ModuleOrigin {
     File {
         is_mod_rs: bool,
         declaration: AstId<ast::Module>,
+        declaration_tree_id: ItemTreeId<Mod>,
         definition: FileId,
     },
     Inline {
+        definition_tree_id: ItemTreeId<Mod>,
         definition: AstId<ast::Module>,
     },
     /// Pseudo-module introduced by a block scope (contains only inner items).
@@ -186,7 +188,7 @@ impl ModuleOrigin {
                 let sf = db.parse(file_id).tree();
                 InFile::new(file_id.into(), ModuleSource::SourceFile(sf))
             }
-            ModuleOrigin::Inline { definition } => InFile::new(
+            ModuleOrigin::Inline { definition, .. } => InFile::new(
                 definition.file_id,
                 ModuleSource::Module(definition.to_node(db.upcast())),
             ),
diff --git a/crates/hir-def/src/nameres/collector.rs b/crates/hir-def/src/nameres/collector.rs
index 317e21538b0..743d130d020 100644
--- a/crates/hir-def/src/nameres/collector.rs
+++ b/crates/hir-def/src/nameres/collector.rs
@@ -1525,7 +1525,7 @@ impl ModCollector<'_, '_> {
             };
 
             match item {
-                ModItem::Mod(m) => self.collect_module(&self.item_tree[m], &attrs),
+                ModItem::Mod(m) => self.collect_module(m, &attrs),
                 ModItem::Import(import_id) => {
                     let imports = Import::from_use(
                         db,
@@ -1700,9 +1700,10 @@ impl ModCollector<'_, '_> {
         }
     }
 
-    fn collect_module(&mut self, module: &Mod, attrs: &Attrs) {
+    fn collect_module(&mut self, module_id: FileItemTreeId<Mod>, attrs: &Attrs) {
         let path_attr = attrs.by_key("path").string_value();
         let is_macro_use = attrs.by_key("macro_use").exists();
+        let module = &self.item_tree[module_id];
         match &module.kind {
             // inline module, just recurse
             ModKind::Inline { items } => {
@@ -1711,6 +1712,7 @@ impl ModCollector<'_, '_> {
                     AstId::new(self.file_id(), module.ast_id),
                     None,
                     &self.item_tree[module.visibility],
+                    module_id,
                 );
 
                 if let Some(mod_dir) = self.mod_dir.descend_into_definition(&module.name, path_attr)
@@ -1748,6 +1750,7 @@ impl ModCollector<'_, '_> {
                                 ast_id,
                                 Some((file_id, is_mod_rs)),
                                 &self.item_tree[module.visibility],
+                                module_id,
                             );
                             ModCollector {
                                 def_collector: self.def_collector,
@@ -1774,6 +1777,7 @@ impl ModCollector<'_, '_> {
                             ast_id,
                             None,
                             &self.item_tree[module.visibility],
+                            module_id,
                         );
                         self.def_collector.def_map.diagnostics.push(
                             DefDiagnostic::unresolved_module(self.module_id, ast_id, candidates),
@@ -1790,6 +1794,7 @@ impl ModCollector<'_, '_> {
         declaration: AstId<ast::Module>,
         definition: Option<(FileId, bool)>,
         visibility: &crate::visibility::RawVisibility,
+        mod_tree_id: FileItemTreeId<Mod>,
     ) -> LocalModuleId {
         let def_map = &mut self.def_collector.def_map;
         let vis = def_map
@@ -1797,10 +1802,16 @@ impl ModCollector<'_, '_> {
             .unwrap_or(Visibility::Public);
         let modules = &mut def_map.modules;
         let origin = match definition {
-            None => ModuleOrigin::Inline { definition: declaration },
-            Some((definition, is_mod_rs)) => {
-                ModuleOrigin::File { declaration, definition, is_mod_rs }
-            }
+            None => ModuleOrigin::Inline {
+                definition: declaration,
+                definition_tree_id: ItemTreeId::new(self.tree_id, mod_tree_id),
+            },
+            Some((definition, is_mod_rs)) => ModuleOrigin::File {
+                declaration,
+                definition,
+                is_mod_rs,
+                declaration_tree_id: ItemTreeId::new(self.tree_id, mod_tree_id),
+            },
         };
 
         let res = modules.alloc(ModuleData::new(origin, vis));