1 files changed, 50 insertions, 16 deletions

diff --git a/crates/ide/src/syntax_highlighting/inject.rs b/crates/ide/src/syntax_highlighting/inject.rs
index 90feb6c18..5722dea3a 100644
--- a/crates/ide/src/syntax_highlighting/inject.rs
+++ b/crates/ide/src/syntax_highlighting/inject.rs
@@ -3,8 +3,8 @@
 use std::{mem, ops::Range};
 
 use either::Either;
-use hir::{HasAttrs, Semantics};
-use ide_db::{call_info::ActiveParameter, defs::Definition};
+use hir::{HasAttrs, InFile, Semantics};
+use ide_db::{call_info::ActiveParameter, defs::Definition, SymbolKind};
 use syntax::{
     ast::{self, AstNode, AttrsOwner, DocCommentsOwner},
     match_ast, AstToken, NodeOrToken, SyntaxNode, SyntaxToken, TextRange, TextSize,
@@ -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) {
@@ -213,13 +225,16 @@ pub(super) fn doc_comment(hl: &mut Highlights, sema: &Semantics<RootDatabase>, n
                     intra_doc_links.extend(
                         extract_definitions_from_markdown(line)
                             .into_iter()
-                            .filter(|(link, ns, _)| {
-                                validate_intra_doc_link(sema.db, &def, link, *ns)
+                            .filter_map(|(link, ns, range)| {
+                                validate_intra_doc_link(sema.db, &def, &link, ns).zip(Some(range))
                             })
-                            .map(|(.., Range { start, end })| {
-                                TextRange::at(
-                                    prev_range_start + TextSize::from(start as u32),
-                                    TextSize::from((end - start) as u32),
+                            .map(|(def, Range { start, end })| {
+                                (
+                                    def,
+                                    TextRange::at(
+                                        prev_range_start + TextSize::from(start as u32),
+                                        TextSize::from((end - start) as u32),
+                                    ),
                                 )
                             }),
                     );
@@ -243,10 +258,13 @@ pub(super) fn doc_comment(hl: &mut Highlights, sema: &Semantics<RootDatabase>, n
         }
     }
 
-    for range in intra_doc_links {
+    for (def, range) in intra_doc_links {
         hl.add(HlRange {
             range,
-            highlight: HlTag::IntraDocLink | HlMod::Documentation,
+            highlight: module_def_to_hl_tag(def)
+                | HlMod::Documentation
+                | HlMod::Injected
+                | HlMod::IntraDocLink,
             binding_hash: None,
         });
     }
@@ -306,7 +324,7 @@ fn validate_intra_doc_link(
     def: &Definition,
     link: &str,
     ns: Option<hir::Namespace>,
-) -> bool {
+) -> Option<hir::ModuleDef> {
     match def {
         Definition::ModuleDef(def) => match def {
             hir::ModuleDef::Module(it) => it.resolve_doc_path(db, &link, ns),
@@ -326,5 +344,21 @@ fn validate_intra_doc_link(
         | Definition::GenericParam(_)
         | Definition::Label(_) => None,
     }
-    .is_some()
+}
+
+fn module_def_to_hl_tag(def: hir::ModuleDef) -> HlTag {
+    let symbol = match def {
+        hir::ModuleDef::Module(_) => SymbolKind::Module,
+        hir::ModuleDef::Function(_) => SymbolKind::Function,
+        hir::ModuleDef::Adt(hir::Adt::Struct(_)) => SymbolKind::Struct,
+        hir::ModuleDef::Adt(hir::Adt::Enum(_)) => SymbolKind::Enum,
+        hir::ModuleDef::Adt(hir::Adt::Union(_)) => SymbolKind::Union,
+        hir::ModuleDef::Variant(_) => SymbolKind::Variant,
+        hir::ModuleDef::Const(_) => SymbolKind::Const,
+        hir::ModuleDef::Static(_) => SymbolKind::Static,
+        hir::ModuleDef::Trait(_) => SymbolKind::Trait,
+        hir::ModuleDef::TypeAlias(_) => SymbolKind::TypeAlias,
+        hir::ModuleDef::BuiltinType(_) => return HlTag::BuiltinType,
+    };
+    HlTag::Symbol(symbol)
 }