Restructure inlay highlights data for proper access

crates/editor/src/display_map.rs

@@ -9,7 +9,7 @@ use crate::{
     MultiBufferSnapshot, ToOffset, ToPoint,
 };
 pub use block_map::{BlockMap, BlockPoint};
-use collections::{HashMap, HashSet};
+use collections::{BTreeMap, HashMap, HashSet};
 use fold_map::FoldMap;
 use gpui::{
     color::Color,
@@ -44,7 +44,7 @@ pub trait ToDisplayPoint {
 }
 
 type TextHighlights = TreeMap<Option<TypeId>, Arc<(HighlightStyle, Vec<Range<Anchor>>)>>;
-type InlayHighlights = TreeMap<Option<TypeId>, Arc<(HighlightStyle, Vec<InlayHighlight>)>>;
+type InlayHighlights = BTreeMap<TypeId, HashMap<InlayId, (HighlightStyle, InlayHighlight)>>;
 
 pub struct DisplayMap {
     buffer: ModelHandle<MultiBuffer>,
@@ -226,11 +226,15 @@ impl DisplayMap {
     pub fn highlight_inlays(
         &mut self,
         type_id: TypeId,
-        ranges: Vec<InlayHighlight>,
+        highlights: Vec<InlayHighlight>,
         style: HighlightStyle,
     ) {
+        for highlight in highlights {
             self.inlay_highlights
-            .insert(Some(type_id), Arc::new((style, ranges)));
+                .entry(type_id)
+                .or_default()
+                .insert(highlight.inlay, (style, highlight));
+        }
     }
 
     pub fn text_highlights(&self, type_id: TypeId) -> Option<(HighlightStyle, &[Range<Anchor>])> {
@@ -239,7 +243,7 @@ impl DisplayMap {
     }
     pub fn clear_highlights(&mut self, type_id: TypeId) -> bool {
         let mut cleared = self.text_highlights.remove(&Some(type_id)).is_some();
-        cleared |= self.inlay_highlights.remove(&Some(type_id)).is_none();
+        cleared |= self.inlay_highlights.remove(&type_id).is_none();
         cleared
     }
 
@@ -756,11 +760,11 @@ impl DisplaySnapshot {
     }
 
     #[cfg(any(test, feature = "test-support"))]
-    pub fn inlay_highlight_ranges<Tag: ?Sized + 'static>(
+    pub fn inlay_highlights<Tag: ?Sized + 'static>(
         &self,
-    ) -> Option<Arc<(HighlightStyle, Vec<InlayHighlight>)>> {
+    ) -> Option<&HashMap<InlayId, (HighlightStyle, InlayHighlight)>> {
         let type_id = TypeId::of::<Tag>();
-        self.inlay_highlights.get(&Some(type_id)).cloned()
+        self.inlay_highlights.get(&type_id)
     }
 }
 

crates/editor/src/display_map/inlay_map.rs

@@ -1,5 +1,4 @@
 use crate::{
-    link_go_to_definition::InlayHighlight,
     multi_buffer::{MultiBufferChunks, MultiBufferRows},
     Anchor, InlayId, MultiBufferSnapshot, ToOffset,
 };
@@ -295,16 +294,13 @@ impl<'a> Iterator for InlayChunks<'a> {
                 prefix
             }
             Transform::Inlay(inlay) => {
-                let mut inlay_highlight_style_and_range = None;
+                let mut inlay_style_and_highlight = None;
+                // type InlayHighlights = BTreeMap<TypeId, HashMap<InlayId, (HighlightStyle, InlayHighlight)>>;
                 if let Some(inlay_highlights) = self.highlights.inlay_highlights {
-                    for (_, style_and_ranges) in inlay_highlights.iter() {
-                        let highlight_style: &HighlightStyle = &style_and_ranges.0;
-                        let inlay_ranges: &Vec<InlayHighlight> = &style_and_ranges.1;
-                        // TODO kb: turn into a map.
-                        if let Some(range) =
-                            inlay_ranges.iter().find(|range| range.inlay == inlay.id)
-                        {
-                            inlay_highlight_style_and_range = Some((highlight_style, range));
+                    for (_, inlay_id_to_data) in inlay_highlights.iter() {
+                        let style_and_highlight = inlay_id_to_data.get(&inlay.id);
+                        if style_and_highlight.is_some() {
+                            inlay_style_and_highlight = style_and_highlight;
                             break;
                         }
                     }
@@ -316,7 +312,7 @@ impl<'a> Iterator for InlayChunks<'a> {
                 };
                 let next_inlay_highlight_endpoint;
                 let offset_in_inlay = self.output_offset - self.transforms.start().0;
-                if let Some((style, highlight)) = inlay_highlight_style_and_range {
+                if let Some((style, highlight)) = inlay_style_and_highlight {
                     let range = &highlight.range;
                     if offset_in_inlay.0 < range.start {
                         next_inlay_highlight_endpoint = range.start - offset_in_inlay.0;
@@ -1176,6 +1172,7 @@ mod tests {
     use super::*;
     use crate::{
         display_map::{InlayHighlights, TextHighlights},
+        link_go_to_definition::InlayHighlight,
         InlayId, MultiBuffer,
     };
     use gpui::AppContext;
@@ -1706,7 +1703,7 @@ mod tests {
                 while inlay_indices.len() < highlight_count.min(inlays.len()) {
                     inlay_indices.insert(rng.gen_range(0..inlays.len()));
                 }
-                let highlight_ranges = inlay_indices
+                let new_highlights = inlay_indices
                     .into_iter()
                     .filter_map(|i| {
                         let (_, inlay) = &inlays[i];
@@ -1736,13 +1733,10 @@ mod tests {
                             }
                         }
                     })
+                    .map(|highlight| (highlight.inlay, (HighlightStyle::default(), highlight)))
                     .collect();
-
-                log::info!("highlighting inlay ranges {highlight_ranges:?}");
-                inlay_highlights.insert(
-                    Some(TypeId::of::<()>()),
-                    Arc::new((HighlightStyle::default(), highlight_ranges)),
-                );
+                log::info!("highlighting inlay ranges {new_highlights:?}");
+                inlay_highlights.insert(TypeId::of::<()>(), new_highlights);
             };
 
             for _ in 0..5 {

crates/editor/src/editor.rs

@@ -8020,12 +8020,12 @@ impl Editor {
 
     pub fn highlight_inlays<T: 'static>(
         &mut self,
-        ranges: Vec<InlayHighlight>,
+        highlights: Vec<InlayHighlight>,
         style: HighlightStyle,
         cx: &mut ViewContext<Self>,
     ) {
         self.display_map.update(cx, |map, _| {
-            map.highlight_inlays(TypeId::of::<T>(), ranges, style)
+            map.highlight_inlays(TypeId::of::<T>(), highlights, style)
         });
         cx.notify();
     }

crates/editor/src/link_go_to_definition.rs

@@ -1194,18 +1194,19 @@ mod tests {
         cx.foreground().run_until_parked();
         cx.update_editor(|editor, cx| {
             let snapshot = editor.snapshot(cx);
-            let actual_ranges = snapshot
-                .inlay_highlight_ranges::<LinkGoToDefinitionState>()
-                .map(|ranges| ranges.as_ref().clone().1)
-                .unwrap_or_default();
+            let actual_highlights = snapshot
+                .inlay_highlights::<LinkGoToDefinitionState>()
+                .into_iter()
+                .flat_map(|highlights| highlights.values().map(|(_, highlight)| highlight))
+                .collect::<Vec<_>>();
 
             let buffer_snapshot = editor.buffer().update(cx, |buffer, cx| buffer.snapshot(cx));
-            let expected_ranges = vec![InlayHighlight {
+            let expected_highlight = InlayHighlight {
                 inlay: InlayId::Hint(0),
                 inlay_position: buffer_snapshot.anchor_at(inlay_range.start, Bias::Right),
                 range: 0..hint_label.len(),
-            }];
-            assert_set_eq!(actual_ranges, expected_ranges);
+            };
+            assert_set_eq!(actual_highlights, vec![&expected_highlight]);
         });
 
         // Unpress cmd causes highlight to go away