Inject highlight into block doc comments

3 files changed, 72 insertions, 28 deletions

diff --git a/crates/ide/src/syntax_highlighting/inject.rs b/crates/ide/src/syntax_highlighting/inject.rs
index d57ce4027..7a4f2645f 100644
--- a/crates/ide/src/syntax_highlighting/inject.rs
+++ b/crates/ide/src/syntax_highlighting/inject.rs
@@ -1,5 +1,7 @@
 //! "Recursive" Syntax highlighting for code in doctests and fixtures.
 
+use std::mem;
+
 use either::Either;
 use hir::{HasAttrs, Semantics};
 use ide_db::call_info::ActiveParameter;
@@ -186,34 +188,44 @@ pub(super) fn doc_comment(hl: &mut Highlights, sema: &Semantics<RootDatabase>, n
             }
         };
 
-        match line.find(RUSTDOC_FENCE) {
-            Some(idx) => {
-                is_codeblock = !is_codeblock;
-                // Check whether code is rust by inspecting fence guards
-                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;
-                continue;
+        let mut pos = TextSize::from(prefix.len() as u32);
+        let mut range_start = range.start();
+        for line in line.split('\n') {
+            let line_len = TextSize::from(line.len() as u32);
+            let prev_range_start = {
+                let next_range_start = range_start + line_len + TextSize::from(1);
+                mem::replace(&mut range_start, next_range_start)
+            };
+            // only first line has the prefix so take it away for future iterations
+            let mut pos = mem::take(&mut pos);
+
+            match line.find(RUSTDOC_FENCE) {
+                Some(idx) => {
+                    is_codeblock = !is_codeblock;
+                    // Check whether code is rust by inspecting fence guards
+                    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;
+                    continue;
+                }
+                None if !is_doctest => continue,
+                None => (),
             }
-            None if !is_doctest => continue,
-            None => (),
-        }
-
-        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);
-        }
-        // lines marked with `#` should be ignored in output, we skip the `#` char
-        if let Some(ws) = line[pos.into()..].chars().next().filter(|&c| c == '#') {
-            pos += TextSize::of(ws);
-        }
 
-        new_comments.push(TextRange::at(range.start(), pos));
+            // whitespace after comment is ignored
+            if let Some(ws) = line[pos.into()..].chars().next().filter(|c| c.is_whitespace()) {
+                pos += TextSize::of(ws);
+            }
+            // lines marked with `#` should be ignored in output, we skip the `#` char
+            if line[pos.into()..].starts_with('#') {
+                pos += TextSize::of('#');
+            }
 
-        inj.add(&line[pos.into()..], TextRange::new(range.start() + pos, range.end()));
-        inj.add_unmapped("\n");
+            new_comments.push(TextRange::at(prev_range_start, pos));
+            inj.add(&line[pos.into()..], TextRange::new(pos, line_len) + prev_range_start);
+            inj.add_unmapped("\n");
+        }
     }
     inj.add_unmapped("\n}");
 
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 45817faf9..d792a23cf 100644
--- a/crates/ide/src/syntax_highlighting/test_data/highlight_doctest.html
+++ b/crates/ide/src/syntax_highlighting/test_data/highlight_doctest.html
@@ -81,7 +81,7 @@ pre                 { color: #DCDCCC; background: #3F3F3F; font-size: 22px; padd
     <span class="comment documentation">/// </span><span class="comment injected">       comment */</span>
     <span class="comment documentation">///</span>
     <span class="comment documentation">/// </span><span class="keyword injected">let</span><span class="none injected"> </span><span class="variable declaration injected">multi_line_string</span><span class="none injected"> </span><span class="operator injected">=</span><span class="none injected"> </span><span class="string_literal injected">"Foo</span>
-    <span class="comment documentation">/// </span><span class="string_literal injected">  bar</span>
+    <span class="comment documentation">/// </span><span class="string_literal injected">  bar</span><span class="escape_sequence injected">\n</span>
     <span class="comment documentation">/// </span><span class="string_literal injected">         "</span><span class="semicolon injected">;</span>
     <span class="comment documentation">///</span>
     <span class="comment documentation">/// ```</span>
@@ -121,4 +121,20 @@ pre                 { color: #DCDCCC; background: #3F3F3F; font-size: 22px; padd
 <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
+<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>
+
+<span class="comment documentation">/**
+It is beyond me why you'd use these when you got ///
+```rust
+</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="keyword">pub</span> <span class="keyword">fn</span> <span class="function declaration">block_comments</span><span class="parenthesis">(</span><span class="parenthesis">)</span> <span class="brace">{</span><span class="brace">}</span>
+
+<span class="comment documentation">/**
+    Really, I don't get it
+    ```rust
+</span><span class="comment documentation"> </span><span class="none injected">   </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="keyword">pub</span> <span class="keyword">fn</span> <span class="function declaration">block_comments2</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 a5ef2d29b..cf0b86ad0 100644
--- a/crates/ide/src/syntax_highlighting/tests.rs
+++ b/crates/ide/src/syntax_highlighting/tests.rs
@@ -516,7 +516,7 @@ impl Foo {
     ///        comment */
     ///
     /// let multi_line_string = "Foo
-    ///   bar
+    ///   bar\n
     ///          ";
     ///
     /// ```
@@ -557,6 +557,22 @@ macro_rules! noop {
 /// let _ = example(&alloc::vec![1, 2, 3]);
 /// ```
 pub fn mix_and_match() {}
+
+/**
+It is beyond me why you'd use these when you got ///
+```rust
+let _ = example(&[1, 2, 3]);
+```
+ */
+pub fn block_comments() {}
+
+/**
+    Really, I don't get it
+    ```rust
+    let _ = example(&[1, 2, 3]);
+    ```
+*/
+pub fn block_comments2() {}
 "#
         .trim(),
         expect_file!["./test_data/highlight_doctest.html"],