internal: Use itemtree for variant and field attribute collection

Author: Veykril

crates/hir-def/src/adt.rs

@@ -136,9 +136,9 @@ impl EnumData {
 
         let enum_ = &item_tree[loc.id.value];
         let mut variants = Arena::new();
-        for var_id in enum_.variants.clone() {
-            if item_tree.attrs(db, krate, var_id.into()).is_cfg_enabled(&cfg_options) {
-                let var = &item_tree[var_id];
+        for tree_id in enum_.variants.clone() {
+            if item_tree.attrs(db, krate, tree_id.into()).is_cfg_enabled(&cfg_options) {
+                let var = &item_tree[tree_id];
                 let var_data = lower_fields(
                     db,
                     krate,

crates/hir-def/src/attr.rs

@@ -7,7 +7,7 @@ use cfg::{CfgExpr, CfgOptions};
 use either::Either;
 use hir_expand::{hygiene::Hygiene, name::AsName, HirFileId, InFile};
 use itertools::Itertools;
-use la_arena::ArenaMap;
+use la_arena::{ArenaMap, Idx, RawIdx};
 use mbe::{syntax_node_to_token_tree, DelimiterKind, Punct};
 use smallvec::{smallvec, SmallVec};
 use syntax::{
@@ -19,12 +19,12 @@ use tt::Subtree;
 use crate::{
     db::DefDatabase,
     intern::Interned,
-    item_tree::{ItemTreeId, ItemTreeNode},
+    item_tree::{Fields, ItemTreeId, ItemTreeNode},
     nameres::ModuleSource,
     path::{ModPath, PathKind},
     src::{HasChildSource, HasSource},
-    AdtId, AttrDefId, EnumId, GenericParamId, HasModule, LocalEnumVariantId, LocalFieldId, Lookup,
-    MacroId, VariantId,
+    AdtId, AttrDefId, EnumId, GenericParamId, LocalEnumVariantId, LocalFieldId, Lookup, MacroId,
+    VariantId,
 };
 
 /// Holds documentation
@@ -201,15 +201,22 @@ impl Attrs {
         db: &dyn DefDatabase,
         e: EnumId,
     ) -> Arc<ArenaMap<LocalEnumVariantId, Attrs>> {
-        let krate = e.lookup(db).container.krate;
-        let src = e.child_source(db);
         let mut res = ArenaMap::default();
 
-        for (id, var) in src.value.iter() {
-            let attrs = RawAttrs::from_attrs_owner(db, src.with_value(var as &dyn ast::HasAttrs))
-                .filter(db, krate);
-
-            res.insert(id, attrs)
+        let loc = e.lookup(db);
+        let krate = loc.container.krate;
+        let item_tree = loc.id.item_tree(db);
+        let enum_ = &item_tree[loc.id.value];
+        let crate_graph = db.crate_graph();
+        let cfg_options = &crate_graph[krate].cfg_options;
+
+        let mut idx = 0;
+        for variant in enum_.variants.clone() {
+            let attrs = item_tree.attrs(db, krate, variant.into());
+            if attrs.is_cfg_enabled(cfg_options) {
+                res.insert(Idx::from_raw(RawIdx::from(idx)), attrs);
+                idx += 1;
+            }
         }
 
         Arc::new(res)
@@ -219,18 +226,63 @@ impl Attrs {
         db: &dyn DefDatabase,
         v: VariantId,
     ) -> Arc<ArenaMap<LocalFieldId, Attrs>> {
-        let krate = v.module(db).krate;
-        let src = v.child_source(db);
         let mut res = ArenaMap::default();
 
-        for (id, fld) in src.value.iter() {
-            let owner: &dyn HasAttrs = match fld {
-                Either::Left(tuple) => tuple,
-                Either::Right(record) => record,
-            };
-            let attrs = RawAttrs::from_attrs_owner(db, src.with_value(owner)).filter(db, krate);
+        let crate_graph = db.crate_graph();
+        let (fields, item_tree, krate) = match v {
+            VariantId::EnumVariantId(it) => {
+                let e = it.parent;
+                let loc = e.lookup(db);
+                let krate = loc.container.krate;
+                let item_tree = loc.id.item_tree(db);
+                let enum_ = &item_tree[loc.id.value];
+
+                let cfg_options = &crate_graph[krate].cfg_options;
+                let variant = 'tri: loop {
+                    let mut idx = 0;
+                    for variant in enum_.variants.clone() {
+                        let attrs = item_tree.attrs(db, krate, variant.into());
+                        if attrs.is_cfg_enabled(cfg_options) {
+                            if it.local_id == Idx::from_raw(RawIdx::from(idx)) {
+                                break 'tri variant;
+                            }
+                            idx += 1;
+                        }
+                    }
+                    return Arc::new(res);
+                };
+                (item_tree[variant].fields.clone(), item_tree, krate)
+            }
+            VariantId::StructId(it) => {
+                let loc = it.lookup(db);
+                let krate = loc.container.krate;
+                let item_tree = loc.id.item_tree(db);
+                let struct_ = &item_tree[loc.id.value];
+                (struct_.fields.clone(), item_tree, krate)
+            }
+            VariantId::UnionId(it) => {
+                let loc = it.lookup(db);
+                let krate = loc.container.krate;
+                let item_tree = loc.id.item_tree(db);
+                let union_ = &item_tree[loc.id.value];
+                (union_.fields.clone(), item_tree, krate)
+            }
+        };
 
-            res.insert(id, attrs);
+        let fields = match fields {
+            Fields::Record(fields) | Fields::Tuple(fields) => fields,
+            Fields::Unit => return Arc::new(res),
+        };
+
+        let cfg_options = &crate_graph[krate].cfg_options;
+
+        let mut idx = 0;
+        for field in fields {
+            let attrs = item_tree.attrs(db, krate, field.into());
+            if attrs.is_cfg_enabled(cfg_options) {
+                res.insert(Idx::from_raw(RawIdx::from(idx)), attrs);
+                idx += 1;
+            }
         }
 
         Arc::new(res)

crates/hir-def/src/data.rs

@@ -1,22 +1,23 @@
 //! Contains basic data about various HIR declarations.
 
-use std::{mem, sync::Arc};
+use std::sync::Arc;
 
-use hir_expand::{name::Name, AstId, ExpandResult, HirFileId, InFile, MacroCallId, MacroDefKind};
+use hir_expand::{name::Name, AstId, ExpandResult, HirFileId, MacroCallId, MacroDefKind};
+use smallvec::SmallVec;
 use syntax::ast;
 
 use crate::{
     attr::Attrs,
     body::{Expander, Mark},
     db::DefDatabase,
     intern::Interned,
-    item_tree::{self, AssocItem, FnFlags, ItemTreeId, ModItem, Param, TreeId},
+    item_tree::{self, AssocItem, FnFlags, ItemTree, ItemTreeId, ModItem, Param, TreeId},
     nameres::{attr_resolution::ResolvedAttr, DefMap},
     type_ref::{TraitRef, TypeBound, TypeRef},
     visibility::RawVisibility,
     AssocItemId, AstIdWithPath, ConstId, ConstLoc, FunctionId, FunctionLoc, HasModule, ImplId,
-    Intern, ItemContainerId, Lookup, Macro2Id, MacroRulesId, ModuleId, ProcMacroId, StaticId,
-    TraitId, TypeAliasId, TypeAliasLoc,
+    Intern, ItemContainerId, ItemLoc, Lookup, Macro2Id, MacroRulesId, ModuleId, ProcMacroId,
+    StaticId, TraitId, TypeAliasId, TypeAliasLoc,
 };
 
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -209,35 +210,31 @@ pub struct TraitData {
 
 impl TraitData {
     pub(crate) fn trait_data_query(db: &dyn DefDatabase, tr: TraitId) -> Arc<TraitData> {
-        let tr_loc = tr.lookup(db);
-        let item_tree = tr_loc.id.item_tree(db);
-        let tr_def = &item_tree[tr_loc.id.value];
+        let tr_loc @ ItemLoc { container: module_id, id: tree_id } = tr.lookup(db);
+        let item_tree = tree_id.item_tree(db);
+        let tr_def = &item_tree[tree_id.value];
         let _cx = stdx::panic_context::enter(format!(
             "trait_data_query({:?} -> {:?} -> {:?})",
             tr, tr_loc, tr_def
         ));
         let name = tr_def.name.clone();
         let is_auto = tr_def.is_auto;
         let is_unsafe = tr_def.is_unsafe;
-        let module_id = tr_loc.container;
         let visibility = item_tree[tr_def.visibility].clone();
         let skip_array_during_method_dispatch = item_tree
-            .attrs(db, tr_loc.container.krate(), ModItem::from(tr_loc.id.value).into())
+            .attrs(db, module_id.krate(), ModItem::from(tree_id.value).into())
             .by_key("rustc_skip_array_during_method_dispatch")
             .exists();
 
-        let mut collector = AssocItemCollector::new(
-            db,
-            module_id,
-            tr_loc.id.file_id(),
-            ItemContainerId::TraitId(tr),
-        );
-        collector.collect(tr_loc.id.tree_id(), &tr_def.items);
+        let mut collector =
+            AssocItemCollector::new(db, module_id, tree_id.file_id(), ItemContainerId::TraitId(tr));
+        collector.collect(&item_tree, tree_id.tree_id(), &tr_def.items);
+        let (items, attribute_calls) = collector.finish();
 
         Arc::new(TraitData {
             name,
-            attribute_calls: collector.take_attr_calls(),
-            items: collector.items,
+            attribute_calls,
+            items,
             is_auto,
             is_unsafe,
             visibility,
@@ -284,25 +281,20 @@ pub struct ImplData {
 impl ImplData {
     pub(crate) fn impl_data_query(db: &dyn DefDatabase, id: ImplId) -> Arc<ImplData> {
         let _p = profile::span("impl_data_query");
-        let impl_loc = id.lookup(db);
+        let ItemLoc { container: module_id, id: tree_id } = id.lookup(db);
 
-        let item_tree = impl_loc.id.item_tree(db);
-        let impl_def = &item_tree[impl_loc.id.value];
+        let item_tree = tree_id.item_tree(db);
+        let impl_def = &item_tree[tree_id.value];
         let target_trait = impl_def.target_trait.clone();
         let self_ty = impl_def.self_ty.clone();
         let is_negative = impl_def.is_negative;
-        let module_id = impl_loc.container;
 
-        let mut collector = AssocItemCollector::new(
-            db,
-            module_id,
-            impl_loc.id.file_id(),
-            ItemContainerId::ImplId(id),
-        );
-        collector.collect(impl_loc.id.tree_id(), &impl_def.items);
+        let mut collector =
+            AssocItemCollector::new(db, module_id, tree_id.file_id(), ItemContainerId::ImplId(id));
+        collector.collect(&item_tree, tree_id.tree_id(), &impl_def.items);
 
-        let attribute_calls = collector.take_attr_calls();
-        let items = collector.items.into_iter().map(|(_, item)| item).collect();
+        let (items, attribute_calls) = collector.finish();
+        let items = items.into_iter().map(|(_, item)| item).collect();
 
         Arc::new(ImplData { target_trait, self_ty, items, is_negative, attribute_calls })
     }
@@ -463,18 +455,19 @@ impl<'a> AssocItemCollector<'a> {
         }
     }
 
-    fn take_attr_calls(&mut self) -> Option<Box<Vec<(AstId<ast::Item>, MacroCallId)>>> {
-        let attribute_calls = mem::take(&mut self.attr_calls);
-        if attribute_calls.is_empty() {
-            None
-        } else {
-            Some(Box::new(attribute_calls))
-        }
+    fn finish(
+        self,
+    ) -> (Vec<(Name, AssocItemId)>, Option<Box<Vec<(AstId<ast::Item>, MacroCallId)>>>) {
+        (
+            self.items,
+            if self.attr_calls.is_empty() { None } else { Some(Box::new(self.attr_calls)) },
+        )
     }
 
     // FIXME: proc-macro diagnostics
-    fn collect(&mut self, tree_id: TreeId, assoc_items: &[AssocItem]) {
-        let item_tree = tree_id.item_tree(self.db);
+    fn collect(&mut self, item_tree: &ItemTree, tree_id: TreeId, assoc_items: &[AssocItem]) {
+        let container = self.container;
+        self.items.reserve(assoc_items.len());
 
         'items: for &item in assoc_items {
             let attrs = item_tree.attrs(self.db, self.module_id.krate, ModItem::from(item).into());
@@ -509,9 +502,9 @@ impl<'a> AssocItemCollector<'a> {
                             continue 'attrs;
                         }
                     }
-                    match self.expander.enter_expand_id(self.db, call_id) {
-                        ExpandResult { value: Some((mark, mac)), .. } => {
-                            self.collect_macro_items(mark, mac);
+                    match self.expander.enter_expand_id::<ast::MacroItems>(self.db, call_id) {
+                        ExpandResult { value: Some((mark, _)), .. } => {
+                            self.collect_macro_items(mark);
                             continue 'items;
                         }
                         ExpandResult { .. } => {}
@@ -522,59 +515,57 @@ impl<'a> AssocItemCollector<'a> {
             match item {
                 AssocItem::Function(id) => {
                     let item = &item_tree[id];
+
                     let def =
-                        FunctionLoc { container: self.container, id: ItemTreeId::new(tree_id, id) }
-                            .intern(self.db);
+                        FunctionLoc { container, id: ItemTreeId::new(tree_id, id) }.intern(self.db);
                     self.items.push((item.name.clone(), def.into()));
                 }
                 AssocItem::Const(id) => {
                     let item = &item_tree[id];
+
                     let name = match item.name.clone() {
                         Some(name) => name,
                         None => continue,
                     };
                     let def =
-                        ConstLoc { container: self.container, id: ItemTreeId::new(tree_id, id) }
-                            .intern(self.db);
+                        ConstLoc { container, id: ItemTreeId::new(tree_id, id) }.intern(self.db);
                     self.items.push((name, def.into()));
                 }
                 AssocItem::TypeAlias(id) => {
                     let item = &item_tree[id];
-                    let def = TypeAliasLoc {
-                        container: self.container,
-                        id: ItemTreeId::new(tree_id, id),
-                    }
-                    .intern(self.db);
+
+                    let def = TypeAliasLoc { container, id: ItemTreeId::new(tree_id, id) }
+                        .intern(self.db);
                     self.items.push((item.name.clone(), def.into()));
                 }
                 AssocItem::MacroCall(call) => {
-                    let call = &item_tree[call];
-                    let ast_id_map = self.db.ast_id_map(self.expander.current_file_id());
                     if let Some(root) = self.db.parse_or_expand(self.expander.current_file_id()) {
+                        let call = &item_tree[call];
+
+                        let ast_id_map = self.db.ast_id_map(self.expander.current_file_id());
                         let call = ast_id_map.get(call.ast_id).to_node(&root);
                         let _cx = stdx::panic_context::enter(format!(
                             "collect_items MacroCall: {}",
                             call
                         ));
-                        let res = self.expander.enter_expand(self.db, call);
+                        let res = self.expander.enter_expand::<ast::MacroItems>(self.db, call);
 
-                        if let Ok(ExpandResult { value: Some((mark, mac)), .. }) = res {
-                            self.collect_macro_items(mark, mac);
+                        if let Ok(ExpandResult { value: Some((mark, _)), .. }) = res {
+                            self.collect_macro_items(mark);
                         }
                     }
                 }
             }
         }
     }
 
-    fn collect_macro_items(&mut self, mark: Mark, mac: ast::MacroItems) {
-        let src: InFile<ast::MacroItems> = self.expander.to_source(mac);
-        let tree_id = item_tree::TreeId::new(src.file_id, None);
+    fn collect_macro_items(&mut self, mark: Mark) {
+        let tree_id = item_tree::TreeId::new(self.expander.current_file_id(), None);
         let item_tree = tree_id.item_tree(self.db);
-        let iter: Vec<_> =
+        let iter: SmallVec<[_; 2]> =
             item_tree.top_level_items().iter().filter_map(ModItem::as_assoc_item).collect();
 
-        self.collect(tree_id, &iter);
+        self.collect(&item_tree, tree_id, &iter);
 
         self.expander.exit(self.db, mark);
     }