1use std::ops::Range;
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>, String)>,
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 const fn is_supported(&self) -> bool {
44 !matches!(self, DataCollectionState::Unsupported)
45 }
46
47 pub const fn is_enabled(&self) -> bool {
48 matches!(self, DataCollectionState::Enabled { .. })
49 }
50
51 pub const 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 suggest(
108 &mut self,
109 buffer: &Entity<Buffer>,
110 cursor_position: language::Anchor,
111 cx: &mut Context<Self>,
112 ) -> Option<EditPrediction>;
113}
114
115pub trait EditPredictionProviderHandle {
116 fn name(&self) -> &'static str;
117 fn display_name(&self) -> &'static str;
118 fn is_enabled(
119 &self,
120 buffer: &Entity<Buffer>,
121 cursor_position: language::Anchor,
122 cx: &App,
123 ) -> bool;
124 fn show_completions_in_menu(&self) -> bool;
125 fn show_tab_accept_marker(&self) -> bool;
126 fn supports_jump_to_edit(&self) -> bool;
127 fn data_collection_state(&self, cx: &App) -> DataCollectionState;
128 fn usage(&self, cx: &App) -> Option<EditPredictionUsage>;
129 fn toggle_data_collection(&self, cx: &mut App);
130 fn is_refreshing(&self, cx: &App) -> bool;
131 fn refresh(
132 &self,
133 buffer: Entity<Buffer>,
134 cursor_position: language::Anchor,
135 debounce: bool,
136 cx: &mut App,
137 );
138 fn cycle(
139 &self,
140 buffer: Entity<Buffer>,
141 cursor_position: language::Anchor,
142 direction: Direction,
143 cx: &mut App,
144 );
145 fn accept(&self, cx: &mut App);
146 fn discard(&self, cx: &mut App);
147 fn suggest(
148 &self,
149 buffer: &Entity<Buffer>,
150 cursor_position: language::Anchor,
151 cx: &mut App,
152 ) -> Option<EditPrediction>;
153}
154
155impl<T> EditPredictionProviderHandle for Entity<T>
156where
157 T: EditPredictionProvider,
158{
159 fn name(&self) -> &'static str {
160 T::name()
161 }
162
163 fn display_name(&self) -> &'static str {
164 T::display_name()
165 }
166
167 fn show_completions_in_menu(&self) -> bool {
168 T::show_completions_in_menu()
169 }
170
171 fn show_tab_accept_marker(&self) -> bool {
172 T::show_tab_accept_marker()
173 }
174
175 fn supports_jump_to_edit(&self) -> bool {
176 T::supports_jump_to_edit()
177 }
178
179 fn data_collection_state(&self, cx: &App) -> DataCollectionState {
180 self.read(cx).data_collection_state(cx)
181 }
182
183 fn usage(&self, cx: &App) -> Option<EditPredictionUsage> {
184 self.read(cx).usage(cx)
185 }
186
187 fn toggle_data_collection(&self, cx: &mut App) {
188 self.update(cx, |this, cx| this.toggle_data_collection(cx))
189 }
190
191 fn is_enabled(
192 &self,
193 buffer: &Entity<Buffer>,
194 cursor_position: language::Anchor,
195 cx: &App,
196 ) -> bool {
197 self.read(cx).is_enabled(buffer, cursor_position, cx)
198 }
199
200 fn is_refreshing(&self, cx: &App) -> bool {
201 self.read(cx).is_refreshing()
202 }
203
204 fn refresh(
205 &self,
206 buffer: Entity<Buffer>,
207 cursor_position: language::Anchor,
208 debounce: bool,
209 cx: &mut App,
210 ) {
211 self.update(cx, |this, cx| {
212 this.refresh(buffer, cursor_position, debounce, cx)
213 })
214 }
215
216 fn cycle(
217 &self,
218 buffer: Entity<Buffer>,
219 cursor_position: language::Anchor,
220 direction: Direction,
221 cx: &mut App,
222 ) {
223 self.update(cx, |this, cx| {
224 this.cycle(buffer, cursor_position, direction, cx)
225 })
226 }
227
228 fn accept(&self, cx: &mut App) {
229 self.update(cx, |this, cx| this.accept(cx))
230 }
231
232 fn discard(&self, cx: &mut App) {
233 self.update(cx, |this, cx| this.discard(cx))
234 }
235
236 fn suggest(
237 &self,
238 buffer: &Entity<Buffer>,
239 cursor_position: language::Anchor,
240 cx: &mut App,
241 ) -> Option<EditPrediction> {
242 self.update(cx, |this, cx| this.suggest(buffer, cursor_position, cx))
243 }
244}
245
246/// Returns edits updated based on user edits since the old snapshot. None is returned if any user
247/// edit is not a prefix of a predicted insertion.
248pub fn interpolate_edits(
249 old_snapshot: &BufferSnapshot,
250 new_snapshot: &BufferSnapshot,
251 current_edits: &[(Range<Anchor>, String)],
252) -> Option<Vec<(Range<Anchor>, String)>> {
253 let mut edits = Vec::new();
254
255 let mut model_edits = current_edits.iter().peekable();
256 for user_edit in new_snapshot.edits_since::<usize>(&old_snapshot.version) {
257 while let Some((model_old_range, _)) = model_edits.peek() {
258 let model_old_range = model_old_range.to_offset(old_snapshot);
259 if model_old_range.end < user_edit.old.start {
260 let (model_old_range, model_new_text) = model_edits.next().unwrap();
261 edits.push((model_old_range.clone(), model_new_text.clone()));
262 } else {
263 break;
264 }
265 }
266
267 if let Some((model_old_range, model_new_text)) = model_edits.peek() {
268 let model_old_offset_range = model_old_range.to_offset(old_snapshot);
269 if user_edit.old == model_old_offset_range {
270 let user_new_text = new_snapshot
271 .text_for_range(user_edit.new.clone())
272 .collect::<String>();
273
274 if let Some(model_suffix) = model_new_text.strip_prefix(&user_new_text) {
275 if !model_suffix.is_empty() {
276 let anchor = old_snapshot.anchor_after(user_edit.old.end);
277 edits.push((anchor..anchor, model_suffix.to_string()));
278 }
279
280 model_edits.next();
281 continue;
282 }
283 }
284 }
285
286 return None;
287 }
288
289 edits.extend(model_edits.cloned());
290
291 if edits.is_empty() { None } else { Some(edits) }
292}