1use crate::{Supermaven, SupermavenCompletionStateId};
2use anyhow::Result;
3use edit_prediction_types::{EditPrediction, EditPredictionDelegate};
4use futures::StreamExt as _;
5use gpui::{App, Context, Entity, EntityId, Task};
6use language::{Anchor, Buffer, BufferSnapshot};
7use std::{
8 ops::{AddAssign, Range},
9 path::Path,
10 sync::Arc,
11 time::Duration,
12};
13use text::{ToOffset, ToPoint};
14use unicode_segmentation::UnicodeSegmentation;
15
16pub const DEBOUNCE_TIMEOUT: Duration = Duration::from_millis(75);
17
18pub struct SupermavenEditPredictionDelegate {
19 supermaven: Entity<Supermaven>,
20 buffer_id: Option<EntityId>,
21 completion_id: Option<SupermavenCompletionStateId>,
22 completion_text: Option<String>,
23 file_extension: Option<String>,
24 pending_refresh: Option<Task<Result<()>>>,
25 completion_position: Option<language::Anchor>,
26}
27
28impl SupermavenEditPredictionDelegate {
29 pub fn new(supermaven: Entity<Supermaven>) -> Self {
30 Self {
31 supermaven,
32 buffer_id: None,
33 completion_id: None,
34 completion_text: None,
35 file_extension: None,
36 pending_refresh: None,
37 completion_position: None,
38 }
39 }
40}
41
42// Computes the edit prediction from the difference between the completion text.
43// This is defined by greedily matching the buffer text against the completion text.
44// Inlays are inserted for parts of the completion text that are not present in the buffer text.
45// For example, given the completion text "axbyc" and the buffer text "xy", the rendered output in the editor would be "[a]x[b]y[c]".
46// The parts in brackets are the inlays.
47fn completion_from_diff(
48 snapshot: BufferSnapshot,
49 completion_text: &str,
50 position: Anchor,
51 delete_range: Range<Anchor>,
52) -> EditPrediction {
53 let buffer_text = snapshot.text_for_range(delete_range).collect::<String>();
54
55 let mut edits: Vec<(Range<language::Anchor>, Arc<str>)> = Vec::new();
56
57 let completion_graphemes: Vec<&str> = completion_text.graphemes(true).collect();
58 let buffer_graphemes: Vec<&str> = buffer_text.graphemes(true).collect();
59
60 let mut offset = position.to_offset(&snapshot);
61
62 let mut i = 0;
63 let mut j = 0;
64 while i < completion_graphemes.len() && j < buffer_graphemes.len() {
65 // find the next instance of the buffer text in the completion text.
66 let k = completion_graphemes[i..]
67 .iter()
68 .position(|c| *c == buffer_graphemes[j]);
69 match k {
70 Some(k) => {
71 if k != 0 {
72 let offset = snapshot.anchor_after(offset);
73 // the range from the current position to item is an inlay.
74 let edit = (
75 offset..offset,
76 completion_graphemes[i..i + k].join("").into(),
77 );
78 edits.push(edit);
79 }
80 i += k + 1;
81 j += 1;
82 offset.add_assign(buffer_graphemes[j - 1].len());
83 }
84 None => {
85 // there are no more matching completions, so drop the remaining
86 // completion text as an inlay.
87 break;
88 }
89 }
90 }
91
92 if j == buffer_graphemes.len() && i < completion_graphemes.len() {
93 let offset = snapshot.anchor_after(offset);
94 // there is leftover completion text, so drop it as an inlay.
95 let edit_range = offset..offset;
96 let edit_text = completion_graphemes[i..].join("");
97 edits.push((edit_range, edit_text.into()));
98 }
99
100 EditPrediction::Local {
101 id: None,
102 edits,
103 edit_preview: None,
104 }
105}
106
107impl EditPredictionDelegate for SupermavenEditPredictionDelegate {
108 fn name() -> &'static str {
109 "supermaven"
110 }
111
112 fn display_name() -> &'static str {
113 "Supermaven"
114 }
115
116 fn show_predictions_in_menu() -> bool {
117 true
118 }
119
120 fn show_tab_accept_marker() -> bool {
121 true
122 }
123
124 fn supports_jump_to_edit() -> bool {
125 false
126 }
127
128 fn is_enabled(&self, _buffer: &Entity<Buffer>, _cursor_position: Anchor, cx: &App) -> bool {
129 self.supermaven.read(cx).is_enabled()
130 }
131
132 fn is_refreshing(&self, _cx: &App) -> bool {
133 self.pending_refresh.is_some() && self.completion_id.is_none()
134 }
135
136 fn refresh(
137 &mut self,
138 buffer_handle: Entity<Buffer>,
139 cursor_position: Anchor,
140 debounce: bool,
141 cx: &mut Context<Self>,
142 ) {
143 // Only make new completion requests when debounce is true (i.e., when text is typed)
144 // When debounce is false (i.e., cursor movement), we should not make new requests
145 if !debounce {
146 return;
147 }
148
149 reset_completion_cache(self, cx);
150
151 let Some(mut completion) = self.supermaven.update(cx, |supermaven, cx| {
152 supermaven.complete(&buffer_handle, cursor_position, cx)
153 }) else {
154 return;
155 };
156
157 self.pending_refresh = Some(cx.spawn(async move |this, cx| {
158 if debounce {
159 cx.background_executor().timer(DEBOUNCE_TIMEOUT).await;
160 }
161
162 while let Some(()) = completion.updates.next().await {
163 this.update(cx, |this, cx| {
164 // Get the completion text and cache it
165 if let Some(text) =
166 this.supermaven
167 .read(cx)
168 .completion(&buffer_handle, cursor_position, cx)
169 {
170 this.completion_text = Some(text.to_string());
171
172 this.completion_position = Some(cursor_position);
173 }
174
175 this.completion_id = Some(completion.id);
176 this.buffer_id = Some(buffer_handle.entity_id());
177 this.file_extension = buffer_handle.read(cx).file().and_then(|file| {
178 Some(
179 Path::new(file.file_name(cx))
180 .extension()?
181 .to_str()?
182 .to_string(),
183 )
184 });
185 cx.notify();
186 })?;
187 }
188 Ok(())
189 }));
190 }
191
192 fn accept(&mut self, _cx: &mut Context<Self>) {
193 reset_completion_cache(self, _cx);
194 }
195
196 fn discard(&mut self, _cx: &mut Context<Self>) {
197 reset_completion_cache(self, _cx);
198 }
199
200 fn suggest(
201 &mut self,
202 buffer: &Entity<Buffer>,
203 cursor_position: Anchor,
204 cx: &mut Context<Self>,
205 ) -> Option<EditPrediction> {
206 if self.buffer_id != Some(buffer.entity_id()) {
207 return None;
208 }
209
210 if self.completion_id.is_none() {
211 return None;
212 }
213
214 let completion_text = if let Some(cached_text) = &self.completion_text {
215 cached_text.as_str()
216 } else {
217 let text = self
218 .supermaven
219 .read(cx)
220 .completion(buffer, cursor_position, cx)?;
221 self.completion_text = Some(text.to_string());
222 text
223 };
224
225 // Check if the cursor is still at the same position as the completion request
226 // If we don't have a completion position stored, don't show the completion
227 if let Some(completion_position) = self.completion_position {
228 if cursor_position != completion_position {
229 return None;
230 }
231 } else {
232 return None;
233 }
234
235 let completion_text = trim_to_end_of_line_unless_leading_newline(completion_text);
236
237 let completion_text = completion_text.trim_end();
238
239 if !completion_text.trim().is_empty() {
240 let snapshot = buffer.read(cx).snapshot();
241
242 // Calculate the range from cursor to end of line correctly
243 let cursor_point = cursor_position.to_point(&snapshot);
244 let end_of_line = snapshot.anchor_after(language::Point::new(
245 cursor_point.row,
246 snapshot.line_len(cursor_point.row),
247 ));
248 let delete_range = cursor_position..end_of_line;
249
250 Some(completion_from_diff(
251 snapshot,
252 completion_text,
253 cursor_position,
254 delete_range,
255 ))
256 } else {
257 None
258 }
259 }
260}
261
262fn reset_completion_cache(
263 provider: &mut SupermavenEditPredictionDelegate,
264 _cx: &mut Context<SupermavenEditPredictionDelegate>,
265) {
266 provider.pending_refresh = None;
267 provider.completion_id = None;
268 provider.completion_text = None;
269 provider.completion_position = None;
270 provider.buffer_id = None;
271}
272
273fn trim_to_end_of_line_unless_leading_newline(text: &str) -> &str {
274 if has_leading_newline(text) {
275 text
276 } else if let Some(i) = text.find('\n') {
277 &text[..i]
278 } else {
279 text
280 }
281}
282
283fn has_leading_newline(text: &str) -> bool {
284 for c in text.chars() {
285 if c == '\n' {
286 return true;
287 }
288 if !c.is_whitespace() {
289 return false;
290 }
291 }
292 false
293}