Track source file IDs in source mapping of Attrs

2 files changed, 28 insertions, 10 deletions

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) {