1use std::{ops::Range, sync::Arc};
2
3use client::EditPredictionUsage;
4use gpui::{App, Context, Entity, SharedString};
5use language::{Anchor, Buffer, BufferSnapshot, OffsetRangeExt};
6
7// TODO: Find a better home for `Direction`.
8//
9// This should live in an ancestor crate of `editor` and `edit_prediction`,
10// but at time of writing there isn't an obvious spot.
11#[derive(Copy, Clone, PartialEq, Eq)]
12pub enum Direction {
13 Prev,
14 Next,
15}
16
17#[derive(Clone)]
18pub enum EditPrediction {
19 /// Edits within the buffer that requested the prediction
20 Local {
21 id: Option<SharedString>,
22 edits: Vec<(Range<language::Anchor>, Arc<str>)>,
23 edit_preview: Option<language::EditPreview>,
24 },
25 /// Jump to a different file from the one that requested the prediction
26 Jump {
27 id: Option<SharedString>,
28 snapshot: language::BufferSnapshot,
29 target: language::Anchor,
30 },
31}
32
33pub enum DataCollectionState {
34 /// The provider doesn't support data collection.
35 Unsupported,
36 /// Data collection is enabled.
37 Enabled { is_project_open_source: bool },
38 /// Data collection is disabled or unanswered.
39 Disabled { is_project_open_source: bool },
40}
41
42impl DataCollectionState {
43 pub fn is_supported(&self) -> bool {
44 !matches!(self, DataCollectionState::Unsupported)
45 }
46
47 pub fn is_enabled(&self) -> bool {
48 matches!(self, DataCollectionState::Enabled { .. })
49 }
50
51 pub fn is_project_open_source(&self) -> bool {
52 match self {
53 Self::Enabled {
54 is_project_open_source,
55 }
56 | Self::Disabled {
57 is_project_open_source,
58 } => *is_project_open_source,
59 _ => false,
60 }
61 }
62}
63
64pub trait EditPredictionProvider: 'static + Sized {
65 fn name() -> &'static str;
66 fn display_name() -> &'static str;
67 fn show_completions_in_menu() -> bool;
68 fn show_tab_accept_marker() -> bool {
69 false
70 }
71 fn supports_jump_to_edit() -> bool {
72 true
73 }
74
75 fn data_collection_state(&self, _cx: &App) -> DataCollectionState {
76 DataCollectionState::Unsupported
77 }
78
79 fn usage(&self, _cx: &App) -> Option<EditPredictionUsage> {
80 None
81 }
82
83 fn toggle_data_collection(&mut self, _cx: &mut App) {}
84 fn is_enabled(
85 &self,
86 buffer: &Entity<Buffer>,
87 cursor_position: language::Anchor,
88 cx: &App,
89 ) -> bool;
90 fn is_refreshing(&self) -> bool;
91 fn refresh(
92 &mut self,
93 buffer: Entity<Buffer>,
94 cursor_position: language::Anchor,
95 debounce: bool,
96 cx: &mut Context<Self>,
97 );
98 fn cycle(
99 &mut self,
100 buffer: Entity<Buffer>,
101 cursor_position: language::Anchor,
102 direction: Direction,
103 cx: &mut Context<Self>,
104 );
105 fn accept(&mut self, cx: &mut Context<Self>);
106 fn discard(&mut self, cx: &mut Context<Self>);
107 fn did_show(&mut self, _cx: &mut Context<Self>) {}
108 fn suggest(
109 &mut self,
110 buffer: &Entity<Buffer>,
111 cursor_position: language::Anchor,
112 cx: &mut Context<Self>,
113 ) -> Option<EditPrediction>;
114}
115
116pub trait EditPredictionProviderHandle {
117 fn name(&self) -> &'static str;
118 fn display_name(&self) -> &'static str;
119 fn is_enabled(
120 &self,
121 buffer: &Entity<Buffer>,
122 cursor_position: language::Anchor,
123 cx: &App,
124 ) -> bool;
125 fn show_completions_in_menu(&self) -> bool;
126 fn show_tab_accept_marker(&self) -> bool;
127 fn supports_jump_to_edit(&self) -> bool;
128 fn data_collection_state(&self, cx: &App) -> DataCollectionState;
129 fn usage(&self, cx: &App) -> Option<EditPredictionUsage>;
130 fn toggle_data_collection(&self, cx: &mut App);
131 fn is_refreshing(&self, cx: &App) -> bool;
132 fn refresh(
133 &self,
134 buffer: Entity<Buffer>,
135 cursor_position: language::Anchor,
136 debounce: bool,
137 cx: &mut App,
138 );
139 fn cycle(
140 &self,
141 buffer: Entity<Buffer>,
142 cursor_position: language::Anchor,
143 direction: Direction,
144 cx: &mut App,
145 );
146 fn did_show(&self, cx: &mut App);
147 fn accept(&self, cx: &mut App);
148 fn discard(&self, cx: &mut App);
149 fn suggest(
150 &self,
151 buffer: &Entity<Buffer>,
152 cursor_position: language::Anchor,
153 cx: &mut App,
154 ) -> Option<EditPrediction>;
155}
156
157impl<T> EditPredictionProviderHandle for Entity<T>
158where
159 T: EditPredictionProvider,
160{
161 fn name(&self) -> &'static str {
162 T::name()
163 }
164
165 fn display_name(&self) -> &'static str {
166 T::display_name()
167 }
168
169 fn show_completions_in_menu(&self) -> bool {
170 T::show_completions_in_menu()
171 }
172
173 fn show_tab_accept_marker(&self) -> bool {
174 T::show_tab_accept_marker()
175 }
176
177 fn supports_jump_to_edit(&self) -> bool {
178 T::supports_jump_to_edit()
179 }
180
181 fn data_collection_state(&self, cx: &App) -> DataCollectionState {
182 self.read(cx).data_collection_state(cx)
183 }
184
185 fn usage(&self, cx: &App) -> Option<EditPredictionUsage> {
186 self.read(cx).usage(cx)
187 }
188
189 fn toggle_data_collection(&self, cx: &mut App) {
190 self.update(cx, |this, cx| this.toggle_data_collection(cx))
191 }
192
193 fn is_enabled(
194 &self,
195 buffer: &Entity<Buffer>,
196 cursor_position: language::Anchor,
197 cx: &App,
198 ) -> bool {
199 self.read(cx).is_enabled(buffer, cursor_position, cx)
200 }
201
202 fn is_refreshing(&self, cx: &App) -> bool {
203 self.read(cx).is_refreshing()
204 }
205
206 fn refresh(
207 &self,
208 buffer: Entity<Buffer>,
209 cursor_position: language::Anchor,
210 debounce: bool,
211 cx: &mut App,
212 ) {
213 self.update(cx, |this, cx| {
214 this.refresh(buffer, cursor_position, debounce, cx)
215 })
216 }
217
218 fn cycle(
219 &self,
220 buffer: Entity<Buffer>,
221 cursor_position: language::Anchor,
222 direction: Direction,
223 cx: &mut App,
224 ) {
225 self.update(cx, |this, cx| {
226 this.cycle(buffer, cursor_position, direction, cx)
227 })
228 }
229
230 fn accept(&self, cx: &mut App) {
231 self.update(cx, |this, cx| this.accept(cx))
232 }
233
234 fn discard(&self, cx: &mut App) {
235 self.update(cx, |this, cx| this.discard(cx))
236 }
237
238 fn did_show(&self, cx: &mut App) {
239 self.update(cx, |this, cx| this.did_show(cx))
240 }
241
242 fn suggest(
243 &self,
244 buffer: &Entity<Buffer>,
245 cursor_position: language::Anchor,
246 cx: &mut App,
247 ) -> Option<EditPrediction> {
248 self.update(cx, |this, cx| this.suggest(buffer, cursor_position, cx))
249 }
250}
251
252/// Returns edits updated based on user edits since the old snapshot. None is returned if any user
253/// edit is not a prefix of a predicted insertion.
254pub fn interpolate_edits(
255 old_snapshot: &BufferSnapshot,
256 new_snapshot: &BufferSnapshot,
257 current_edits: &[(Range<Anchor>, Arc<str>)],
258) -> Option<Vec<(Range<Anchor>, Arc<str>)>> {
259 let mut edits = Vec::new();
260
261 let mut model_edits = current_edits.iter().peekable();
262 for user_edit in new_snapshot.edits_since::<usize>(&old_snapshot.version) {
263 while let Some((model_old_range, _)) = model_edits.peek() {
264 let model_old_range = model_old_range.to_offset(old_snapshot);
265 if model_old_range.end < user_edit.old.start {
266 let (model_old_range, model_new_text) = model_edits.next().unwrap();
267 edits.push((model_old_range.clone(), model_new_text.clone()));
268 } else {
269 break;
270 }
271 }
272
273 if let Some((model_old_range, model_new_text)) = model_edits.peek() {
274 let model_old_offset_range = model_old_range.to_offset(old_snapshot);
275 if user_edit.old == model_old_offset_range {
276 let user_new_text = new_snapshot
277 .text_for_range(user_edit.new.clone())
278 .collect::<String>();
279
280 if let Some(model_suffix) = model_new_text.strip_prefix(&user_new_text) {
281 if !model_suffix.is_empty() {
282 let anchor = old_snapshot.anchor_after(user_edit.old.end);
283 edits.push((anchor..anchor, model_suffix.into()));
284 }
285
286 model_edits.next();
287 continue;
288 }
289 }
290 }
291
292 return None;
293 }
294
295 edits.extend(model_edits.cloned());
296
297 if edits.is_empty() { None } else { Some(edits) }
298}