3 files changed, 76 insertions, 26 deletions

diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index e4c84afbf..739e3f5e3 100644
--- a/crates/hir_def/src/attr.rs
+++ b/crates/hir_def/src/attr.rs
@@ -283,8 +283,51 @@ impl Attrs {
     /// Constructs a map that maps the lowered `Attr`s in this `Attrs` back to its original syntax nodes.
     ///
     /// `owner` must be the original owner of the attributes.
-    pub fn source_map(&self, owner: &dyn ast::AttrsOwner) -> AttrSourceMap {
-        AttrSourceMap { attrs: collect_attrs(owner).collect() }
+    // FIXME: figure out a better api that doesnt require the for_module hack
+    pub fn source_map(&self, owner: InFile<&dyn ast::AttrsOwner>) -> AttrSourceMap {
+        // FIXME: This doesn't work correctly for modules, as the attributes there can have up to
+        // two different owners
+        AttrSourceMap {
+            attrs: collect_attrs(owner.value)
+                .map(|attr| InFile::new(owner.file_id, attr))
+                .collect(),
+        }
+    }
+
+    pub fn source_map_for_module(
+        &self,
+        db: &dyn DefDatabase,
+        module: crate::ModuleId,
+    ) -> AttrSourceMap {
+        let def_map = module.def_map(db);
+        let mod_data = &def_map[module.local_id];
+        let attrs = 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();
+                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)),
+                    )
+                }
+                attrs
+            }
+            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()
+            }
+        };
+        AttrSourceMap { attrs }
     }
 
     pub fn by_key(&self, key: &'static str) -> AttrQuery<'_> {
@@ -379,7 +422,7 @@ fn inner_attributes(
 }
 
 pub struct AttrSourceMap {
-    attrs: Vec<Either<ast::Attr, ast::Comment>>,
+    attrs: Vec<InFile<Either<ast::Attr, ast::Comment>>>,
 }
 
 impl AttrSourceMap {
@@ -389,10 +432,11 @@ impl 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) -> &Either<ast::Attr, ast::Comment> {
+    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))
+            .as_ref()
     }
 }
 
@@ -428,18 +472,6 @@ impl Attr {
         Some(Attr { index, path, input })
     }
 
-    /// Maps this lowered `Attr` back to its original syntax node.
-    ///
-    /// `owner` must be the original owner of the attribute.
-    ///
-    /// Note that the returned syntax node might be a `#[cfg_attr]`, or a doc comment, instead of
-    /// the attribute represented by `Attr`.
-    pub fn to_src(&self, owner: &dyn ast::AttrsOwner) -> Either<ast::Attr, ast::Comment> {
-        collect_attrs(owner).nth(self.index as usize).unwrap_or_else(|| {
-            panic!("cannot find `Attr` at index {} in {}", self.index, owner.syntax())
-        })
-    }
-
     /// Parses this attribute as a `#[derive]`, returns an iterator that yields all contained paths
     /// to derive macros.
     ///
diff --git a/crates/ide/src/syntax_highlighting.rs b/crates/ide/src/syntax_highlighting.rs
index ba3447b3a..e25b698e0 100644
--- a/crates/ide/src/syntax_highlighting.rs
+++ b/crates/ide/src/syntax_highlighting.rs
@@ -12,7 +12,7 @@ mod html;
 #[cfg(test)]
 mod tests;
 
-use hir::{Name, Semantics};
+use hir::{InFile, Name, Semantics};
 use ide_db::{RootDatabase, SymbolKind};
 use rustc_hash::FxHashMap;
 use syntax::{
@@ -73,14 +73,20 @@ pub(crate) fn highlight(
     };
 
     let mut hl = highlights::Highlights::new(root.text_range());
-    traverse(&mut hl, &sema, &root, range_to_highlight, syntactic_name_ref_highlighting);
+    traverse(
+        &mut hl,
+        &sema,
+        InFile::new(file_id.into(), &root),
+        range_to_highlight,
+        syntactic_name_ref_highlighting,
+    );
     hl.to_vec()
 }
 
 fn traverse(
     hl: &mut Highlights,
     sema: &Semantics<RootDatabase>,
-    root: &SyntaxNode,
+    root: InFile<&SyntaxNode>,
     range_to_highlight: TextRange,
     syntactic_name_ref_highlighting: bool,
 ) {
@@ -93,7 +99,7 @@ fn traverse(
 
     // Walk all nodes, keeping track of whether we are inside a macro or not.
     // If in macro, expand it first and highlight the expanded code.
-    for event in root.preorder_with_tokens() {
+    for event in root.value.preorder_with_tokens() {
         let event_range = match &event {
             WalkEvent::Enter(it) | WalkEvent::Leave(it) => it.text_range(),
         };
@@ -150,7 +156,7 @@ fn traverse(
             WalkEvent::Enter(it) => it,
             WalkEvent::Leave(it) => {
                 if let Some(node) = it.as_node() {
-                    inject::doc_comment(hl, sema, node);
+                    inject::doc_comment(hl, sema, root.with_value(node));
                 }
                 continue;
             }
diff --git a/crates/ide/src/syntax_highlighting/inject.rs b/crates/ide/src/syntax_highlighting/inject.rs
index 947cc974c..e6dbd307e 100644
--- a/crates/ide/src/syntax_highlighting/inject.rs
+++ b/crates/ide/src/syntax_highlighting/inject.rs
@@ -3,7 +3,7 @@
 use std::{mem, ops::Range};
 
 use either::Either;
-use hir::{HasAttrs, Semantics};
+use hir::{HasAttrs, InFile, Semantics};
 use ide_db::{call_info::ActiveParameter, defs::Definition};
 use syntax::{
     ast::{self, AstNode, AttrsOwner, DocCommentsOwner},
@@ -148,8 +148,12 @@ fn doc_attributes<'node>(
 }
 
 /// Injection of syntax highlighting of doctests.
-pub(super) fn doc_comment(hl: &mut Highlights, sema: &Semantics<RootDatabase>, node: &SyntaxNode) {
-    let (owner, attributes, def) = match doc_attributes(sema, node) {
+pub(super) fn doc_comment(
+    hl: &mut Highlights,
+    sema: &Semantics<RootDatabase>,
+    node: InFile<&SyntaxNode>,
+) {
+    let (owner, attributes, def) = match doc_attributes(sema, node.value) {
         Some(it) => it,
         None => return,
     };
@@ -157,7 +161,12 @@ pub(super) fn doc_comment(hl: &mut Highlights, sema: &Semantics<RootDatabase>, n
     let mut inj = Injector::default();
     inj.add_unmapped("fn doctest() {\n");
 
-    let attrs_source_map = attributes.source_map(&owner);
+    let attrs_source_map = match def {
+        Definition::ModuleDef(hir::ModuleDef::Module(module)) => {
+            attributes.source_map_for_module(sema.db, module.into())
+        }
+        _ => attributes.source_map(node.with_value(&owner)),
+    };
 
     let mut is_codeblock = false;
     let mut is_doctest = false;
@@ -168,7 +177,10 @@ pub(super) fn doc_comment(hl: &mut Highlights, sema: &Semantics<RootDatabase>, n
     let mut intra_doc_links = Vec::new();
     let mut string;
     for attr in attributes.by_key("doc").attrs() {
-        let src = attrs_source_map.source_of(&attr);
+        let InFile { file_id, value: src } = attrs_source_map.source_of(&attr);
+        if file_id != node.file_id {
+            continue;
+        }
         let (line, range, prefix) = match &src {
             Either::Left(it) => {
                 string = match find_doc_string_in_attr(attr, it) {