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    file_extension: Option<String>,
 23    pending_refresh: Option<Task<Result<()>>>,
 24}
 25
 26impl SupermavenCompletionProvider {
 27    pub fn new(supermaven: Entity<Supermaven>) -> Self {
 28        Self {
 29            supermaven,
 30            buffer_id: None,
 31            completion_id: None,
 32            file_extension: None,
 33            pending_refresh: None,
 34        }
 35    }
 36}
 37
 38// Computes the edit prediction from the difference between the completion text.
 39// this is defined by greedily matching the buffer text against the completion text, with any leftover buffer placed at the end.
 40// for example, given the completion text "moo cows are cool" and the buffer text "cowsre pool", the completion state would be
 41// the inlays "moo ", " a", and "cool" which will render as "[moo ]cows[ a]re [cool]pool" in the editor.
 42fn completion_from_diff(
 43    snapshot: BufferSnapshot,
 44    completion_text: &str,
 45    position: Anchor,
 46    delete_range: Range<Anchor>,
 47) -> EditPrediction {
 48    let buffer_text = snapshot.text_for_range(delete_range).collect::<String>();
 49
 50    let mut edits: Vec<(Range<language::Anchor>, String)> = Vec::new();
 51
 52    let completion_graphemes: Vec<&str> = completion_text.graphemes(true).collect();
 53    let buffer_graphemes: Vec<&str> = buffer_text.graphemes(true).collect();
 54
 55    let mut offset = position.to_offset(&snapshot);
 56
 57    let mut i = 0;
 58    let mut j = 0;
 59    while i < completion_graphemes.len() && j < buffer_graphemes.len() {
 60        // find the next instance of the buffer text in the completion text.
 61        let k = completion_graphemes[i..]
 62            .iter()
 63            .position(|c| *c == buffer_graphemes[j]);
 64        match k {
 65            Some(k) => {
 66                if k != 0 {
 67                    let offset = snapshot.anchor_after(offset);
 68                    // the range from the current position to item is an inlay.
 69                    let edit = (offset..offset, completion_graphemes[i..i + k].join(""));
 70                    edits.push(edit);
 71                }
 72                i += k + 1;
 73                j += 1;
 74                offset.add_assign(buffer_graphemes[j - 1].len());
 75            }
 76            None => {
 77                // there are no more matching completions, so drop the remaining
 78                // completion text as an inlay.
 79                break;
 80            }
 81        }
 82    }
 83
 84    if j == buffer_graphemes.len() && i < completion_graphemes.len() {
 85        let offset = snapshot.anchor_after(offset);
 86        // there is leftover completion text, so drop it as an inlay.
 87        let edit_range = offset..offset;
 88        let edit_text = completion_graphemes[i..].join("");
 89        edits.push((edit_range, edit_text));
 90    }
 91
 92    EditPrediction {
 93        id: None,
 94        edits,
 95        edit_preview: None,
 96    }
 97}
 98
 99impl EditPredictionProvider for SupermavenCompletionProvider {
100    fn name() -> &'static str {
101        "supermaven"
102    }
103
104    fn display_name() -> &'static str {
105        "Supermaven"
106    }
107
108    fn show_completions_in_menu() -> bool {
109        true
110    }
111
112    fn show_tab_accept_marker() -> bool {
113        true
114    }
115
116    fn supports_jump_to_edit() -> bool {
117        false
118    }
119
120    fn is_enabled(&self, _buffer: &Entity<Buffer>, _cursor_position: Anchor, cx: &App) -> bool {
121        self.supermaven.read(cx).is_enabled()
122    }
123
124    fn is_refreshing(&self) -> bool {
125        self.pending_refresh.is_some() && self.completion_id.is_none()
126    }
127
128    fn refresh(
129        &mut self,
130        _project: Option<Entity<Project>>,
131        buffer_handle: Entity<Buffer>,
132        cursor_position: Anchor,
133        debounce: bool,
134        cx: &mut Context<Self>,
135    ) {
136        let Some(mut completion) = self.supermaven.update(cx, |supermaven, cx| {
137            supermaven.complete(&buffer_handle, cursor_position, cx)
138        }) else {
139            return;
140        };
141
142        self.pending_refresh = Some(cx.spawn(async move |this, cx| {
143            if debounce {
144                cx.background_executor().timer(DEBOUNCE_TIMEOUT).await;
145            }
146
147            while let Some(()) = completion.updates.next().await {
148                this.update(cx, |this, cx| {
149                    this.completion_id = Some(completion.id);
150                    this.buffer_id = Some(buffer_handle.entity_id());
151                    this.file_extension = buffer_handle.read(cx).file().and_then(|file| {
152                        Some(
153                            Path::new(file.file_name(cx))
154                                .extension()?
155                                .to_str()?
156                                .to_string(),
157                        )
158                    });
159                    this.pending_refresh = None;
160                    cx.notify();
161                })?;
162            }
163            Ok(())
164        }));
165    }
166
167    fn cycle(
168        &mut self,
169        _buffer: Entity<Buffer>,
170        _cursor_position: Anchor,
171        _direction: Direction,
172        _cx: &mut Context<Self>,
173    ) {
174    }
175
176    fn accept(&mut self, _cx: &mut Context<Self>) {
177        self.pending_refresh = None;
178        self.completion_id = None;
179    }
180
181    fn discard(&mut self, _cx: &mut Context<Self>) {
182        self.pending_refresh = None;
183        self.completion_id = None;
184    }
185
186    fn suggest(
187        &mut self,
188        buffer: &Entity<Buffer>,
189        cursor_position: Anchor,
190        cx: &mut Context<Self>,
191    ) -> Option<EditPrediction> {
192        let completion_text = self
193            .supermaven
194            .read(cx)
195            .completion(buffer, cursor_position, cx)?;
196
197        let completion_text = trim_to_end_of_line_unless_leading_newline(completion_text);
198
199        let completion_text = completion_text.trim_end();
200
201        if !completion_text.trim().is_empty() {
202            let snapshot = buffer.read(cx).snapshot();
203            let mut point = cursor_position.to_point(&snapshot);
204            point.column = snapshot.line_len(point.row);
205            let range = cursor_position..snapshot.anchor_after(point);
206
207            Some(completion_from_diff(
208                snapshot,
209                completion_text,
210                cursor_position,
211                range,
212            ))
213        } else {
214            None
215        }
216    }
217}
218
219fn trim_to_end_of_line_unless_leading_newline(text: &str) -> &str {
220    if has_leading_newline(text) {
221        text
222    } else if let Some(i) = text.find('\n') {
223        &text[..i]
224    } else {
225        text
226    }
227}
228
229fn has_leading_newline(text: &str) -> bool {
230    for c in text.chars() {
231        if c == '\n' {
232            return true;
233        }
234        if !c.is_whitespace() {
235            return false;
236        }
237    }
238    false
239}