Switch over to inlay map for Copilot suggestions

Co-Authored-By: Antonio Scandurra <antonio@zed.dev>

crates/editor/src/display_map.rs

@@ -6,12 +6,12 @@ mod tab_map;
 mod wrap_map;
 
 use crate::{
-    display_map::inlay_map::InlayProperties, inlay_cache::InlayId, Anchor, AnchorRangeExt,
-    MultiBuffer, MultiBufferSnapshot, ToOffset, ToPoint,
+    inlay_cache::{Inlay, InlayId, InlayProperties},
+    Anchor, AnchorRangeExt, MultiBuffer, MultiBufferSnapshot, ToOffset, ToPoint,
 };
 pub use block_map::{BlockMap, BlockPoint};
 use collections::{HashMap, HashSet};
-use fold_map::{FoldMap, FoldOffset};
+use fold_map::FoldMap;
 use gpui::{
     color::Color,
     fonts::{FontId, HighlightStyle},
@@ -26,6 +26,7 @@ pub use suggestion_map::Suggestion;
 use suggestion_map::SuggestionMap;
 use sum_tree::{Bias, TreeMap};
 use tab_map::TabMap;
+use text::Rope;
 use wrap_map::WrapMap;
 
 pub use block_map::{
@@ -244,33 +245,6 @@ impl DisplayMap {
         self.text_highlights.remove(&Some(type_id))
     }
 
-    pub fn has_suggestion(&self) -> bool {
-        self.suggestion_map.has_suggestion()
-    }
-
-    pub fn replace_suggestion<T>(
-        &mut self,
-        new_suggestion: Option<Suggestion<T>>,
-        cx: &mut ModelContext<Self>,
-    ) -> Option<Suggestion<FoldOffset>>
-    where
-        T: ToPoint,
-    {
-        let snapshot = self.buffer.read(cx).snapshot(cx);
-        let edits = self.buffer_subscription.consume().into_inner();
-        let tab_size = Self::tab_size(&self.buffer, cx);
-        let (snapshot, edits) = self.fold_map.read(snapshot, edits);
-        let (snapshot, edits, old_suggestion) =
-            self.suggestion_map.replace(new_suggestion, snapshot, edits);
-        let (snapshot, edits) = self.inlay_map.sync(snapshot, edits);
-        let (snapshot, edits) = self.tab_map.sync(snapshot, edits, tab_size);
-        let (snapshot, edits) = self
-            .wrap_map
-            .update(cx, |map, cx| map.sync(snapshot, edits, cx));
-        self.block_map.read(snapshot, edits);
-        old_suggestion
-    }
-
     pub fn set_font(&self, font_id: FontId, font_size: f32, cx: &mut ModelContext<Self>) -> bool {
         self.wrap_map
             .update(cx, |map, cx| map.set_font(font_id, font_size, cx))
@@ -285,10 +259,10 @@ impl DisplayMap {
             .update(cx, |map, cx| map.set_wrap_width(width, cx))
     }
 
-    pub fn splice_inlays(
+    pub fn splice_inlays<T: Into<Rope>>(
         &mut self,
         to_remove: Vec<InlayId>,
-        to_insert: Vec<(InlayId, Anchor, project::InlayHint)>,
+        to_insert: Vec<(InlayId, InlayProperties<T>)>,
         cx: &mut ModelContext<Self>,
     ) {
         let buffer_snapshot = self.buffer.read(cx).snapshot(cx);
@@ -303,28 +277,7 @@ impl DisplayMap {
             .update(cx, |map, cx| map.sync(snapshot, edits, cx));
         self.block_map.read(snapshot, edits);
 
-        let new_inlays: Vec<(InlayId, InlayProperties<String>)> = to_insert
-            .into_iter()
-            .map(|(inlay_id, hint_anchor, hint)| {
-                let mut text = hint.text();
-                // TODO kb styling instead?
-                if hint.padding_right {
-                    text.push(' ');
-                }
-                if hint.padding_left {
-                    text.insert(0, ' ');
-                }
-
-                (
-                    inlay_id,
-                    InlayProperties {
-                        position: hint_anchor.bias_left(&buffer_snapshot),
-                        text,
-                    },
-                )
-            })
-            .collect();
-        let (snapshot, edits) = self.inlay_map.splice(to_remove, new_inlays);
+        let (snapshot, edits) = self.inlay_map.splice(to_remove, to_insert);
         let (snapshot, edits) = self.tab_map.sync(snapshot, edits, tab_size);
         let (snapshot, edits) = self
             .wrap_map

crates/editor/src/display_map/inlay_map.rs

@@ -5,7 +5,10 @@ use super::{
     },
     TextHighlights,
 };
-use crate::{inlay_cache::InlayId, Anchor, MultiBufferSnapshot, ToPoint};
+use crate::{
+    inlay_cache::{Inlay, InlayId, InlayProperties},
+    MultiBufferSnapshot, ToPoint,
+};
 use collections::{BTreeSet, HashMap};
 use gpui::fonts::HighlightStyle;
 use language::{Chunk, Edit, Point, Rope, TextSummary};
@@ -140,19 +143,6 @@ pub struct InlayChunks<'a> {
     highlight_style: Option<HighlightStyle>,
 }
 
-#[derive(Debug, Clone)]
-pub struct Inlay {
-    pub(super) id: InlayId,
-    pub(super) position: Anchor,
-    pub(super) text: Rope,
-}
-
-#[derive(Debug, Clone)]
-pub struct InlayProperties<T> {
-    pub position: Anchor,
-    pub text: T,
-}
-
 impl<'a> Iterator for InlayChunks<'a> {
     type Item = Chunk<'a>;
 
@@ -431,6 +421,21 @@ impl InlayMap {
         snapshot.version += 1;
 
         let mut edits = BTreeSet::new();
+
+        self.inlays.retain(|inlay| !to_remove.contains(&inlay.id));
+        for inlay_id in to_remove {
+            if let Some(inlay) = self.inlays_by_id.remove(&inlay_id) {
+                let buffer_point = inlay.position.to_point(snapshot.buffer_snapshot());
+                let fold_point = snapshot
+                    .suggestion_snapshot
+                    .fold_snapshot
+                    .to_fold_point(buffer_point, Bias::Left);
+                let suggestion_point = snapshot.suggestion_snapshot.to_suggestion_point(fold_point);
+                let suggestion_offset = snapshot.suggestion_snapshot.to_offset(suggestion_point);
+                edits.insert(suggestion_offset);
+            }
+        }
+
         for (id, properties) in to_insert {
             let inlay = Inlay {
                 id,
@@ -458,20 +463,6 @@ impl InlayMap {
             edits.insert(suggestion_offset);
         }
 
-        self.inlays.retain(|inlay| !to_remove.contains(&inlay.id));
-        for inlay_id in to_remove {
-            if let Some(inlay) = self.inlays_by_id.remove(&inlay_id) {
-                let buffer_point = inlay.position.to_point(snapshot.buffer_snapshot());
-                let fold_point = snapshot
-                    .suggestion_snapshot
-                    .fold_snapshot
-                    .to_fold_point(buffer_point, Bias::Left);
-                let suggestion_point = snapshot.suggestion_snapshot.to_suggestion_point(fold_point);
-                let suggestion_offset = snapshot.suggestion_snapshot.to_offset(suggestion_point);
-                edits.insert(suggestion_offset);
-            }
-        }
-
         let suggestion_snapshot = snapshot.suggestion_snapshot.clone();
         let suggestion_edits = edits
             .into_iter()

crates/editor/src/editor.rs

@@ -54,7 +54,9 @@ use gpui::{
 };
 use highlight_matching_bracket::refresh_matching_bracket_highlights;
 use hover_popover::{hide_hover, HoverState};
-use inlay_cache::{InlayCache, InlayRefreshReason, InlaySplice, QueryInlaysRange};
+use inlay_cache::{
+    Inlay, InlayCache, InlayId, InlayProperties, InlayRefreshReason, InlaySplice, QueryInlaysRange,
+};
 pub use items::MAX_TAB_TITLE_LEN;
 use itertools::Itertools;
 pub use language::{char_kind, CharKind};
@@ -95,6 +97,7 @@ use std::{
     time::{Duration, Instant},
 };
 pub use sum_tree::Bias;
+use text::Rope;
 use theme::{DiagnosticStyle, Theme, ThemeSettings};
 use util::{post_inc, RangeExt, ResultExt, TryFutureExt};
 use workspace::{ItemNavHistory, ViewId, Workspace};
@@ -1061,6 +1064,7 @@ pub struct CopilotState {
     cycled: bool,
     completions: Vec<copilot::Completion>,
     active_completion_index: usize,
+    suggestion: Option<Inlay>,
 }
 
 impl Default for CopilotState {
@@ -1072,6 +1076,7 @@ impl Default for CopilotState {
             completions: Default::default(),
             active_completion_index: 0,
             cycled: false,
+            suggestion: None,
         }
     }
 }
@@ -2597,9 +2602,7 @@ impl Editor {
 
         if !settings::get::<EditorSettings>(cx).inlay_hints.enabled {
             let to_remove = self.inlay_cache.clear();
-            self.display_map.update(cx, |display_map, cx| {
-                display_map.splice_inlays(to_remove, Vec::new(), cx);
-            });
+            self.splice_inlay_hints(to_remove, Vec::new(), cx);
             return;
         }
 
@@ -2609,9 +2612,7 @@ impl Editor {
                     to_remove,
                     to_insert,
                 } = self.inlay_cache.apply_settings(new_settings);
-                self.display_map.update(cx, |display_map, cx| {
-                    display_map.splice_inlays(to_remove, to_insert, cx);
-                });
+                self.splice_inlay_hints(to_remove, to_insert, cx);
             }
             InlayRefreshReason::Regular => {
                 let buffer_handle = self.buffer().clone();
@@ -2652,9 +2653,7 @@ impl Editor {
                         .context("inlay cache hint fetch")?;
 
                     editor.update(&mut cx, |editor, cx| {
-                        editor.display_map.update(cx, |display_map, cx| {
-                            display_map.splice_inlays(to_remove, to_insert, cx);
-                        });
+                        editor.splice_inlay_hints(to_remove, to_insert, cx)
                     })
                 })
                 .detach_and_log_err(cx);
@@ -2662,6 +2661,40 @@ impl Editor {
         }
     }
 
+    fn splice_inlay_hints(
+        &self,
+        to_remove: Vec<InlayId>,
+        to_insert: Vec<(InlayId, Anchor, project::InlayHint)>,
+        cx: &mut ViewContext<Self>,
+    ) {
+        let buffer = self.buffer.read(cx).read(cx);
+        let new_inlays: Vec<(InlayId, InlayProperties<String>)> = to_insert
+            .into_iter()
+            .map(|(inlay_id, hint_anchor, hint)| {
+                let mut text = hint.text();
+                // TODO kb styling instead?
+                if hint.padding_right {
+                    text.push(' ');
+                }
+                if hint.padding_left {
+                    text.insert(0, ' ');
+                }
+
+                (
+                    inlay_id,
+                    InlayProperties {
+                        position: hint_anchor.bias_left(&buffer),
+                        text,
+                    },
+                )
+            })
+            .collect();
+        drop(buffer);
+        self.display_map.update(cx, |display_map, cx| {
+            display_map.splice_inlays(to_remove, new_inlays, cx);
+        });
+    }
+
     fn trigger_on_type_formatting(
         &self,
         input: String,
@@ -3312,10 +3345,7 @@ impl Editor {
     }
 
     fn accept_copilot_suggestion(&mut self, cx: &mut ViewContext<Self>) -> bool {
-        if let Some(suggestion) = self
-            .display_map
-            .update(cx, |map, cx| map.replace_suggestion::<usize>(None, cx))
-        {
+        if let Some(suggestion) = self.take_active_copilot_suggestion(cx) {
             if let Some((copilot, completion)) =
                 Copilot::global(cx).zip(self.copilot_state.active_completion())
             {
@@ -3334,7 +3364,7 @@ impl Editor {
     }
 
     fn discard_copilot_suggestion(&mut self, cx: &mut ViewContext<Self>) -> bool {
-        if self.has_active_copilot_suggestion(cx) {
+        if let Some(suggestion) = self.take_active_copilot_suggestion(cx) {
             if let Some(copilot) = Copilot::global(cx) {
                 copilot
                     .update(cx, |copilot, cx| {
@@ -3345,8 +3375,9 @@ impl Editor {
                 self.report_copilot_event(None, false, cx)
             }
 
-            self.display_map
-                .update(cx, |map, cx| map.replace_suggestion::<usize>(None, cx));
+            self.display_map.update(cx, |map, cx| {
+                map.splice_inlays::<&str>(vec![suggestion.id], Vec::new(), cx)
+            });
             cx.notify();
             true
         } else {
@@ -3367,7 +3398,26 @@ impl Editor {
     }
 
     fn has_active_copilot_suggestion(&self, cx: &AppContext) -> bool {
-        self.display_map.read(cx).has_suggestion()
+        if let Some(suggestion) = self.copilot_state.suggestion.as_ref() {
+            let buffer = self.buffer.read(cx).read(cx);
+            suggestion.position.is_valid(&buffer)
+        } else {
+            false
+        }
+    }
+
+    fn take_active_copilot_suggestion(&mut self, cx: &mut ViewContext<Self>) -> Option<Inlay> {
+        let suggestion = self.copilot_state.suggestion.take()?;
+        self.display_map.update(cx, |map, cx| {
+            map.splice_inlays::<&str>(vec![suggestion.id], Default::default(), cx);
+        });
+        let buffer = self.buffer.read(cx).read(cx);
+
+        if suggestion.position.is_valid(&buffer) {
+            Some(suggestion)
+        } else {
+            None
+        }
     }
 
     fn update_visible_copilot_suggestion(&mut self, cx: &mut ViewContext<Self>) {
@@ -3384,14 +3434,28 @@ impl Editor {
             .copilot_state
             .text_for_active_completion(cursor, &snapshot)
         {
+            let text = Rope::from(text);
+            let mut to_remove = Vec::new();
+            if let Some(suggestion) = self.copilot_state.suggestion.take() {
+                to_remove.push(suggestion.id);
+            }
+
+            let to_insert = vec![(
+                // TODO kb check how can I get the unique id for the suggestion
+                // Move the generation of the id inside the map
+                InlayId(usize::MAX),
+                InlayProperties {
+                    position: cursor,
+                    text: text.clone(),
+                },
+            )];
             self.display_map.update(cx, move |map, cx| {
-                map.replace_suggestion(
-                    Some(Suggestion {
-                        position: cursor,
-                        text: text.trim_end().into(),
-                    }),
-                    cx,
-                )
+                map.splice_inlays(to_remove, to_insert, cx)
+            });
+            self.copilot_state.suggestion = Some(Inlay {
+                id: InlayId(usize::MAX),
+                position: cursor,
+                text,
             });
             cx.notify();
         } else {

crates/editor/src/inlay_cache.rs

@@ -14,6 +14,19 @@ use util::post_inc;
 
 use collections::{hash_map, BTreeMap, HashMap, HashSet};
 
+#[derive(Debug, Clone)]
+pub struct Inlay {
+    pub id: InlayId,
+    pub position: Anchor,
+    pub text: text::Rope,
+}
+
+#[derive(Debug, Clone)]
+pub struct InlayProperties<T> {
+    pub position: Anchor,
+    pub text: T,
+}
+
 #[derive(Debug, Copy, Clone)]
 pub enum InlayRefreshReason {
     Settings(editor_settings::InlayHints),