Implement inlay highlighting

crates/editor/src/display_map.rs

@@ -5,11 +5,11 @@ mod tab_map;
 mod wrap_map;
 
 use crate::{
-    link_go_to_definition::InlayRange, Anchor, AnchorRangeExt, InlayBackgroundHighlight, InlayId,
+    link_go_to_definition::InlayRange, Anchor, AnchorRangeExt, InlayId, MultiBuffer,
-    MultiBuffer, MultiBufferSnapshot, ToOffset, ToPoint,
+    MultiBufferSnapshot, ToOffset, ToPoint,
 };
 pub use block_map::{BlockMap, BlockPoint};
-use collections::{BTreeMap, HashMap, HashSet};
+use collections::{HashMap, HashSet};
 use fold_map::FoldMap;
 use gpui::{
     color::Color,
@@ -304,6 +304,16 @@ impl DisplayMap {
     }
 }
 
+#[derive(Debug, Default)]
+pub struct Highlights<'a> {
+    pub text_highlights: Option<&'a TextHighlights>,
+    pub inlay_highlights: Option<&'a InlayHighlights>,
+    pub inlay_background_highlights:
+        Option<TreeMap<Option<TypeId>, Arc<(HighlightStyle, &'a [InlayRange])>>>,
+    pub inlay_highlight_style: Option<HighlightStyle>,
+    pub suggestion_highlight_style: Option<HighlightStyle>,
+}
+
 pub struct DisplaySnapshot {
     pub buffer_snapshot: MultiBufferSnapshot,
     pub fold_snapshot: fold_map::FoldSnapshot,
@@ -316,15 +326,6 @@ pub struct DisplaySnapshot {
     clip_at_line_ends: bool,
 }
 
-#[derive(Debug, Default)]
-pub struct Highlights<'a> {
-    pub text_highlights: Option<&'a TextHighlights>,
-    pub inlay_highlights: Option<&'a InlayHighlights>,
-    pub inlay_background_highlights: Option<&'a BTreeMap<TypeId, InlayBackgroundHighlight>>,
-    pub inlay_highlight_style: Option<HighlightStyle>,
-    pub suggestion_highlight_style: Option<HighlightStyle>,
-}
-
 impl DisplaySnapshot {
     #[cfg(test)]
     pub fn fold_count(&self) -> usize {
@@ -480,7 +481,9 @@ impl DisplaySnapshot {
         &'a self,
         display_rows: Range<u32>,
         language_aware: bool,
-        inlay_background_highlights: Option<&'a BTreeMap<TypeId, InlayBackgroundHighlight>>,
+        inlay_background_highlights: Option<
+            TreeMap<Option<TypeId>, Arc<(HighlightStyle, &'a [InlayRange])>>,
+        >,
         inlay_highlight_style: Option<HighlightStyle>,
         suggestion_highlight_style: Option<HighlightStyle>,
     ) -> DisplayChunks<'_> {

crates/editor/src/display_map/inlay_map.rs

@@ -1,4 +1,5 @@
 use crate::{
+    link_go_to_definition::InlayRange,
     multi_buffer::{MultiBufferChunks, MultiBufferRows},
     Anchor, InlayId, MultiBufferSnapshot, ToOffset,
 };
@@ -10,22 +11,23 @@ use std::{
     cmp,
     iter::Peekable,
     ops::{Add, AddAssign, Range, Sub, SubAssign},
+    sync::Arc,
     vec,
 };
-use sum_tree::{Bias, Cursor, SumTree};
+use sum_tree::{Bias, Cursor, SumTree, TreeMap};
 use text::{Patch, Rope};
 
 use super::Highlights;
 
 pub struct InlayMap {
     snapshot: InlaySnapshot,
-    inlays: Vec<Inlay>,
 }
 
 #[derive(Clone)]
 pub struct InlaySnapshot {
     pub buffer: MultiBufferSnapshot,
     transforms: SumTree<Transform>,
+    inlays: Vec<Inlay>,
     pub version: usize,
 }
 
@@ -399,13 +401,13 @@ impl InlayMap {
         let snapshot = InlaySnapshot {
             buffer: buffer.clone(),
             transforms: SumTree::from_iter(Some(Transform::Isomorphic(buffer.text_summary())), &()),
+            inlays: Vec::new(),
             version,
         };
 
         (
             Self {
                 snapshot: snapshot.clone(),
-                inlays: Vec::new(),
             },
             snapshot,
         )
@@ -474,7 +476,7 @@ impl InlayMap {
                 );
                 let new_start = InlayOffset(new_transforms.summary().output.len);
 
-                let start_ix = match self.inlays.binary_search_by(|probe| {
+                let start_ix = match snapshot.inlays.binary_search_by(|probe| {
                     probe
                         .position
                         .to_offset(&buffer_snapshot)
@@ -484,7 +486,7 @@ impl InlayMap {
                     Ok(ix) | Err(ix) => ix,
                 };
 
-                for inlay in &self.inlays[start_ix..] {
+                for inlay in &snapshot.inlays[start_ix..] {
                     let buffer_offset = inlay.position.to_offset(&buffer_snapshot);
                     if buffer_offset > buffer_edit.new.end {
                         break;
@@ -554,7 +556,7 @@ impl InlayMap {
         let snapshot = &mut self.snapshot;
         let mut edits = BTreeSet::new();
 
-        self.inlays.retain(|inlay| {
+        snapshot.inlays.retain(|inlay| {
             let retain = !to_remove.contains(&inlay.id);
             if !retain {
                 let offset = inlay.position.to_offset(&snapshot.buffer);
@@ -570,13 +572,13 @@ impl InlayMap {
             }
 
             let offset = inlay_to_insert.position.to_offset(&snapshot.buffer);
-            match self.inlays.binary_search_by(|probe| {
+            match snapshot.inlays.binary_search_by(|probe| {
                 probe
                     .position
                     .cmp(&inlay_to_insert.position, &snapshot.buffer)
             }) {
                 Ok(ix) | Err(ix) => {
-                    self.inlays.insert(ix, inlay_to_insert);
+                    snapshot.inlays.insert(ix, inlay_to_insert);
                 }
             }
 
@@ -596,7 +598,7 @@ impl InlayMap {
     }
 
     pub fn current_inlays(&self) -> impl Iterator<Item = &Inlay> {
-        self.inlays.iter()
+        self.snapshot.inlays.iter()
     }
 
     #[cfg(test)]
@@ -612,7 +614,7 @@ impl InlayMap {
         let mut to_insert = Vec::new();
         let snapshot = &mut self.snapshot;
         for i in 0..rng.gen_range(1..=5) {
-            if self.inlays.is_empty() || rng.gen() {
+            if snapshot.inlays.is_empty() || rng.gen() {
                 let position = snapshot.buffer.random_byte_range(0, rng).start;
                 let bias = if rng.gen() { Bias::Left } else { Bias::Right };
                 let len = if rng.gen_bool(0.01) {
@@ -643,7 +645,8 @@ impl InlayMap {
                 });
             } else {
                 to_remove.push(
-                    self.inlays
+                    snapshot
+                        .inlays
                         .iter()
                         .choose(rng)
                         .map(|inlay| inlay.id)
@@ -997,61 +1000,50 @@ impl InlaySnapshot {
         cursor.seek(&range.start, Bias::Right, &());
 
         let mut highlight_endpoints = Vec::new();
+        // TODO kb repeat this for all other highlights?
         if let Some(text_highlights) = highlights.text_highlights {
             if !text_highlights.is_empty() {
-                while cursor.start().0 < range.end {
+                self.apply_text_highlights(
-                    let transform_start = self.buffer.anchor_after(
+                    &mut cursor,
-                        self.to_buffer_offset(cmp::max(range.start, cursor.start().0)),
+                    &range,
-                    );
+                    text_highlights,
-                    let transform_end = {
+                    &mut highlight_endpoints,
-                        let overshoot = InlayOffset(range.end.0 - cursor.start().0 .0);
+                );
-                        self.buffer.anchor_before(self.to_buffer_offset(cmp::min(
+                cursor.seek(&range.start, Bias::Right, &());
-                            cursor.end(&()).0,
+            }
-                            cursor.start().0 + overshoot,
+        }
-                        )))
+        if let Some(inlay_highlights) = highlights.inlay_highlights {
-                    };
+            let adjusted_highlights = TreeMap::from_ordered_entries(inlay_highlights.iter().map(
-
+                |(type_id, styled_ranges)| {
-                    for (tag, text_highlights) in text_highlights.iter() {
+                    (
-                        let style = text_highlights.0;
+                        *type_id,
-                        let ranges = &text_highlights.1;
+                        Arc::new((styled_ranges.0, styled_ranges.1.as_slice())),
-
+                    )
-                        let start_ix = match ranges.binary_search_by(|probe| {
+                },
-                            let cmp = probe.end.cmp(&transform_start, &self.buffer);
+            ));
-                            if cmp.is_gt() {
+            if !inlay_highlights.is_empty() {
-                                cmp::Ordering::Greater
+                self.apply_inlay_highlights(
-                            } else {
+                    &mut cursor,
-                                cmp::Ordering::Less
+                    &range,
-                            }
+                    &adjusted_highlights,
-                        }) {
+                    &mut highlight_endpoints,
-                            Ok(i) | Err(i) => i,
+                );
-                        };
+                cursor.seek(&range.start, Bias::Right, &());
-                        for range in &ranges[start_ix..] {
+            }
-                            if range.start.cmp(&transform_end, &self.buffer).is_ge() {
+        }
-                                break;
+        if let Some(inlay_background_highlights) = highlights.inlay_background_highlights.as_ref() {
-                            }
+            if !inlay_background_highlights.is_empty() {
-
+                self.apply_inlay_highlights(
-                            highlight_endpoints.push(HighlightEndpoint {
+                    &mut cursor,
-                                offset: self.to_inlay_offset(range.start.to_offset(&self.buffer)),
+                    &range,
-                                is_start: true,
+                    inlay_background_highlights,
-                                tag: *tag,
+                    &mut highlight_endpoints,
-                                style,
+                );
-                            });
-                            highlight_endpoints.push(HighlightEndpoint {
-                                offset: self.to_inlay_offset(range.end.to_offset(&self.buffer)),
-                                is_start: false,
-                                tag: *tag,
-                                style,
-                            });
-                        }
-                    }
-
-                    cursor.next(&());
-                }
-                highlight_endpoints.sort();
                 cursor.seek(&range.start, Bias::Right, &());
             }
         }
 
+        highlight_endpoints.sort();
         let buffer_range = self.to_buffer_offset(range.start)..self.to_buffer_offset(range.end);
         let buffer_chunks = self.buffer.chunks(buffer_range, language_aware);
 
@@ -1071,6 +1063,137 @@ impl InlaySnapshot {
         }
     }
 
+    fn apply_text_highlights(
+        &self,
+        cursor: &mut Cursor<'_, Transform, (InlayOffset, usize)>,
+        range: &Range<InlayOffset>,
+        text_highlights: &TreeMap<Option<TypeId>, Arc<(HighlightStyle, Vec<Range<Anchor>>)>>,
+        highlight_endpoints: &mut Vec<HighlightEndpoint>,
+    ) {
+        while cursor.start().0 < range.end {
+            let transform_start = self
+                .buffer
+                .anchor_after(self.to_buffer_offset(cmp::max(range.start, cursor.start().0)));
+            let transform_end =
+                {
+                    let overshoot = InlayOffset(range.end.0 - cursor.start().0 .0);
+                    self.buffer.anchor_before(self.to_buffer_offset(cmp::min(
+                        cursor.end(&()).0,
+                        cursor.start().0 + overshoot,
+                    )))
+                };
+
+            for (tag, text_highlights) in text_highlights.iter() {
+                let style = text_highlights.0;
+                let ranges = &text_highlights.1;
+
+                let start_ix = match ranges.binary_search_by(|probe| {
+                    let cmp = probe.end.cmp(&transform_start, &self.buffer);
+                    if cmp.is_gt() {
+                        cmp::Ordering::Greater
+                    } else {
+                        cmp::Ordering::Less
+                    }
+                }) {
+                    Ok(i) | Err(i) => i,
+                };
+                for range in &ranges[start_ix..] {
+                    if range.start.cmp(&transform_end, &self.buffer).is_ge() {
+                        break;
+                    }
+
+                    highlight_endpoints.push(HighlightEndpoint {
+                        offset: self.to_inlay_offset(range.start.to_offset(&self.buffer)),
+                        is_start: true,
+                        tag: *tag,
+                        style,
+                    });
+                    highlight_endpoints.push(HighlightEndpoint {
+                        offset: self.to_inlay_offset(range.end.to_offset(&self.buffer)),
+                        is_start: false,
+                        tag: *tag,
+                        style,
+                    });
+                }
+            }
+
+            cursor.next(&());
+        }
+    }
+
+    fn apply_inlay_highlights(
+        &self,
+        cursor: &mut Cursor<'_, Transform, (InlayOffset, usize)>,
+        range: &Range<InlayOffset>,
+        inlay_highlights: &TreeMap<Option<TypeId>, Arc<(HighlightStyle, &[InlayRange])>>,
+        highlight_endpoints: &mut Vec<HighlightEndpoint>,
+    ) {
+        while cursor.start().0 < range.end {
+            let transform_start = self
+                .buffer
+                .anchor_after(self.to_buffer_offset(cmp::max(range.start, cursor.start().0)));
+            let transform_start = self.to_inlay_offset(transform_start.to_offset(&self.buffer));
+            let transform_end =
+                {
+                    let overshoot = InlayOffset(range.end.0 - cursor.start().0 .0);
+                    self.buffer.anchor_before(self.to_buffer_offset(cmp::min(
+                        cursor.end(&()).0,
+                        cursor.start().0 + overshoot,
+                    )))
+                };
+            let transform_end = self.to_inlay_offset(transform_end.to_offset(&self.buffer));
+
+            // TODO kb add a map
+            let hint_for_id = |id| self.inlays.iter().find(|inlay| inlay.id == id);
+
+            for (tag, inlay_highlights) in inlay_highlights.iter() {
+                let style = inlay_highlights.0;
+                let ranges = inlay_highlights
+                    .1
+                    .iter()
+                    .filter_map(|range| Some((hint_for_id(range.inlay)?, range)))
+                    .map(|(hint, range)| {
+                        let hint_start =
+                            self.to_inlay_offset(hint.position.to_offset(&self.buffer));
+                        let highlight_start = InlayOffset(hint_start.0 + range.highlight_start);
+                        let highlight_end = InlayOffset(hint_start.0 + range.highlight_end);
+                        highlight_start..highlight_end
+                    })
+                    .collect::<Vec<_>>();
+
+                let start_ix = match ranges.binary_search_by(|probe| {
+                    if probe.end > transform_start {
+                        cmp::Ordering::Greater
+                    } else {
+                        cmp::Ordering::Less
+                    }
+                }) {
+                    Ok(i) | Err(i) => i,
+                };
+                for range in &ranges[start_ix..] {
+                    if range.start >= transform_end {
+                        break;
+                    }
+
+                    highlight_endpoints.push(HighlightEndpoint {
+                        offset: range.start,
+                        is_start: true,
+                        tag: *tag,
+                        style,
+                    });
+                    highlight_endpoints.push(HighlightEndpoint {
+                        offset: range.end,
+                        is_start: false,
+                        tag: *tag,
+                        style,
+                    });
+                }
+            }
+
+            cursor.next(&());
+        }
+    }
+
     #[cfg(test)]
     pub fn text(&self) -> String {
         self.chunks(Default::default()..self.len(), false, Highlights::default())
@@ -1495,7 +1618,12 @@ mod tests {
 
         // The inlays can be manually removed.
         let (inlay_snapshot, _) = inlay_map.splice(
-            inlay_map.inlays.iter().map(|inlay| inlay.id).collect(),
+            inlay_map
+                .snapshot
+                .inlays
+                .iter()
+                .map(|inlay| inlay.id)
+                .collect(),
             Vec::new(),
         );
         assert_eq!(inlay_snapshot.text(), "abxJKLyDzefghi");
@@ -1587,6 +1715,7 @@ mod tests {
             log::info!("inlay text: {:?}", inlay_snapshot.text());
 
             let inlays = inlay_map
+                .snapshot
                 .inlays
                 .iter()
                 .filter(|inlay| inlay.position.is_valid(&buffer_snapshot))

crates/editor/src/editor.rs

@@ -99,6 +99,7 @@ use std::{
     time::{Duration, Instant},
 };
 pub use sum_tree::Bias;
+use sum_tree::TreeMap;
 use text::Rope;
 use theme::{DiagnosticStyle, Theme, ThemeSettings};
 use util::{post_inc, RangeExt, ResultExt, TryFutureExt};
@@ -581,7 +582,7 @@ pub struct Editor {
     placeholder_text: Option<Arc<str>>,
     highlighted_rows: Option<Range<u32>>,
     background_highlights: BTreeMap<TypeId, BackgroundHighlight>,
-    inlay_background_highlights: BTreeMap<TypeId, InlayBackgroundHighlight>,
+    inlay_background_highlights: TreeMap<Option<TypeId>, InlayBackgroundHighlight>,
     nav_history: Option<ItemNavHistory>,
     context_menu: Option<ContextMenu>,
     mouse_context_menu: ViewHandle<context_menu::ContextMenu>,
@@ -7823,7 +7824,7 @@ impl Editor {
         cx: &mut ViewContext<Self>,
     ) {
         self.inlay_background_highlights
-            .insert(TypeId::of::<T>(), (color_fetcher, ranges));
+            .insert(Some(TypeId::of::<T>()), (color_fetcher, ranges));
         cx.notify();
     }
 
@@ -7832,7 +7833,9 @@ impl Editor {
         cx: &mut ViewContext<Self>,
     ) -> Option<BackgroundHighlight> {
         let text_highlights = self.background_highlights.remove(&TypeId::of::<T>());
-        let inlay_highlights = self.inlay_background_highlights.remove(&TypeId::of::<T>());
+        let inlay_highlights = self
+            .inlay_background_highlights
+            .remove(&Some(TypeId::of::<T>()));
         if text_highlights.is_some() || inlay_highlights.is_some() {
             cx.notify();
         }

crates/editor/src/element.rs

@@ -54,6 +54,7 @@ use std::{
     ops::Range,
     sync::Arc,
 };
+use sum_tree::TreeMap;
 use text::Point;
 use workspace::item::Item;
 
@@ -1592,11 +1593,28 @@ impl EditorElement {
                 .collect()
         } else {
             let style = &self.style;
+            let theme = theme::current(cx);
+            let inlay_background_highlights =
+                TreeMap::from_ordered_entries(editor.inlay_background_highlights.iter().map(
+                    |(type_id, (color_fetcher, ranges))| {
+                        let color = Some(color_fetcher(&theme));
+                        (
+                            *type_id,
+                            Arc::new((
+                                HighlightStyle {
+                                    color,
+                                    ..HighlightStyle::default()
+                                },
+                                ranges.as_slice(),
+                            )),
+                        )
+                    },
+                ));
             let chunks = snapshot
                 .chunks(
                     rows.clone(),
                     true,
-                    Some(&editor.inlay_background_highlights),
+                    Some(inlay_background_highlights),
                     Some(style.theme.hint),
                     Some(style.theme.suggestion),
                 )