2 files changed, 84 insertions, 59 deletions

diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index 469811ce8..aadd4e44a 100644
--- a/crates/hir_def/src/attr.rs
+++ b/crates/hir_def/src/attr.rs
@@ -101,17 +101,13 @@ impl RawAttrs {
         hygiene: &Hygiene,
     ) -> Self {
         let entries = collect_attrs(owner)
-            .enumerate()
-            .flat_map(|(i, attr)| {
-                let index = AttrId(i as u32);
-                match attr {
-                    Either::Left(attr) => Attr::from_src(db, 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))),
-                    }),
-                }
+            .flat_map(|(id, attr)| match attr {
+                Either::Left(attr) => Attr::from_src(db, attr, hygiene, id),
+                Either::Right(comment) => comment.doc_comment().map(|doc| Attr {
+                    id,
+                    input: Some(AttrInput::Literal(SmolStr::new(doc))),
+                    path: Interned::new(ModPath::from(hir_expand::name!(doc))),
+                }),
             })
             .collect::<Arc<_>>();
 
@@ -124,6 +120,7 @@ impl RawAttrs {
     }
 
     pub(crate) fn merge(&self, other: Self) -> Self {
+        // FIXME: This needs to fixup `AttrId`s
         match (&self.entries, &other.entries) {
             (None, None) => Self::EMPTY,
             (Some(entries), None) | (None, Some(entries)) => {
@@ -375,39 +372,26 @@ impl AttrsWithOwner {
 
                 let def_map = module.def_map(db);
                 let mod_data = &def_map[module.local_id];
-                let attrs = match mod_data.declaration_source(db) {
+                match mod_data.declaration_source(db) {
                     Some(it) => {
-                        let mut attrs: Vec<_> = collect_attrs(&it.value as &dyn ast::AttrsOwner)
-                            .map(|attr| InFile::new(it.file_id, attr))
-                            .collect();
+                        let mut map = AttrSourceMap::new(InFile::new(it.file_id, &it.value));
                         if let InFile { file_id, value: ModuleSource::SourceFile(file) } =
                             mod_data.definition_source(db)
                         {
-                            attrs.extend(
-                                collect_attrs(&file as &dyn ast::AttrsOwner)
-                                    .map(|attr| InFile::new(file_id, attr)),
-                            )
+                            map.merge(AttrSourceMap::new(InFile::new(file_id, &file)));
                         }
-                        attrs
+                        return map;
                     }
                     None => {
                         let InFile { file_id, value } = mod_data.definition_source(db);
-                        match &value {
-                            ModuleSource::SourceFile(file) => {
-                                collect_attrs(file as &dyn ast::AttrsOwner)
-                            }
-                            ModuleSource::Module(module) => {
-                                collect_attrs(module as &dyn ast::AttrsOwner)
-                            }
-                            ModuleSource::BlockExpr(block) => {
-                                collect_attrs(block as &dyn ast::AttrsOwner)
-                            }
-                        }
-                        .map(|attr| InFile::new(file_id, attr))
-                        .collect()
+                        let attrs_owner = match &value {
+                            ModuleSource::SourceFile(file) => file as &dyn ast::AttrsOwner,
+                            ModuleSource::Module(module) => module as &dyn ast::AttrsOwner,
+                            ModuleSource::BlockExpr(block) => block as &dyn ast::AttrsOwner,
+                        };
+                        return AttrSourceMap::new(InFile::new(file_id, attrs_owner));
                     }
-                };
-                return AttrSourceMap { attrs };
+                }
             }
             AttrDefId::FieldId(id) => {
                 let map = db.fields_attrs_source_map(id.parent);
@@ -462,11 +446,7 @@ impl AttrsWithOwner {
             },
         };
 
-        AttrSourceMap {
-            attrs: collect_attrs(&owner.value)
-                .map(|attr| InFile::new(owner.file_id, attr))
-                .collect(),
-        }
+        AttrSourceMap::new(owner.as_ref().map(|node| node as &dyn AttrsOwner))
     }
 
     pub fn docs_with_rangemap(
@@ -518,7 +498,7 @@ impl AttrsWithOwner {
         if buf.is_empty() {
             None
         } else {
-            Some((Documentation(buf), DocsRangeMap { mapping, source: self.source_map(db).attrs }))
+            Some((Documentation(buf), DocsRangeMap { mapping, source_map: self.source_map(db) }))
         }
     }
 }
@@ -559,27 +539,59 @@ fn inner_attributes(
 }
 
 pub struct AttrSourceMap {
-    attrs: Vec<InFile<Either<ast::Attr, ast::Comment>>>,
+    attrs: Vec<InFile<ast::Attr>>,
+    doc_comments: Vec<InFile<ast::Comment>>,
 }
 
 impl AttrSourceMap {
+    fn new(owner: InFile<&dyn ast::AttrsOwner>) -> Self {
+        let mut attrs = Vec::new();
+        let mut doc_comments = Vec::new();
+        for (_, attr) in collect_attrs(owner.value) {
+            match attr {
+                Either::Left(attr) => attrs.push(owner.with_value(attr)),
+                Either::Right(comment) => doc_comments.push(owner.with_value(comment)),
+            }
+        }
+
+        Self { attrs, doc_comments }
+    }
+
+    fn merge(&mut self, other: Self) {
+        self.attrs.extend(other.attrs);
+        self.doc_comments.extend(other.doc_comments);
+    }
+
     /// Maps the lowered `Attr` back to its original syntax node.
     ///
     /// `attr` must come from the `owner` used for AttrSourceMap
     ///
     /// Note that the returned syntax node might be a `#[cfg_attr]`, or a doc comment, instead of
     /// the attribute represented by `Attr`.
-    pub fn source_of(&self, attr: &Attr) -> InFile<&Either<ast::Attr, ast::Comment>> {
-        self.attrs
-            .get(attr.id.0 as usize)
-            .unwrap_or_else(|| panic!("cannot find `Attr` at index {:?}", attr.id))
-            .as_ref()
+    pub fn source_of(&self, attr: &Attr) -> InFile<Either<ast::Attr, ast::Comment>> {
+        self.source_of_id(attr.id)
+    }
+
+    fn source_of_id(&self, id: AttrId) -> InFile<Either<ast::Attr, ast::Comment>> {
+        if id.is_doc_comment {
+            self.doc_comments
+                .get(id.ast_index as usize)
+                .unwrap_or_else(|| panic!("cannot find doc comment at index {:?}", id))
+                .clone()
+                .map(|attr| Either::Right(attr))
+        } else {
+            self.attrs
+                .get(id.ast_index as usize)
+                .unwrap_or_else(|| panic!("cannot find `Attr` at index {:?}", id))
+                .clone()
+                .map(|attr| Either::Left(attr))
+        }
     }
 }
 
 /// A struct to map text ranges from [`Documentation`] back to TextRanges in the syntax tree.
 pub struct DocsRangeMap {
-    source: Vec<InFile<Either<ast::Attr, ast::Comment>>>,
+    source_map: AttrSourceMap,
     // (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
@@ -596,7 +608,7 @@ impl DocsRangeMap {
 
         let relative_range = range - line_docs_range.start();
 
-        let &InFile { file_id, value: ref source } = &self.source[idx.0 as usize];
+        let &InFile { file_id, value: ref source } = &self.source_map.source_of_id(idx);
         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
@@ -616,7 +628,10 @@ impl DocsRangeMap {
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-pub struct AttrId(pub u32);
+pub(crate) struct AttrId {
+    is_doc_comment: bool,
+    pub(crate) ast_index: u32,
+}
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Attr {
@@ -752,22 +767,32 @@ fn attrs_from_item_tree<N: ItemTreeNode>(id: ItemTreeId<N>, db: &dyn DefDatabase
 
 fn collect_attrs(
     owner: &dyn ast::AttrsOwner,
-) -> impl Iterator<Item = Either<ast::Attr, ast::Comment>> {
+) -> impl Iterator<Item = (AttrId, Either<ast::Attr, ast::Comment>)> {
     let (inner_attrs, inner_docs) = inner_attributes(owner.syntax())
         .map_or((None, None), |(attrs, docs)| (Some(attrs), Some(docs)));
 
     let outer_attrs = owner.attrs().filter(|attr| attr.kind().is_outer());
-    let attrs = outer_attrs
-        .chain(inner_attrs.into_iter().flatten())
-        .map(|attr| (attr.syntax().text_range().start(), Either::Left(attr)));
+    let attrs =
+        outer_attrs.chain(inner_attrs.into_iter().flatten()).enumerate().map(|(idx, attr)| {
+            (
+                AttrId { ast_index: idx as u32, is_doc_comment: false },
+                attr.syntax().text_range().start(),
+                Either::Left(attr),
+            )
+        });
 
     let outer_docs =
         ast::CommentIter::from_syntax_node(owner.syntax()).filter(ast::Comment::is_outer);
-    let docs = outer_docs
-        .chain(inner_docs.into_iter().flatten())
-        .map(|docs_text| (docs_text.syntax().text_range().start(), Either::Right(docs_text)));
+    let docs =
+        outer_docs.chain(inner_docs.into_iter().flatten()).enumerate().map(|(idx, docs_text)| {
+            (
+                AttrId { ast_index: idx as u32, is_doc_comment: true },
+                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
-    docs.chain(attrs).sorted_by_key(|&(offset, _)| offset).map(|(_, attr)| attr)
+    docs.chain(attrs).sorted_by_key(|&(_, offset, _)| offset).map(|(id, _, attr)| (id, attr))
 }
 
 pub(crate) fn variants_attrs_source_map(
diff --git a/crates/hir_def/src/lib.rs b/crates/hir_def/src/lib.rs
index 869e0280e..a82ea5957 100644
--- a/crates/hir_def/src/lib.rs
+++ b/crates/hir_def/src/lib.rs
@@ -753,7 +753,7 @@ fn derive_macro_as_call_id(
             MacroCallKind::Derive {
                 ast_id: item_attr.ast_id,
                 derive_name: last_segment.to_string(),
-                derive_attr_index: derive_attr.0,
+                derive_attr_index: derive_attr.ast_index,
             },
         )
         .into();