Merge #5063

5063: Store field/variant attrs in ItemTree and use it for adt.rs queries r=jonas-schievink a=jonas-schievink



Co-authored-by: Jonas Schievink <[email protected]>

7 files changed, 162 insertions, 61 deletions

diff --git a/crates/ra_hir_def/src/adt.rs b/crates/ra_hir_def/src/adt.rs
index 2bc34d449..4994a2125 100644
--- a/crates/ra_hir_def/src/adt.rs
+++ b/crates/ra_hir_def/src/adt.rs
@@ -8,12 +8,12 @@ use hir_expand::{
     InFile,
 };
 use ra_arena::{map::ArenaMap, Arena};
-use ra_prof::profile;
 use ra_syntax::ast::{self, NameOwner, TypeAscriptionOwner, VisibilityOwner};
 
 use crate::{
     body::{CfgExpander, LowerCtx},
     db::DefDatabase,
+    item_tree::{Field, Fields, ItemTree},
     src::HasChildSource,
     src::HasSource,
     trace::Trace,
@@ -22,6 +22,7 @@ use crate::{
     EnumId, HasModule, LocalEnumVariantId, LocalFieldId, Lookup, ModuleId, StructId, UnionId,
     VariantId,
 };
+use ra_cfg::CfgOptions;
 
 /// Note that we use `StructData` for unions as well!
 #[derive(Debug, Clone, PartialEq, Eq)]
@@ -59,39 +60,48 @@ pub struct FieldData {
 
 impl StructData {
     pub(crate) fn struct_data_query(db: &dyn DefDatabase, id: StructId) -> Arc<StructData> {
-        let src = id.lookup(db).source(db);
+        let loc = id.lookup(db);
+        let item_tree = db.item_tree(loc.id.file_id);
+        let cfg_options = db.crate_graph()[loc.container.module(db).krate].cfg_options.clone();
 
-        let name = src.value.name().map_or_else(Name::missing, |n| n.as_name());
-        let variant_data =
-            VariantData::new(db, src.map(|s| s.kind()), id.lookup(db).container.module(db));
-        let variant_data = Arc::new(variant_data);
-        Arc::new(StructData { name, variant_data })
+        let strukt = &item_tree[loc.id.value];
+        let variant_data = lower_fields(&item_tree, &cfg_options, &strukt.fields);
+
+        Arc::new(StructData { name: strukt.name.clone(), variant_data: Arc::new(variant_data) })
     }
     pub(crate) fn union_data_query(db: &dyn DefDatabase, id: UnionId) -> Arc<StructData> {
-        let src = id.lookup(db).source(db);
-        let name = src.value.name().map_or_else(Name::missing, |n| n.as_name());
-        let variant_data = VariantData::new(
-            db,
-            src.map(|s| {
-                s.record_field_def_list()
-                    .map(ast::StructKind::Record)
-                    .unwrap_or(ast::StructKind::Unit)
-            }),
-            id.lookup(db).container.module(db),
-        );
-        let variant_data = Arc::new(variant_data);
-        Arc::new(StructData { name, variant_data })
+        let loc = id.lookup(db);
+        let item_tree = db.item_tree(loc.id.file_id);
+        let cfg_options = db.crate_graph()[loc.container.module(db).krate].cfg_options.clone();
+
+        let union = &item_tree[loc.id.value];
+        let variant_data = lower_fields(&item_tree, &cfg_options, &union.fields);
+
+        Arc::new(StructData { name: union.name.clone(), variant_data: Arc::new(variant_data) })
     }
 }
 
 impl EnumData {
     pub(crate) fn enum_data_query(db: &dyn DefDatabase, e: EnumId) -> Arc<EnumData> {
-        let _p = profile("enum_data_query");
-        let src = e.lookup(db).source(db);
-        let name = src.value.name().map_or_else(Name::missing, |n| n.as_name());
-        let mut trace = Trace::new_for_arena();
-        lower_enum(db, &mut trace, &src, e.lookup(db).container.module(db));
-        Arc::new(EnumData { name, variants: trace.into_arena() })
+        let loc = e.lookup(db);
+        let item_tree = db.item_tree(loc.id.file_id);
+        let cfg_options = db.crate_graph()[loc.container.module(db).krate].cfg_options.clone();
+
+        let enum_ = &item_tree[loc.id.value];
+        let mut variants = Arena::new();
+        for var_id in enum_.variants.clone() {
+            if item_tree.attrs(var_id.into()).is_cfg_enabled(&cfg_options) {
+                let var = &item_tree[var_id];
+                let var_data = lower_fields(&item_tree, &cfg_options, &var.fields);
+
+                variants.alloc(EnumVariantData {
+                    name: var.name.clone(),
+                    variant_data: Arc::new(var_data),
+                });
+            }
+        }
+
+        Arc::new(EnumData { name: enum_.name.clone(), variants })
     }
 
     pub fn variant(&self, name: &Name) -> Option<LocalEnumVariantId> {
@@ -251,3 +261,35 @@ fn lower_struct(
         ast::StructKind::Unit => StructKind::Unit,
     }
 }
+
+fn lower_fields(item_tree: &ItemTree, cfg_options: &CfgOptions, fields: &Fields) -> VariantData {
+    match fields {
+        Fields::Record(flds) => {
+            let mut arena = Arena::new();
+            for field_id in flds.clone() {
+                if item_tree.attrs(field_id.into()).is_cfg_enabled(cfg_options) {
+                    arena.alloc(lower_field(item_tree, &item_tree[field_id]));
+                }
+            }
+            VariantData::Record(arena)
+        }
+        Fields::Tuple(flds) => {
+            let mut arena = Arena::new();
+            for field_id in flds.clone() {
+                if item_tree.attrs(field_id.into()).is_cfg_enabled(cfg_options) {
+                    arena.alloc(lower_field(item_tree, &item_tree[field_id]));
+                }
+            }
+            VariantData::Tuple(arena)
+        }
+        Fields::Unit => VariantData::Unit,
+    }
+}
+
+fn lower_field(item_tree: &ItemTree, field: &Field) -> FieldData {
+    FieldData {
+        name: field.name.clone(),
+        type_ref: field.type_ref.clone(),
+        visibility: item_tree[field.visibility].clone(),
+    }
+}
diff --git a/crates/ra_hir_def/src/attr.rs b/crates/ra_hir_def/src/attr.rs
index 197737ffc..e228e2145 100644
--- a/crates/ra_hir_def/src/attr.rs
+++ b/crates/ra_hir_def/src/attr.rs
@@ -208,5 +208,5 @@ where
 fn attrs_from_item_tree<N: ItemTreeNode>(id: ItemTreeId<N>, db: &dyn DefDatabase) -> Attrs {
     let tree = db.item_tree(id.file_id);
     let mod_item = N::id_to_mod_item(id.value);
-    tree.attrs(mod_item).clone()
+    tree.attrs(mod_item.into()).clone()
 }
diff --git a/crates/ra_hir_def/src/data.rs b/crates/ra_hir_def/src/data.rs
index f9e5701db..282ade2a3 100644
--- a/crates/ra_hir_def/src/data.rs
+++ b/crates/ra_hir_def/src/data.rs
@@ -40,7 +40,7 @@ impl FunctionData {
             name: func.name.clone(),
             params: func.params.to_vec(),
             ret_type: func.ret_type.clone(),
-            attrs: item_tree.attrs(loc.id.value.into()).clone(),
+            attrs: item_tree.attrs(ModItem::from(loc.id.value).into()).clone(),
             has_self_param: func.has_self_param,
             is_unsafe: func.is_unsafe,
             visibility: item_tree[func.visibility].clone(),
@@ -224,7 +224,7 @@ fn collect_items(
         match item {
             AssocItem::Function(id) => {
                 let item = &item_tree[id];
-                let attrs = item_tree.attrs(id.into());
+                let attrs = item_tree.attrs(ModItem::from(id).into());
                 if !attrs.is_cfg_enabled(&cfg_options) {
                     continue;
                 }
diff --git a/crates/ra_hir_def/src/item_tree.rs b/crates/ra_hir_def/src/item_tree.rs
index d7bc64e6c..3e603bd55 100644
--- a/crates/ra_hir_def/src/item_tree.rs
+++ b/crates/ra_hir_def/src/item_tree.rs
@@ -5,6 +5,7 @@ mod lower;
 mod tests;
 
 use std::{
+    any::type_name,
     fmt::{self, Debug},
     hash::{Hash, Hasher},
     marker::PhantomData,
@@ -178,8 +179,8 @@ impl ItemTree {
         self.attrs.get(&AttrOwner::TopLevel).unwrap_or(&Attrs::EMPTY)
     }
 
-    pub fn attrs(&self, of: ModItem) -> &Attrs {
-        self.attrs.get(&AttrOwner::ModItem(of)).unwrap_or(&Attrs::EMPTY)
+    pub fn attrs(&self, of: AttrOwner) -> &Attrs {
+        self.attrs.get(&of).unwrap_or(&Attrs::EMPTY)
     }
 
     /// Returns the lowered inner items that `ast` corresponds to.
@@ -282,15 +283,32 @@ struct ItemTreeData {
 }
 
 #[derive(Debug, Eq, PartialEq, Hash)]
-enum AttrOwner {
+pub enum AttrOwner {
     /// Attributes on an item.
     ModItem(ModItem),
     /// Inner attributes of the source file.
     TopLevel,
+
+    Variant(Idx<Variant>),
+    Field(Idx<Field>),
     // FIXME: Store variant and field attrs, and stop reparsing them in `attrs_query`.
 }
 
-/// Trait implemented by all nodes in the item tree.
+macro_rules! from_attrs {
+    ( $( $var:ident($t:ty) ),+ ) => {
+        $(
+            impl From<$t> for AttrOwner {
+                fn from(t: $t) -> AttrOwner {
+                    AttrOwner::$var(t)
+                }
+            }
+        )+
+    };
+}
+
+from_attrs!(ModItem(ModItem), Variant(Idx<Variant>), Field(Idx<Field>));
+
+/// Trait implemented by all item nodes in the item tree.
 pub trait ItemTreeNode: Clone {
     type Source: AstNode + Into<ast::ModuleItem>;
 
@@ -523,7 +541,7 @@ pub struct Enum {
     pub name: Name,
     pub visibility: RawVisibilityId,
     pub generic_params: GenericParamsId,
-    pub variants: Range<Idx<Variant>>,
+    pub variants: IdRange<Variant>,
     pub ast_id: FileAstId<ast::EnumDef>,
 }
 
@@ -681,10 +699,48 @@ pub struct Variant {
     pub fields: Fields,
 }
 
+pub struct IdRange<T> {
+    range: Range<u32>,
+    _p: PhantomData<T>,
+}
+
+impl<T> IdRange<T> {
+    fn new(range: Range<Idx<T>>) -> Self {
+        Self { range: range.start.into_raw().into()..range.end.into_raw().into(), _p: PhantomData }
+    }
+}
+
+impl<T> Iterator for IdRange<T> {
+    type Item = Idx<T>;
+    fn next(&mut self) -> Option<Self::Item> {
+        self.range.next().map(|raw| Idx::from_raw(raw.into()))
+    }
+}
+
+impl<T> fmt::Debug for IdRange<T> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_tuple(&format!("IdRange::<{}>", type_name::<T>())).field(&self.range).finish()
+    }
+}
+
+impl<T> Clone for IdRange<T> {
+    fn clone(&self) -> Self {
+        Self { range: self.range.clone(), _p: PhantomData }
+    }
+}
+
+impl<T> PartialEq for IdRange<T> {
+    fn eq(&self, other: &Self) -> bool {
+        self.range == other.range
+    }
+}
+
+impl<T> Eq for IdRange<T> {}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub enum Fields {
-    Record(Range<Idx<Field>>),
-    Tuple(Range<Idx<Field>>),
+    Record(IdRange<Field>),
+    Tuple(IdRange<Field>),
     Unit,
 }
 
diff --git a/crates/ra_hir_def/src/item_tree/lower.rs b/crates/ra_hir_def/src/item_tree/lower.rs
index f10ad25f7..5149dd141 100644
--- a/crates/ra_hir_def/src/item_tree/lower.rs
+++ b/crates/ra_hir_def/src/item_tree/lower.rs
@@ -126,15 +126,15 @@ impl Ctx {
 
         if !attrs.is_empty() {
             for item in items.iter().flat_map(|items| &items.0) {
-                self.add_attrs(*item, attrs.clone());
+                self.add_attrs((*item).into(), attrs.clone());
             }
         }
 
         items
     }
 
-    fn add_attrs(&mut self, item: ModItem, attrs: Attrs) {
-        match self.tree.attrs.entry(AttrOwner::ModItem(item)) {
+    fn add_attrs(&mut self, item: AttrOwner, attrs: Attrs) {
+        match self.tree.attrs.entry(item) {
             Entry::Occupied(mut entry) => {
                 *entry.get_mut() = entry.get().merge(attrs);
             }
@@ -196,15 +196,16 @@ impl Ctx {
         }
     }
 
-    fn lower_record_fields(&mut self, fields: &ast::RecordFieldDefList) -> Range<Idx<Field>> {
+    fn lower_record_fields(&mut self, fields: &ast::RecordFieldDefList) -> IdRange<Field> {
         let start = self.next_field_idx();
         for field in fields.fields() {
             if let Some(data) = self.lower_record_field(&field) {
-                self.data().fields.alloc(data);
+                let idx = self.data().fields.alloc(data);
+                self.add_attrs(idx.into(), Attrs::new(&field, &self.hygiene));
             }
         }
         let end = self.next_field_idx();
-        start..end
+        IdRange::new(start..end)
     }
 
     fn lower_record_field(&mut self, field: &ast::RecordFieldDef) -> Option<Field> {
@@ -215,15 +216,16 @@ impl Ctx {
         Some(res)
     }
 
-    fn lower_tuple_fields(&mut self, fields: &ast::TupleFieldDefList) -> Range<Idx<Field>> {
+    fn lower_tuple_fields(&mut self, fields: &ast::TupleFieldDefList) -> IdRange<Field> {
         let start = self.next_field_idx();
         for (i, field) in fields.fields().enumerate() {
             if let Some(data) = self.lower_tuple_field(i, &field) {
-                self.data().fields.alloc(data);
+                let idx = self.data().fields.alloc(data);
+                self.add_attrs(idx.into(), Attrs::new(&field, &self.hygiene));
             }
         }
         let end = self.next_field_idx();
-        start..end
+        IdRange::new(start..end)
     }
 
     fn lower_tuple_field(&mut self, idx: usize, field: &ast::TupleFieldDef) -> Option<Field> {
@@ -242,7 +244,7 @@ impl Ctx {
             Some(record_field_def_list) => {
                 self.lower_fields(&StructKind::Record(record_field_def_list))
             }
-            None => Fields::Record(self.next_field_idx()..self.next_field_idx()),
+            None => Fields::Record(IdRange::new(self.next_field_idx()..self.next_field_idx())),
         };
         let ast_id = self.source_ast_id_map.ast_id(union);
         let res = Union { name, visibility, generic_params, fields, ast_id };
@@ -255,22 +257,23 @@ impl Ctx {
         let generic_params = self.lower_generic_params(GenericsOwner::Enum, enum_);
         let variants = match &enum_.variant_list() {
             Some(variant_list) => self.lower_variants(variant_list),
-            None => self.next_variant_idx()..self.next_variant_idx(),
+            None => IdRange::new(self.next_variant_idx()..self.next_variant_idx()),
         };
         let ast_id = self.source_ast_id_map.ast_id(enum_);
         let res = Enum { name, visibility, generic_params, variants, ast_id };
         Some(id(self.data().enums.alloc(res)))
     }
 
-    fn lower_variants(&mut self, variants: &ast::EnumVariantList) -> Range<Idx<Variant>> {
+    fn lower_variants(&mut self, variants: &ast::EnumVariantList) -> IdRange<Variant> {
         let start = self.next_variant_idx();
         for variant in variants.variants() {
             if let Some(data) = self.lower_variant(&variant) {
-                self.data().variants.alloc(data);
+                let idx = self.data().variants.alloc(data);
+                self.add_attrs(idx.into(), Attrs::new(&variant, &self.hygiene));
             }
         }
         let end = self.next_variant_idx();
-        start..end
+        IdRange::new(start..end)
     }
 
     fn lower_variant(&mut self, variant: &ast::EnumVariant) -> Option<Variant> {
@@ -419,7 +422,7 @@ impl Ctx {
                         let attrs = Attrs::new(&item, &this.hygiene);
                         this.collect_inner_items(item.syntax());
                         this.lower_assoc_item(&item).map(|item| {
-                            this.add_attrs(item.into(), attrs);
+                            this.add_attrs(ModItem::from(item).into(), attrs);
                             item
                         })
                     })
@@ -453,7 +456,7 @@ impl Ctx {
                 self.collect_inner_items(item.syntax());
                 let assoc = self.lower_assoc_item(&item)?;
                 let attrs = Attrs::new(&item, &self.hygiene);
-                self.add_attrs(assoc.into(), attrs);
+                self.add_attrs(ModItem::from(assoc).into(), attrs);
                 Some(assoc)
             })
             .collect();
@@ -539,7 +542,7 @@ impl Ctx {
                 .filter_map(|item| {
                     self.collect_inner_items(item.syntax());
                     let attrs = Attrs::new(&item, &self.hygiene);
-                    let id = match item {
+                    let id: ModItem = match item {
                         ast::ExternItem::FnDef(ast) => {
                             let func = self.lower_function(&ast)?;
                             func.into()
@@ -549,7 +552,7 @@ impl Ctx {
                             statik.into()
                         }
                     };
-                    self.add_attrs(id, attrs);
+                    self.add_attrs(id.into(), attrs);
                     Some(id)
                 })
                 .collect()
diff --git a/crates/ra_hir_def/src/item_tree/tests.rs b/crates/ra_hir_def/src/item_tree/tests.rs
index dc035d809..08559fb92 100644
--- a/crates/ra_hir_def/src/item_tree/tests.rs
+++ b/crates/ra_hir_def/src/item_tree/tests.rs
@@ -92,7 +92,7 @@ fn print_item_tree(ra_fixture: &str) -> String {
 }
 
 fn fmt_mod_item(out: &mut String, tree: &ItemTree, item: ModItem) {
-    let attrs = tree.attrs(item);
+    let attrs = tree.attrs(item.into());
     if !attrs.is_empty() {
         format_to!(out, "#[{:?}]\n", attrs);
     }
@@ -237,13 +237,13 @@ Trait { name: Name(Text("Tr")), visibility: RawVisibilityId("pub(self)"), generi
 #[Attrs { entries: Some([Attr { path: ModPath { kind: Plain, segments: [Name(Text("struct0"))] }, input: None }]) }]
 Struct { name: Name(Text("Struct0")), visibility: RawVisibilityId("pub(self)"), generic_params: GenericParamsId(1), fields: Unit, ast_id: FileAstId::<ra_syntax::ast::generated::nodes::StructDef>(3), kind: Unit }
 #[Attrs { entries: Some([Attr { path: ModPath { kind: Plain, segments: [Name(Text("struct1"))] }, input: None }]) }]
-Struct { name: Name(Text("Struct1")), visibility: RawVisibilityId("pub(self)"), generic_params: GenericParamsId(2), fields: Tuple(Idx::<Field>(0)..Idx::<Field>(1)), ast_id: FileAstId::<ra_syntax::ast::generated::nodes::StructDef>(4), kind: Tuple }
+Struct { name: Name(Text("Struct1")), visibility: RawVisibilityId("pub(self)"), generic_params: GenericParamsId(2), fields: Tuple(IdRange::<ra_hir_def::item_tree::Field>(0..1)), ast_id: FileAstId::<ra_syntax::ast::generated::nodes::StructDef>(4), kind: Tuple }
 #[Attrs { entries: Some([Attr { path: ModPath { kind: Plain, segments: [Name(Text("struct2"))] }, input: None }]) }]
-Struct { name: Name(Text("Struct2")), visibility: RawVisibilityId("pub(self)"), generic_params: GenericParamsId(3), fields: Record(Idx::<Field>(1)..Idx::<Field>(2)), ast_id: FileAstId::<ra_syntax::ast::generated::nodes::StructDef>(5), kind: Record }
+Struct { name: Name(Text("Struct2")), visibility: RawVisibilityId("pub(self)"), generic_params: GenericParamsId(3), fields: Record(IdRange::<ra_hir_def::item_tree::Field>(1..2)), ast_id: FileAstId::<ra_syntax::ast::generated::nodes::StructDef>(5), kind: Record }
 #[Attrs { entries: Some([Attr { path: ModPath { kind: Plain, segments: [Name(Text("en"))] }, input: None }]) }]
-Enum { name: Name(Text("En")), visibility: RawVisibilityId("pub(self)"), generic_params: GenericParamsId(4294967295), variants: Idx::<Variant>(0)..Idx::<Variant>(1), ast_id: FileAstId::<ra_syntax::ast::generated::nodes::EnumDef>(6) }
+Enum { name: Name(Text("En")), visibility: RawVisibilityId("pub(self)"), generic_params: GenericParamsId(4294967295), variants: IdRange::<ra_hir_def::item_tree::Variant>(0..1), ast_id: FileAstId::<ra_syntax::ast::generated::nodes::EnumDef>(6) }
 #[Attrs { entries: Some([Attr { path: ModPath { kind: Plain, segments: [Name(Text("un"))] }, input: None }]) }]
-Union { name: Name(Text("Un")), visibility: RawVisibilityId("pub(self)"), generic_params: GenericParamsId(4294967295), fields: Record(Idx::<Field>(3)..Idx::<Field>(4)), ast_id: FileAstId::<ra_syntax::ast::generated::nodes::UnionDef>(7) }
+Union { name: Name(Text("Un")), visibility: RawVisibilityId("pub(self)"), generic_params: GenericParamsId(4294967295), fields: Record(IdRange::<ra_hir_def::item_tree::Field>(3..4)), ast_id: FileAstId::<ra_syntax::ast::generated::nodes::UnionDef>(7) }
     "###);
 }
 
diff --git a/crates/ra_hir_def/src/nameres/collector.rs b/crates/ra_hir_def/src/nameres/collector.rs
index 94da700ad..2ced4f66b 100644
--- a/crates/ra_hir_def/src/nameres/collector.rs
+++ b/crates/ra_hir_def/src/nameres/collector.rs
@@ -742,7 +742,7 @@ impl ModCollector<'_, '_> {
         // `#[macro_use] extern crate` is hoisted to imports macros before collecting
         // any other items.
         for item in items {
-            if self.is_cfg_enabled(self.item_tree.attrs(*item)) {
+            if self.is_cfg_enabled(self.item_tree.attrs((*item).into())) {
                 if let ModItem::ExternCrate(id) = item {
                     let import = self.item_tree[*id].clone();
                     if import.is_macro_use {
@@ -753,7 +753,7 @@ impl ModCollector<'_, '_> {
         }
 
         for &item in items {
-            let attrs = self.item_tree.attrs(item);
+            let attrs = self.item_tree.attrs(item.into());
             if self.is_cfg_enabled(attrs) {
                 let module =
                     ModuleId { krate: self.def_collector.def_map.krate, local_id: self.module_id };