supermaven_completion_provider.rs

  1use crate::{Supermaven, SupermavenCompletionStateId};
  2use anyhow::Result;
  3use edit_prediction::{Direction, EditPrediction, EditPredictionProvider};
  4use futures::StreamExt as _;
  5use gpui::{App, Context, Entity, EntityId, Task};
  6use language::{Anchor, Buffer, BufferSnapshot};
  7use project::Project;
  8use std::{
  9    ops::{AddAssign, Range},
 10    path::Path,
 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 SupermavenCompletionProvider {
 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 SupermavenCompletionProvider {
 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>, String)> = 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 = (offset..offset, completion_graphemes[i..i + k].join(""));
 75                    edits.push(edit);
 76                }
 77                i += k + 1;
 78                j += 1;
 79                offset.add_assign(buffer_graphemes[j - 1].len());
 80            }
 81            None => {
 82                // there are no more matching completions, so drop the remaining
 83                // completion text as an inlay.
 84                break;
 85            }
 86        }
 87    }
 88
 89    if j == buffer_graphemes.len() && i < completion_graphemes.len() {
 90        let offset = snapshot.anchor_after(offset);
 91        // there is leftover completion text, so drop it as an inlay.
 92        let edit_range = offset..offset;
 93        let edit_text = completion_graphemes[i..].join("");
 94        edits.push((edit_range, edit_text));
 95    }
 96
 97    EditPrediction {
 98        id: None,
 99        edits,
100        edit_preview: None,
101    }
102}
103
104impl EditPredictionProvider for SupermavenCompletionProvider {
105    fn name() -> &'static str {
106        "supermaven"
107    }
108
109    fn display_name() -> &'static str {
110        "Supermaven"
111    }
112
113    fn show_completions_in_menu() -> bool {
114        true
115    }
116
117    fn show_tab_accept_marker() -> bool {
118        true
119    }
120
121    fn supports_jump_to_edit() -> bool {
122        false
123    }
124
125    fn is_enabled(&self, _buffer: &Entity<Buffer>, _cursor_position: Anchor, cx: &App) -> bool {
126        self.supermaven.read(cx).is_enabled()
127    }
128
129    fn is_refreshing(&self) -> bool {
130        self.pending_refresh.is_some() && self.completion_id.is_none()
131    }
132
133    fn refresh(
134        &mut self,
135        _project: Option<Entity<Project>>,
136        buffer_handle: Entity<Buffer>,
137        cursor_position: Anchor,
138        debounce: bool,
139        cx: &mut Context<Self>,
140    ) {
141        // Only make new completion requests when debounce is true (i.e., when text is typed)
142        // When debounce is false (i.e., cursor movement), we should not make new requests
143        if !debounce {
144            return;
145        }
146
147        reset_completion_cache(self, cx);
148
149        let Some(mut completion) = self.supermaven.update(cx, |supermaven, cx| {
150            supermaven.complete(&buffer_handle, cursor_position, cx)
151        }) else {
152            return;
153        };
154
155        self.pending_refresh = Some(cx.spawn(async move |this, cx| {
156            if debounce {
157                cx.background_executor().timer(DEBOUNCE_TIMEOUT).await;
158            }
159
160            while let Some(()) = completion.updates.next().await {
161                this.update(cx, |this, cx| {
162                    // Get the completion text and cache it
163                    if let Some(text) =
164                        this.supermaven
165                            .read(cx)
166                            .completion(&buffer_handle, cursor_position, cx)
167                    {
168                        this.completion_text = Some(text.to_string());
169
170                        this.completion_position = Some(cursor_position);
171                    }
172
173                    this.completion_id = Some(completion.id);
174                    this.buffer_id = Some(buffer_handle.entity_id());
175                    this.file_extension = buffer_handle.read(cx).file().and_then(|file| {
176                        Some(
177                            Path::new(file.file_name(cx))
178                                .extension()?
179                                .to_str()?
180                                .to_string(),
181                        )
182                    });
183                    cx.notify();
184                })?;
185            }
186            Ok(())
187        }));
188    }
189
190    fn cycle(
191        &mut self,
192        _buffer: Entity<Buffer>,
193        _cursor_position: Anchor,
194        _direction: Direction,
195        _cx: &mut Context<Self>,
196    ) {
197    }
198
199    fn accept(&mut self, _cx: &mut Context<Self>) {
200        reset_completion_cache(self, _cx);
201    }
202
203    fn discard(&mut self, _cx: &mut Context<Self>) {
204        reset_completion_cache(self, _cx);
205    }
206
207    fn suggest(
208        &mut self,
209        buffer: &Entity<Buffer>,
210        cursor_position: Anchor,
211        cx: &mut Context<Self>,
212    ) -> Option<EditPrediction> {
213        if self.buffer_id != Some(buffer.entity_id()) {
214            return None;
215        }
216
217        if self.completion_id.is_none() {
218            return None;
219        }
220
221        let completion_text = if let Some(cached_text) = &self.completion_text {
222            cached_text.as_str()
223        } else {
224            let text = self
225                .supermaven
226                .read(cx)
227                .completion(buffer, cursor_position, cx)?;
228            self.completion_text = Some(text.to_string());
229            text
230        };
231
232        // Check if the cursor is still at the same position as the completion request
233        // If we don't have a completion position stored, don't show the completion
234        if let Some(completion_position) = self.completion_position {
235            if cursor_position != completion_position {
236                return None;
237            }
238        } else {
239            return None;
240        }
241
242        let completion_text = trim_to_end_of_line_unless_leading_newline(completion_text);
243
244        let completion_text = completion_text.trim_end();
245
246        if !completion_text.trim().is_empty() {
247            let snapshot = buffer.read(cx).snapshot();
248
249            // Calculate the range from cursor to end of line correctly
250            let cursor_point = cursor_position.to_point(&snapshot);
251            let end_of_line = snapshot.anchor_after(language::Point::new(
252                cursor_point.row,
253                snapshot.line_len(cursor_point.row),
254            ));
255            let delete_range = cursor_position..end_of_line;
256
257            Some(completion_from_diff(
258                snapshot,
259                completion_text,
260                cursor_position,
261                delete_range,
262            ))
263        } else {
264            None
265        }
266    }
267}
268
269fn reset_completion_cache(
270    provider: &mut SupermavenCompletionProvider,
271    _cx: &mut Context<SupermavenCompletionProvider>,
272) {
273    provider.pending_refresh = None;
274    provider.completion_id = None;
275    provider.completion_text = None;
276    provider.completion_position = None;
277    provider.buffer_id = None;
278}
279
280fn trim_to_end_of_line_unless_leading_newline(text: &str) -> &str {
281    if has_leading_newline(text) {
282        text
283    } else if let Some(i) = text.find('\n') {
284        &text[..i]
285    } else {
286        text
287    }
288}
289
290fn has_leading_newline(text: &str) -> bool {
291    for c in text.chars() {
292        if c == '\n' {
293            return true;
294        }
295        if !c.is_whitespace() {
296            return false;
297        }
298    }
299    false
300}