Collect field/variant attrs in ItemTree

6 files changed, 40 insertions, 20 deletions

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 f99e05432..fd874750e 100644
--- a/crates/ra_hir_def/src/item_tree.rs
+++ b/crates/ra_hir_def/src/item_tree.rs
@@ -178,8 +178,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 +282,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>;
 
diff --git a/crates/ra_hir_def/src/item_tree/lower.rs b/crates/ra_hir_def/src/item_tree/lower.rs
index e2e00323c..230e1f768 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);
             }
@@ -200,7 +200,8 @@ impl Ctx {
         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();
@@ -219,7 +220,8 @@ impl Ctx {
         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();
@@ -266,7 +268,8 @@ impl Ctx {
         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();
@@ -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 fd7ffee24..18df42050 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);
     }
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 };