3 files changed, 119 insertions, 4 deletions

diff --git a/crates/ide/src/syntax_highlighting/injection.rs b/crates/ide/src/syntax_highlighting/injection.rs
index 6cbd683c6..d6be9708d 100644
--- a/crates/ide/src/syntax_highlighting/injection.rs
+++ b/crates/ide/src/syntax_highlighting/injection.rs
@@ -22,7 +22,8 @@ pub(super) fn highlight_injection(
         return None;
     }
     let value = literal.value()?;
-    let (analysis, tmp_file_id) = Analysis::from_single_file(value.into_owned());
+    let marker_info = MarkerInfo::new(&*value);
+    let (analysis, tmp_file_id) = Analysis::from_single_file(marker_info.cleaned_text.clone());
 
     if let Some(range) = literal.open_quote_text_range() {
         acc.add(HighlightedRange {
@@ -33,9 +34,10 @@ pub(super) fn highlight_injection(
     }
 
     for mut h in analysis.highlight(tmp_file_id).unwrap() {
-        if let Some(r) = literal.map_range_up(h.range) {
-            h.range = r;
-            acc.add(h)
+        let range = marker_info.map_range_up(h.range);
+        if let Some(range) = literal.map_range_up(range) {
+            h.range = range;
+            acc.add(h);
         }
     }
 
@@ -50,6 +52,52 @@ pub(super) fn highlight_injection(
     Some(())
 }
 
+/// Data to remove `$0` from string and map ranges
+#[derive(Default, Debug)]
+struct MarkerInfo {
+    cleaned_text: String,
+    markers: Vec<TextRange>,
+}
+
+impl MarkerInfo {
+    fn new(mut text: &str) -> Self {
+        let marker = "$0";
+
+        let mut res = MarkerInfo::default();
+        let mut offset: TextSize = 0.into();
+        while !text.is_empty() {
+            let idx = text.find(marker).unwrap_or(text.len());
+            let (chunk, next) = text.split_at(idx);
+            text = next;
+            res.cleaned_text.push_str(chunk);
+            offset += TextSize::of(chunk);
+
+            if let Some(next) = text.strip_prefix(marker) {
+                text = next;
+
+                let marker_len = TextSize::of(marker);
+                res.markers.push(TextRange::at(offset, marker_len));
+                offset += marker_len;
+            }
+        }
+        res
+    }
+    fn map_range_up(&self, range: TextRange) -> TextRange {
+        TextRange::new(
+            self.map_offset_up(range.start(), true),
+            self.map_offset_up(range.end(), false),
+        )
+    }
+    fn map_offset_up(&self, mut offset: TextSize, start: bool) -> TextSize {
+        for r in &self.markers {
+            if r.start() < offset || (start && r.start() == offset) {
+                offset += r.len()
+            }
+        }
+        offset
+    }
+}
+
 /// Mapping from extracted documentation code to original code
 type RangesMap = BTreeMap<TextSize, TextSize>;
 
diff --git a/crates/ide/src/syntax_highlighting/test_data/injection.html b/crates/ide/src/syntax_highlighting/test_data/injection.html
new file mode 100644
index 000000000..a54d303b4
--- /dev/null
+++ b/crates/ide/src/syntax_highlighting/test_data/injection.html
@@ -0,0 +1,48 @@
+
+<style>
+body                { margin: 0; }
+pre                 { color: #DCDCCC; background: #3F3F3F; font-size: 22px; padding: 0.4em; }
+
+.lifetime           { color: #DFAF8F; font-style: italic; }
+.label              { color: #DFAF8F; font-style: italic; }
+.comment            { color: #7F9F7F; }
+.documentation      { color: #629755; }
+.injected           { opacity: 0.65 ; }
+.struct, .enum      { color: #7CB8BB; }
+.enum_variant       { color: #BDE0F3; }
+.string_literal     { color: #CC9393; }
+.field              { color: #94BFF3; }
+.function           { color: #93E0E3; }
+.function.unsafe    { color: #BC8383; }
+.operator.unsafe    { color: #BC8383; }
+.parameter          { color: #94BFF3; }
+.text               { color: #DCDCCC; }
+.type               { color: #7CB8BB; }
+.builtin_type       { color: #8CD0D3; }
+.type_param         { color: #DFAF8F; }
+.attribute          { color: #94BFF3; }
+.numeric_literal    { color: #BFEBBF; }
+.bool_literal       { color: #BFE6EB; }
+.macro              { color: #94BFF3; }
+.module             { color: #AFD8AF; }
+.value_param        { color: #DCDCCC; }
+.variable           { color: #DCDCCC; }
+.format_specifier   { color: #CC696B; }
+.mutable            { text-decoration: underline; }
+.escape_sequence    { color: #94BFF3; }
+.keyword            { color: #F0DFAF; font-weight: bold; }
+.keyword.unsafe     { color: #BC8383; font-weight: bold; }
+.control            { font-style: italic; }
+
+.unresolved_reference { color: #FC5555; text-decoration: wavy underline; }
+</style>
+<pre><code><span class="keyword">fn</span> <span class="function declaration">f</span><span class="punctuation">(</span><span class="value_param declaration">ra_fixture</span><span class="punctuation">:</span> <span class="operator">&</span><span class="builtin_type">str</span><span class="punctuation">)</span> <span class="punctuation">{</span><span class="punctuation">}</span>
+<span class="keyword">fn</span> <span class="function declaration">main</span><span class="punctuation">(</span><span class="punctuation">)</span> <span class="punctuation">{</span>
+    <span class="function">f</span><span class="punctuation">(</span><span class="string_literal">r"</span>
+<span class="keyword">fn</span> <span class="function declaration">foo</span><span class="punctuation">(</span><span class="punctuation">)</span> <span class="punctuation">{</span>
+    <span class="function">foo</span><span class="punctuation">(</span>$0<span class="punctuation">{</span>
+        <span class="numeric_literal">92</span>
+    <span class="punctuation">}</span>$0<span class="punctuation">)</span>
+<span class="punctuation">}</span><span class="string_literal">"</span><span class="punctuation">)</span><span class="punctuation">;</span>
+<span class="punctuation">}</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 30b5b648e..9e1a3974c 100644
--- a/crates/ide/src/syntax_highlighting/tests.rs
+++ b/crates/ide/src/syntax_highlighting/tests.rs
@@ -555,6 +555,25 @@ impl t for foo {
     )
 }
 
+#[test]
+fn test_injection() {
+    check_highlighting(
+        r##"
+fn f(ra_fixture: &str) {}
+fn main() {
+    f(r"
+fn foo() {
+    foo(\$0{
+        92
+    }\$0)
+}");
+}
+    "##,
+        expect_file!["./test_data/injection.html"],
+        false,
+    );
+}
+
 /// Highlights the code given by the `ra_fixture` argument, renders the
 /// result as HTML, and compares it with the HTML file given as `snapshot`.
 /// Note that the `snapshot` file is overwritten by the rendered HTML.