Simplify highlighting infra

This also fixes the killer whale bug

crates/ide/src/syntax_highlighting/injection.rs

@@ -1,18 +1,18 @@
 //! Syntax highlighting injections such as highlighting of documentation tests.
 
-use std::{collections::BTreeMap, convert::TryFrom};
+use std::convert::TryFrom;
 
 use hir::Semantics;
 use ide_db::call_info::ActiveParameter;
 use itertools::Itertools;
 use syntax::{ast, AstToken, SyntaxNode, SyntaxToken, TextRange, TextSize};
 
-use crate::{Analysis, Highlight, HighlightModifier, HighlightTag, HighlightedRange, RootDatabase};
+use crate::{Analysis, HighlightModifier, HighlightTag, HighlightedRange, RootDatabase};
 
-use super::HighlightedRangeStack;
+use super::{highlights::Highlights, injector::Injector};
 
 pub(super) fn highlight_injection(
-    acc: &mut HighlightedRangeStack,
+    acc: &mut Highlights,
     sema: &Semantics<RootDatabase>,
     literal: ast::String,
     expanded: SyntaxToken,
@@ -98,9 +98,6 @@ impl MarkerInfo {
     }
 }
 
-/// Mapping from extracted documentation code to original code
-type RangesMap = BTreeMap<TextSize, TextSize>;
-
 const RUSTDOC_FENCE: &'static str = "```";
 const RUSTDOC_FENCE_TOKENS: &[&'static str] = &[
     "",
@@ -119,20 +116,20 @@ const RUSTDOC_FENCE_TOKENS: &[&'static str] = &[
 /// Lastly, a vector of new comment highlight ranges (spanning only the
 /// comment prefix) is returned which is used in the syntax highlighting
 /// injection to replace the previous (line-spanning) comment ranges.
-pub(super) fn extract_doc_comments(
-    node: &SyntaxNode,
-) -> Option<(String, RangesMap, Vec<HighlightedRange>)> {
+pub(super) fn extract_doc_comments(node: &SyntaxNode) -> Option<(Vec<HighlightedRange>, Injector)> {
+    let mut inj = Injector::default();
     // wrap the doctest into function body to get correct syntax highlighting
     let prefix = "fn doctest() {\n";
     let suffix = "}\n";
-    // Mapping from extracted documentation code to original code
-    let mut range_mapping: RangesMap = BTreeMap::new();
-    let mut line_start = TextSize::try_from(prefix.len()).unwrap();
+
+    let mut line_start = TextSize::of(prefix);
     let mut is_codeblock = false;
     let mut is_doctest = false;
     // Replace the original, line-spanning comment ranges by new, only comment-prefix
     // spanning comment ranges.
     let mut new_comments = Vec::new();
+
+    inj.add_unmapped(prefix);
     let doctest = node
         .children_with_tokens()
         .filter_map(|el| el.into_token().and_then(ast::Comment::cast))
@@ -169,7 +166,6 @@ pub(super) fn extract_doc_comments(
                 pos
             };
 
-            range_mapping.insert(line_start, range.start() + TextSize::try_from(pos).unwrap());
             new_comments.push(HighlightedRange {
                 range: TextRange::new(
                     range.start(),
@@ -179,62 +175,43 @@ pub(super) fn extract_doc_comments(
                 binding_hash: None,
             });
             line_start += range.len() - TextSize::try_from(pos).unwrap();
-            line_start += TextSize::try_from('\n'.len_utf8()).unwrap();
+            line_start += TextSize::of("\n");
 
+            inj.add(
+                &line[pos..],
+                TextRange::new(range.start() + TextSize::try_from(pos).unwrap(), range.end()),
+            );
+            inj.add_unmapped("\n");
             line[pos..].to_owned()
         })
         .join("\n");
+    inj.add_unmapped(suffix);
 
     if doctest.is_empty() {
         return None;
     }
 
-    let doctest = format!("{}{}{}", prefix, doctest, suffix);
-    Some((doctest, range_mapping, new_comments))
+    Some((new_comments, inj))
 }
 
 /// Injection of syntax highlighting of doctests.
 pub(super) fn highlight_doc_comment(
-    text: String,
-    range_mapping: RangesMap,
     new_comments: Vec<HighlightedRange>,
-    stack: &mut HighlightedRangeStack,
+    inj: Injector,
+    stack: &mut Highlights,
 ) {
-    let (analysis, tmp_file_id) = Analysis::from_single_file(text);
-
-    stack.push();
-    for mut h in analysis.with_db(|db| super::highlight(db, tmp_file_id, None, true)).unwrap() {
-        // Determine start offset and end offset in case of multi-line ranges
-        let mut start_offset = None;
-        let mut end_offset = None;
-        for (line_start, orig_line_start) in range_mapping.range(..h.range.end()).rev() {
-            // It's possible for orig_line_start - line_start to be negative. Add h.range.start()
-            // here and remove it from the end range after the loop below so that the values are
-            // always non-negative.
-            let offset = h.range.start() + orig_line_start - line_start;
-            if line_start <= &h.range.start() {
-                start_offset.get_or_insert(offset);
-                break;
-            } else {
-                end_offset.get_or_insert(offset);
-            }
-        }
-        if let Some(start_offset) = start_offset {
-            h.range = TextRange::new(
-                start_offset,
-                h.range.end() + end_offset.unwrap_or(start_offset) - h.range.start(),
-            );
-
-            h.highlight |= HighlightModifier::Injected;
-            stack.add(h);
-        }
-    }
-
-    // Inject the comment prefix highlight ranges
-    stack.push();
+    let (analysis, tmp_file_id) = Analysis::from_single_file(inj.text().to_string());
     for comment in new_comments {
         stack.add(comment);
     }
-    stack.pop_and_inject(None);
-    stack.pop_and_inject(Some(Highlight::from(HighlightTag::Dummy) | HighlightModifier::Injected));
+
+    for h in analysis.with_db(|db| super::highlight(db, tmp_file_id, None, true)).unwrap() {
+        for r in inj.map_range_up(h.range) {
+            stack.add(HighlightedRange {
+                range: r,
+                highlight: h.highlight | HighlightModifier::Injected,
+                binding_hash: h.binding_hash,
+            });
+        }
+    }
 }