10 files changed, 205 insertions, 27 deletions

diff --git a/crates/hir/src/attrs.rs b/crates/hir/src/attrs.rs
index b9c695921..dab8da7bb 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/lib.rs b/crates/hir/src/lib.rs
index 861b7329e..b41a36a78 100644
--- a/crates/hir/src/lib.rs
+++ b/crates/hir/src/lib.rs
@@ -89,7 +89,7 @@ pub use crate::{
 pub use {
     hir_def::{
         adt::StructKind,
-        attr::{Attrs, Documentation},
+        attr::{Attr, Attrs, Documentation},
         body::scope::ExprScopes,
         find_path::PrefixKind,
         import_map,
diff --git a/crates/hir/src/semantics.rs b/crates/hir/src/semantics.rs
index 00b076175..e0eb2a66d 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 b0b4b5052..7ba53ee5c 100644
--- a/crates/hir_def/src/attr.rs
+++ b/crates/hir_def/src/attr.rs
@@ -136,16 +136,15 @@ impl RawAttrs {
         let new_attrs = self
             .iter()
             .flat_map(|attr| -> SmallVec<[_; 1]> {
-                let attr = attr.clone();
                 let is_cfg_attr =
                     attr.path.as_ident().map_or(false, |name| *name == hir_expand::name![cfg_attr]);
                 if !is_cfg_attr {
-                    return smallvec![attr];
+                    return smallvec![attr.clone()];
                 }
 
                 let subtree = match &attr.input {
                     Some(AttrInput::TokenTree(it)) => it,
-                    _ => return smallvec![attr],
+                    _ => return smallvec![attr.clone()],
                 };
 
                 // Input subtree is: `(cfg, $(attr),+)`
@@ -157,11 +156,13 @@ impl RawAttrs {
                 let cfg = parts.next().unwrap();
                 let cfg = Subtree { delimiter: subtree.delimiter, token_trees: cfg.to_vec() };
                 let cfg = CfgExpr::parse(&cfg);
+                let index = attr.index;
                 let attrs = parts.filter(|a| !a.is_empty()).filter_map(|attr| {
                     let tree = Subtree { delimiter: None, token_trees: attr.to_vec() };
                     let attr = ast::Attr::parse(&format!("#[{}]", tree)).ok()?;
-                    let hygiene = Hygiene::new_unhygienic(); // FIXME
-                    Attr::from_src(attr, &hygiene)
+                    // FIXME hygiene
+                    let hygiene = Hygiene::new_unhygienic();
+                    Attr::from_src(attr, &hygiene).map(|attr| Attr { index, ..attr })
                 });
 
                 let cfg_options = &crate_graph[krate].cfg_options;
@@ -293,6 +294,13 @@ impl Attrs {
         Arc::new(res)
     }
 
+    /// 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 AttrsOwner) -> AttrSourceMap {
+        AttrSourceMap { attrs: collect_attrs(owner).collect() }
+    }
+
     pub fn by_key(&self, key: &'static str) -> AttrQuery<'_> {
         AttrQuery { attrs: self, key }
     }
@@ -365,6 +373,24 @@ fn inner_attributes(
     Some((attrs, docs))
 }
 
+pub struct AttrSourceMap {
+    attrs: Vec<Either<ast::Attr, ast::Comment>>,
+}
+
+impl AttrSourceMap {
+    /// 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) -> &Either<ast::Attr, ast::Comment> {
+        self.attrs
+            .get(attr.index as usize)
+            .unwrap_or_else(|| panic!("cannot find `Attr` at index {}", attr.index))
+    }
+}
+
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct Attr {
     index: u32,
@@ -448,6 +474,13 @@ impl Attr {
             _ => None,
         }
     }
+
+    pub fn string_value(&self) -> Option<&SmolStr> {
+        match self.input.as_ref()? {
+            AttrInput::Literal(it) => Some(it),
+            _ => None,
+        }
+    }
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -475,7 +508,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 4f825523c..d57ce4027 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, DocCommentsOwner},
+    match_ast, AstToken, NodeOrToken, SyntaxNode, SyntaxToken, TextRange, TextSize,
+};
 
 use crate::{Analysis, HlMod, HlRange, HlTag, RootDatabase};
 
@@ -81,16 +85,75 @@ const RUSTDOC_FENCE_TOKENS: &[&'static str] = &[
     "edition2021",
 ];
 
+// Basically an owned dyn AttrsOwner without extra Boxing
+struct AttrsOwnerNode {
+    node: SyntaxNode,
+}
+
+impl AttrsOwnerNode {
+    fn new<N: DocCommentsOwner>(node: N) -> Self {
+        AttrsOwnerNode { node: node.syntax().clone() }
+    }
+}
+
+impl AttrsOwner for AttrsOwnerNode {}
+impl AstNode for AttrsOwnerNode {
+    fn can_cast(_: syntax::SyntaxKind) -> bool
+    where
+        Self: Sized,
+    {
+        false
+    }
+    fn cast(_: SyntaxNode) -> Option<Self>
+    where
+        Self: Sized,
+    {
+        None
+    }
+    fn syntax(&self) -> &SyntaxNode {
+        &self.node
+    }
+}
+
+fn doc_attributes<'node>(
+    sema: &Semantics<RootDatabase>,
+    node: &'node SyntaxNode,
+) -> Option<(AttrsOwnerNode, hir::Attrs)> {
+    match_ast! {
+        match node {
+            ast::SourceFile(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db))),
+            ast::Fn(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db))),
+            ast::Struct(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db))),
+            ast::Union(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db))),
+            ast::RecordField(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db))),
+            ast::TupleField(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db))),
+            ast::Enum(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db))),
+            ast::Variant(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db))),
+            ast::Trait(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db))),
+            ast::Module(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db))),
+            ast::Static(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db))),
+            ast::Const(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db))),
+            ast::TypeAlias(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db))),
+            ast::Impl(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), def.attrs(sema.db))),
+            ast::MacroRules(it) => sema.to_def(&it).map(|def| (AttrsOwnerNode::new(it), 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 attrs_source_map = attributes.source_map(&owner);
 
     let mut inj = Injector::default();
     inj.add_unmapped("fn doctest() {\n");
@@ -101,12 +164,33 @@ pub(super) fn doc_comment(hl: &mut Highlights, node: &SyntaxNode) {
     // Replace the original, line-spanning comment ranges by new, only comment-prefix
     // spanning comment ranges.
     let mut new_comments = Vec::new();
-    for comment in doc_comments {
-        match comment.text().find(RUSTDOC_FENCE) {
+    let mut string;
+    for attr in attributes.by_key("doc").attrs() {
+        let src = attrs_source_map.source_of(&attr);
+        let (line, range, prefix) = match &src {
+            Either::Left(it) => {
+                string = match find_doc_string_in_attr(attr, it) {
+                    Some(it) => it,
+                    None => continue,
+                };
+                let text_range = string.syntax().text_range();
+                let text_range = TextRange::new(
+                    text_range.start() + TextSize::from(1),
+                    text_range.end() - TextSize::from(1),
+                );
+                let text = string.text();
+                (&text[1..text.len() - 1], text_range, "")
+            }
+            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 +200,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);
@@ -156,3 +237,27 @@ pub(super) fn doc_comment(hl: &mut Highlights, node: &SyntaxNode) {
         });
     }
 }
+
+fn find_doc_string_in_attr(attr: &hir::Attr, it: &ast::Attr) -> Option<ast::String> {
+    match it.literal() {
+        // #[doc = lit]
+        Some(lit) => match lit.kind() {
+            ast::LiteralKind::String(it) => Some(it),
+            _ => None,
+        },
+        // #[cfg_attr(..., doc = "", ...)]
+        None => {
+            // We gotta hunt the string token manually here
+            let text = attr.string_value()?;
+            // FIXME: We just pick the first string literal that has the same text as the doc attribute
+            // This means technically we might highlight the wrong one
+            it.syntax()
+                .descendants_with_tokens()
+                .filter_map(NodeOrToken::into_token)
+                .filter_map(ast::String::cast)
+                .find(|string| {
+                    string.text().get(1..string.text().len() - 1).map_or(false, |it| it == text)
+                })
+        }
+    }
+}
diff --git a/crates/ide/src/syntax_highlighting/test_data/highlight_doctest.html b/crates/ide/src/syntax_highlighting/test_data/highlight_doctest.html
index 5e877df88..45817faf9 100644
--- a/crates/ide/src/syntax_highlighting/test_data/highlight_doctest.html
+++ b/crates/ide/src/syntax_highlighting/test_data/highlight_doctest.html
@@ -105,4 +105,20 @@ pre                 { color: #DCDCCC; background: #3F3F3F; font-size: 22px; padd
     <span class="parenthesis">(</span><span class="punctuation">$</span>expr<span class="colon">:</span>expr<span class="parenthesis">)</span> <span class="operator">=</span><span class="angle">&gt;</span> <span class="brace">{</span>
         <span class="punctuation">$</span>expr
     <span class="brace">}</span>
-<span class="brace">}</span></code></pre>
\ No newline at end of file
+<span class="brace">}</span>
+
+<span class="comment documentation">/// ```rust</span>
+<span class="comment documentation">/// </span><span class="keyword injected">let</span><span class="none injected"> </span><span class="punctuation injected">_</span><span class="none injected"> </span><span class="operator injected">=</span><span class="none injected"> </span><span class="function injected">example</span><span class="parenthesis injected">(</span><span class="operator injected">&</span><span class="bracket injected">[</span><span class="numeric_literal injected">1</span><span class="comma injected">,</span><span class="none injected"> </span><span class="numeric_literal injected">2</span><span class="comma injected">,</span><span class="none injected"> </span><span class="numeric_literal injected">3</span><span class="bracket injected">]</span><span class="parenthesis injected">)</span><span class="semicolon injected">;</span>
+<span class="comment documentation">/// ```</span>
+<span class="comment documentation">///</span>
+<span class="comment documentation">/// ```</span>
+<span class="comment documentation">/// </span><span class="keyword control injected">loop</span><span class="none injected"> </span><span class="brace injected">{</span><span class="brace injected">}</span>
+<span class="attribute attribute">#</span><span class="attribute attribute">[</span><span class="function attribute">cfg_attr</span><span class="parenthesis attribute">(</span><span class="attribute attribute">not</span><span class="parenthesis attribute">(</span><span class="attribute attribute">feature </span><span class="operator attribute">=</span><span class="attribute attribute"> </span><span class="string_literal attribute">"false"</span><span class="parenthesis attribute">)</span><span class="comma attribute">,</span><span class="attribute attribute"> doc </span><span class="operator attribute">=</span><span class="attribute attribute"> </span><span class="string_literal attribute">"</span><span class="keyword control injected">loop</span><span class="none injected"> </span><span class="brace injected">{</span><span class="brace injected">}</span><span class="string_literal attribute">"</span><span class="parenthesis attribute">)</span><span class="attribute attribute">]</span>
+<span class="attribute attribute">#</span><span class="attribute attribute">[</span><span class="function attribute">doc</span><span class="attribute attribute"> </span><span class="operator attribute">=</span><span class="attribute attribute"> </span><span class="string_literal attribute">"</span><span class="keyword control injected">loop</span><span class="none injected"> </span><span class="brace injected">{</span><span class="brace injected">}</span><span class="string_literal attribute">"</span><span class="attribute attribute">]</span>
+<span class="comment documentation">/// ```</span>
+<span class="comment documentation">///</span>
+<span class="attribute attribute">#</span><span class="attribute attribute">[</span><span class="function attribute">cfg_attr</span><span class="parenthesis attribute">(</span><span class="attribute attribute">feature </span><span class="operator attribute">=</span><span class="attribute attribute"> </span><span class="string_literal attribute">"alloc"</span><span class="comma attribute">,</span><span class="attribute attribute"> doc </span><span class="operator attribute">=</span><span class="attribute attribute"> </span><span class="string_literal attribute">"```rust"</span><span class="parenthesis attribute">)</span><span class="attribute attribute">]</span>
+<span class="attribute attribute">#</span><span class="attribute attribute">[</span><span class="function attribute">cfg_attr</span><span class="parenthesis attribute">(</span><span class="attribute attribute">not</span><span class="parenthesis attribute">(</span><span class="attribute attribute">feature </span><span class="operator attribute">=</span><span class="attribute attribute"> </span><span class="string_literal attribute">"alloc"</span><span class="parenthesis attribute">)</span><span class="comma attribute">,</span><span class="attribute attribute"> doc </span><span class="operator attribute">=</span><span class="attribute attribute"> </span><span class="string_literal attribute">"```ignore"</span><span class="parenthesis attribute">)</span><span class="attribute attribute">]</span>
+<span class="comment documentation">/// </span><span class="keyword injected">let</span><span class="none injected"> </span><span class="punctuation injected">_</span><span class="none injected"> </span><span class="operator injected">=</span><span class="none injected"> </span><span class="function injected">example</span><span class="parenthesis injected">(</span><span class="operator injected">&</span><span class="none injected">alloc::</span><span class="macro injected">vec!</span><span class="bracket injected">[</span><span class="numeric_literal injected">1</span><span class="comma injected">,</span><span class="none injected"> </span><span class="numeric_literal injected">2</span><span class="comma injected">,</span><span class="none injected"> </span><span class="numeric_literal injected">3</span><span class="bracket injected">]</span><span class="parenthesis injected">)</span><span class="semicolon injected">;</span>
+<span class="comment documentation">/// ```</span>
+<span class="keyword">pub</span> <span class="keyword">fn</span> <span class="function declaration">mix_and_match</span><span class="parenthesis">(</span><span class="parenthesis">)</span> <span class="brace">{</span><span class="brace">}</span></code></pre>
\ No newline at end of file
diff --git a/crates/ide/src/syntax_highlighting/tests.rs b/crates/ide/src/syntax_highlighting/tests.rs
index 9d0cd1af5..a5ef2d29b 100644
--- a/crates/ide/src/syntax_highlighting/tests.rs
+++ b/crates/ide/src/syntax_highlighting/tests.rs
@@ -541,6 +541,22 @@ macro_rules! noop {
         $expr
     }
 }
+
+/// ```rust
+/// let _ = example(&[1, 2, 3]);
+/// ```
+///
+/// ```
+/// loop {}
+#[cfg_attr(not(feature = "false"), doc = "loop {}")]
+#[doc = "loop {}"]
+/// ```
+///
+#[cfg_attr(feature = "alloc", doc = "```rust")]
+#[cfg_attr(not(feature = "alloc"), doc = "```ignore")]
+/// let _ = example(&alloc::vec![1, 2, 3]);
+/// ```
+pub fn mix_and_match() {}
 "#
         .trim(),
         expect_file!["./test_data/highlight_doctest.html"],
diff --git a/crates/syntax/src/ast/traits.rs b/crates/syntax/src/ast/traits.rs
index 13a769d51..96d4cc997 100644
--- a/crates/syntax/src/ast/traits.rs
+++ b/crates/syntax/src/ast/traits.rs
@@ -72,7 +72,7 @@ pub trait AttrsOwner: AstNode {
     }
 }
 
-pub trait DocCommentsOwner: AstNode {
+pub trait DocCommentsOwner: AttrsOwner {
     fn doc_comments(&self) -> CommentIter {
         CommentIter { iter: self.syntax().children_with_tokens() }
     }