Move doc-comment highlight injection from AST to HIR

6 files changed, 63 insertions, 18 deletions

diff --git a/crates/hir/src/attrs.rs b/crates/hir/src/attrs.rs
index 9e6a3e155..505fc05e7 100644
--- a/crates/hir/src/attrs.rs
+++ b/crates/hir/src/attrs.rs
@@ -11,8 +11,8 @@ use hir_ty::db::HirDatabase;
 use syntax::ast;
 
 use crate::{
-    Adt, Const, ConstParam, Enum, Field, Function, GenericParam, LifetimeParam, MacroDef, Module,
-    ModuleDef, Static, Struct, Trait, TypeAlias, TypeParam, Union, Variant,
+    Adt, Const, ConstParam, Enum, Field, Function, GenericParam, Impl, LifetimeParam, MacroDef,
+    Module, ModuleDef, Static, Struct, Trait, TypeAlias, TypeParam, Union, Variant,
 };
 
 pub trait HasAttrs {
@@ -64,6 +64,7 @@ impl_has_attrs![
     (Adt, AdtId),
     (Module, ModuleId),
     (GenericParam, GenericParamId),
+    (Impl, ImplId),
 ];
 
 macro_rules! impl_has_attrs_enum {
diff --git a/crates/hir/src/semantics.rs b/crates/hir/src/semantics.rs
index 03c9371b5..c7e0d0be3 100644
--- a/crates/hir/src/semantics.rs
+++ b/crates/hir/src/semantics.rs
@@ -752,6 +752,7 @@ macro_rules! to_def_impls {
 
 to_def_impls![
     (crate::Module, ast::Module, module_to_def),
+    (crate::Module, ast::SourceFile, source_file_to_def),
     (crate::Struct, ast::Struct, struct_to_def),
     (crate::Enum, ast::Enum, enum_to_def),
     (crate::Union, ast::Union, union_to_def),
diff --git a/crates/hir/src/semantics/source_to_def.rs b/crates/hir/src/semantics/source_to_def.rs
index e9d820140..c6ad5ecb5 100644
--- a/crates/hir/src/semantics/source_to_def.rs
+++ b/crates/hir/src/semantics/source_to_def.rs
@@ -71,6 +71,12 @@ impl SourceToDefCtx<'_, '_> {
         Some(def_map.module_id(child_id))
     }
 
+    pub(super) fn source_file_to_def(&mut self, src: InFile<ast::SourceFile>) -> Option<ModuleId> {
+        let _p = profile::span("source_file_to_def");
+        let file_id = src.file_id.original_file(self.db.upcast());
+        self.file_to_def(file_id).get(0).copied()
+    }
+
     pub(super) fn trait_to_def(&mut self, src: InFile<ast::Trait>) -> Option<TraitId> {
         self.to_def(src, keys::TRAIT)
     }
diff --git a/crates/hir_def/src/attr.rs b/crates/hir_def/src/attr.rs
index 7b41b148c..505c4cd17 100644
--- a/crates/hir_def/src/attr.rs
+++ b/crates/hir_def/src/attr.rs
@@ -475,7 +475,7 @@ impl<'a> AttrQuery<'a> {
         self.attrs().next().is_some()
     }
 
-    pub(crate) fn attrs(self) -> impl Iterator<Item = &'a Attr> {
+    pub fn attrs(self) -> impl Iterator<Item = &'a Attr> {
         let key = self.key;
         self.attrs
             .iter()
diff --git a/crates/ide/src/syntax_highlighting.rs b/crates/ide/src/syntax_highlighting.rs
index 870146d24..ba3447b3a 100644
--- a/crates/ide/src/syntax_highlighting.rs
+++ b/crates/ide/src/syntax_highlighting.rs
@@ -150,7 +150,7 @@ fn traverse(
             WalkEvent::Enter(it) => it,
             WalkEvent::Leave(it) => {
                 if let Some(node) = it.as_node() {
-                    inject::doc_comment(hl, node);
+                    inject::doc_comment(hl, sema, node);
                 }
                 continue;
             }
diff --git a/crates/ide/src/syntax_highlighting/inject.rs b/crates/ide/src/syntax_highlighting/inject.rs
index 8cdc3688f..5b065c09f 100644
--- a/crates/ide/src/syntax_highlighting/inject.rs
+++ b/crates/ide/src/syntax_highlighting/inject.rs
@@ -1,8 +1,12 @@
 //! "Recursive" Syntax highlighting for code in doctests and fixtures.
 
-use hir::Semantics;
+use either::Either;
+use hir::{HasAttrs, Semantics};
 use ide_db::call_info::ActiveParameter;
-use syntax::{ast, AstToken, SyntaxNode, SyntaxToken, TextRange, TextSize};
+use syntax::{
+    ast::{self, AstNode, AttrsOwner},
+    match_ast, AstToken, SyntaxNode, SyntaxToken, TextRange, TextSize,
+};
 
 use crate::{Analysis, HlMod, HlRange, HlTag, RootDatabase};
 
@@ -81,16 +85,46 @@ const RUSTDOC_FENCE_TOKENS: &[&'static str] = &[
     "edition2021",
 ];
 
+fn doc_attributes<'node>(
+    sema: &Semantics<RootDatabase>,
+    node: &'node SyntaxNode,
+) -> Option<(Box<dyn AttrsOwner>, hir::Attrs)> {
+    match_ast! {
+        match node {
+            ast::SourceFile(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            ast::Fn(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            ast::Struct(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            ast::Union(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            ast::RecordField(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            ast::TupleField(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            ast::Enum(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            ast::Variant(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            ast::Trait(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            ast::Module(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            ast::Static(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            ast::Const(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            ast::TypeAlias(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            ast::Impl(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            ast::MacroRules(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            ast::MacroRules(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            // ast::MacroDef(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            // ast::Use(it) => sema.to_def(&it).map(|def| (Box::new(it) as _, def.attrs(sema.db))),
+            _ => return None
+        }
+    }
+}
+
 /// Injection of syntax highlighting of doctests.
-pub(super) fn doc_comment(hl: &mut Highlights, node: &SyntaxNode) {
-    let doc_comments = node
-        .children_with_tokens()
-        .filter_map(|it| it.into_token().and_then(ast::Comment::cast))
-        .filter(|it| it.kind().doc.is_some());
+pub(super) fn doc_comment(hl: &mut Highlights, sema: &Semantics<RootDatabase>, node: &SyntaxNode) {
+    let (owner, attributes) = match doc_attributes(sema, node) {
+        Some(it) => it,
+        None => return,
+    };
 
-    if !doc_comments.clone().any(|it| it.text().contains(RUSTDOC_FENCE)) {
+    if attributes.docs().map_or(true, |docs| !String::from(docs).contains(RUSTDOC_FENCE)) {
         return;
     }
+    let doc_comments = attributes.by_key("doc").attrs().map(|attr| attr.to_src(&*owner));
 
     let mut inj = Injector::default();
     inj.add_unmapped("fn doctest() {\n");
@@ -102,11 +136,17 @@ pub(super) fn doc_comment(hl: &mut Highlights, node: &SyntaxNode) {
     // spanning comment ranges.
     let mut new_comments = Vec::new();
     for comment in doc_comments {
-        match comment.text().find(RUSTDOC_FENCE) {
+        let (line, range, prefix) = match &comment {
+            Either::Left(_) => continue, // FIXME
+            Either::Right(comment) => {
+                (comment.text(), comment.syntax().text_range(), comment.prefix())
+            }
+        };
+        match line.find(RUSTDOC_FENCE) {
             Some(idx) => {
                 is_codeblock = !is_codeblock;
                 // Check whether code is rust by inspecting fence guards
-                let guards = &comment.text()[idx + RUSTDOC_FENCE.len()..];
+                let guards = &line[idx + RUSTDOC_FENCE.len()..];
                 let is_rust =
                     guards.split(',').all(|sub| RUSTDOC_FENCE_TOKENS.contains(&sub.trim()));
                 is_doctest = is_codeblock && is_rust;
@@ -116,10 +156,7 @@ pub(super) fn doc_comment(hl: &mut Highlights, node: &SyntaxNode) {
             None => (),
         }
 
-        let line: &str = comment.text();
-        let range = comment.syntax().text_range();
-
-        let mut pos = TextSize::of(comment.prefix());
+        let mut pos = TextSize::of(prefix);
         // whitespace after comment is ignored
         if let Some(ws) = line[pos.into()..].chars().next().filter(|c| c.is_whitespace()) {
             pos += TextSize::of(ws);