1 files changed, 72 insertions, 33 deletions

diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index ab77d924a..786fad6e1 100644
--- a/crates/hir_def/src/attr.rs
+++ b/crates/hir_def/src/attr.rs
@@ -9,14 +9,14 @@ use std::{
 use base_db::CrateId;
 use cfg::{CfgExpr, CfgOptions};
 use either::Either;
-use hir_expand::{hygiene::Hygiene, name::AsName, AstId, InFile};
+use hir_expand::{hygiene::Hygiene, name::AsName, AstId, AttrId, InFile};
 use itertools::Itertools;
 use la_arena::ArenaMap;
 use mbe::ast_to_token_tree;
 use smallvec::{smallvec, SmallVec};
 use syntax::{
     ast::{self, AstNode, AttrsOwner},
-    match_ast, AstToken, SmolStr, SyntaxNode, TextRange, TextSize,
+    match_ast, AstPtr, AstToken, SmolStr, SyntaxNode, TextRange, TextSize,
 };
 use tt::Subtree;
 
@@ -98,13 +98,16 @@ impl RawAttrs {
     pub(crate) fn new(owner: &dyn ast::AttrsOwner, hygiene: &Hygiene) -> Self {
         let entries = collect_attrs(owner)
             .enumerate()
-            .flat_map(|(i, attr)| match attr {
-                Either::Left(attr) => Attr::from_src(attr, hygiene, i as u32),
-                Either::Right(comment) => comment.doc_comment().map(|doc| Attr {
-                    index: i as u32,
-                    input: Some(AttrInput::Literal(SmolStr::new(doc))),
-                    path: Interned::new(ModPath::from(hir_expand::name!(doc))),
-                }),
+            .flat_map(|(i, attr)| {
+                let index = AttrId(i as u32);
+                match attr {
+                    Either::Left(attr) => Attr::from_src(attr, hygiene, index),
+                    Either::Right(comment) => comment.doc_comment().map(|doc| Attr {
+                        id: index,
+                        input: Some(AttrInput::Literal(SmolStr::new(doc))),
+                        path: Interned::new(ModPath::from(hir_expand::name!(doc))),
+                    }),
+                }
             })
             .collect::<Arc<_>>();
 
@@ -161,7 +164,7 @@ impl RawAttrs {
                 let cfg = parts.next().unwrap();
                 let cfg = Subtree { delimiter: subtree.delimiter, token_trees: cfg.to_vec() };
                 let cfg = CfgExpr::parse(&cfg);
-                let index = attr.index;
+                let index = attr.id;
                 let attrs = parts.filter(|a| !a.is_empty()).filter_map(|attr| {
                     let tree = Subtree { delimiter: None, token_trees: attr.to_vec() };
                     let attr = ast::Attr::parse(&format!("#[{}]", tree)).ok()?;
@@ -215,12 +218,11 @@ impl Attrs {
         let mut res = ArenaMap::default();
 
         for (id, fld) in src.value.iter() {
-            let attrs = match fld {
-                Either::Left(_tuple) => Attrs::default(),
-                Either::Right(record) => {
-                    RawAttrs::from_attrs_owner(db, src.with_value(record)).filter(db, krate)
-                }
+            let owner: &dyn AttrsOwner = match fld {
+                Either::Left(tuple) => tuple,
+                Either::Right(record) => record,
             };
+            let attrs = RawAttrs::from_attrs_owner(db, src.with_value(owner)).filter(db, krate);
 
             res.insert(id, attrs);
         }
@@ -404,10 +406,14 @@ impl AttrsWithOwner {
                 return AttrSourceMap { attrs };
             }
             AttrDefId::FieldId(id) => {
-                id.parent.child_source(db).map(|source| match &source[id.local_id] {
-                    Either::Left(field) => ast::AttrsOwnerNode::new(field.clone()),
-                    Either::Right(field) => ast::AttrsOwnerNode::new(field.clone()),
-                })
+                let map = db.fields_attrs_source_map(id.parent);
+                let file_id = id.parent.file_id(db);
+                let root = db.parse_or_expand(file_id).unwrap();
+                let owner = match &map[id.local_id] {
+                    Either::Left(it) => ast::AttrsOwnerNode::new(it.to_node(&root)),
+                    Either::Right(it) => ast::AttrsOwnerNode::new(it.to_node(&root)),
+                };
+                InFile::new(file_id, owner)
             }
             AttrDefId::AdtId(adt) => match adt {
                 AdtId::StructId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
@@ -415,10 +421,12 @@ impl AttrsWithOwner {
                 AdtId::EnumId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
             },
             AttrDefId::FunctionId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
-            AttrDefId::EnumVariantId(id) => id
-                .parent
-                .child_source(db)
-                .map(|source| ast::AttrsOwnerNode::new(source[id.local_id].clone())),
+            AttrDefId::EnumVariantId(id) => {
+                let map = db.variants_attrs_source_map(id.parent);
+                let file_id = id.parent.lookup(db).id.file_id();
+                let root = db.parse_or_expand(file_id).unwrap();
+                InFile::new(file_id, ast::AttrsOwnerNode::new(map[id.local_id].to_node(&root)))
+            }
             AttrDefId::StaticId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
             AttrDefId::ConstId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
             AttrDefId::TraitId(id) => id.lookup(db).source(db).map(ast::AttrsOwnerNode::new),
@@ -463,7 +471,7 @@ impl AttrsWithOwner {
     ) -> Option<(Documentation, DocsRangeMap)> {
         // FIXME: code duplication in `docs` above
         let docs = self.by_key("doc").attrs().flat_map(|attr| match attr.input.as_ref()? {
-            AttrInput::Literal(s) => Some((s, attr.index)),
+            AttrInput::Literal(s) => Some((s, attr.id)),
             AttrInput::TokenTree(_) => None,
         });
         let indent = docs
@@ -555,8 +563,8 @@ impl AttrSourceMap {
     /// the attribute represented by `Attr`.
     pub fn source_of(&self, attr: &Attr) -> InFile<&Either<ast::Attr, ast::Comment>> {
         self.attrs
-            .get(attr.index as usize)
-            .unwrap_or_else(|| panic!("cannot find `Attr` at index {}", attr.index))
+            .get(attr.id.0 as usize)
+            .unwrap_or_else(|| panic!("cannot find `Attr` at index {:?}", attr.id))
             .as_ref()
     }
 }
@@ -567,7 +575,7 @@ pub struct DocsRangeMap {
     // (docstring-line-range, attr_index, attr-string-range)
     // a mapping from the text range of a line of the [`Documentation`] to the attribute index and
     // the original (untrimmed) syntax doc line
-    mapping: Vec<(TextRange, u32, TextRange)>,
+    mapping: Vec<(TextRange, AttrId, TextRange)>,
 }
 
 impl DocsRangeMap {
@@ -580,7 +588,7 @@ impl DocsRangeMap {
 
         let relative_range = range - line_docs_range.start();
 
-        let &InFile { file_id, value: ref source } = &self.source[idx as usize];
+        let &InFile { file_id, value: ref source } = &self.source[idx.0 as usize];
         match source {
             Either::Left(_) => None, // FIXME, figure out a nice way to handle doc attributes here
             // as well as for whats done in syntax highlight doc injection
@@ -601,7 +609,7 @@ impl DocsRangeMap {
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Attr {
-    index: u32,
+    pub(crate) id: AttrId,
     pub(crate) path: Interned<ModPath>,
     pub(crate) input: Option<AttrInput>,
 }
@@ -615,7 +623,7 @@ pub enum AttrInput {
 }
 
 impl Attr {
-    fn from_src(ast: ast::Attr, hygiene: &Hygiene, index: u32) -> Option<Attr> {
+    fn from_src(ast: ast::Attr, hygiene: &Hygiene, id: AttrId) -> Option<Attr> {
         let path = Interned::new(ModPath::from_src(ast.path()?, hygiene)?);
         let input = if let Some(ast::Expr::Literal(lit)) = ast.expr() {
             let value = match lit.kind() {
@@ -628,7 +636,7 @@ impl Attr {
         } else {
             None
         };
-        Some(Attr { index, path, input })
+        Some(Attr { id, path, input })
     }
 
     /// Parses this attribute as a `#[derive]`, returns an iterator that yields all contained paths
@@ -743,7 +751,38 @@ fn collect_attrs(
         .chain(inner_docs.into_iter().flatten())
         .map(|docs_text| (docs_text.syntax().text_range().start(), Either::Right(docs_text)));
     // sort here by syntax node offset because the source can have doc attributes and doc strings be interleaved
-    let attrs: Vec<_> = docs.chain(attrs).sorted_by_key(|&(offset, _)| offset).collect();
+    docs.chain(attrs).sorted_by_key(|&(offset, _)| offset).map(|(_, attr)| attr)
+}
+
+pub(crate) fn variants_attrs_source_map(
+    db: &dyn DefDatabase,
+    def: EnumId,
+) -> Arc<ArenaMap<LocalEnumVariantId, AstPtr<ast::Variant>>> {
+    let mut res = ArenaMap::default();
+    let child_source = def.child_source(db);
+
+    for (idx, variant) in child_source.value.iter() {
+        res.insert(idx, AstPtr::new(variant));
+    }
+
+    Arc::new(res)
+}
+
+pub(crate) fn fields_attrs_source_map(
+    db: &dyn DefDatabase,
+    def: VariantId,
+) -> Arc<ArenaMap<LocalFieldId, Either<AstPtr<ast::TupleField>, AstPtr<ast::RecordField>>>> {
+    let mut res = ArenaMap::default();
+    let child_source = def.child_source(db);
+
+    for (idx, variant) in child_source.value.iter() {
+        res.insert(
+            idx,
+            variant
+                .as_ref()
+                .either(|l| Either::Left(AstPtr::new(l)), |r| Either::Right(AstPtr::new(r))),
+        );
+    }
 
-    attrs.into_iter().map(|(_, attr)| attr)
+    Arc::new(res)
 }