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