1use edit_prediction::EditPredictionProvider;
2use gpui::{Entity, prelude::*};
3use indoc::indoc;
4use multi_buffer::{Anchor, MultiBufferSnapshot, ToPoint};
5use project::Project;
6use std::ops::Range;
7use text::{Point, ToOffset};
8
9use crate::{
10 EditPrediction, editor_tests::init_test, test::editor_test_context::EditorTestContext,
11};
12
13#[gpui::test]
14async fn test_edit_prediction_insert(cx: &mut gpui::TestAppContext) {
15 init_test(cx, |_| {});
16
17 let mut cx = EditorTestContext::new(cx).await;
18 let provider = cx.new(|_| FakeEditPredictionProvider::default());
19 assign_editor_completion_provider(provider.clone(), &mut cx);
20 cx.set_state("let absolute_zero_celsius = ˇ;");
21
22 propose_edits(&provider, vec![(28..28, "-273.15")], &mut cx);
23 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
24
25 assert_editor_active_edit_completion(&mut cx, |_, edits| {
26 assert_eq!(edits.len(), 1);
27 assert_eq!(edits[0].1.as_str(), "-273.15");
28 });
29
30 accept_completion(&mut cx);
31
32 cx.assert_editor_state("let absolute_zero_celsius = -273.15ˇ;")
33}
34
35#[gpui::test]
36async fn test_edit_prediction_modification(cx: &mut gpui::TestAppContext) {
37 init_test(cx, |_| {});
38
39 let mut cx = EditorTestContext::new(cx).await;
40 let provider = cx.new(|_| FakeEditPredictionProvider::default());
41 assign_editor_completion_provider(provider.clone(), &mut cx);
42 cx.set_state("let pi = ˇ\"foo\";");
43
44 propose_edits(&provider, vec![(9..14, "3.14159")], &mut cx);
45 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
46
47 assert_editor_active_edit_completion(&mut cx, |_, edits| {
48 assert_eq!(edits.len(), 1);
49 assert_eq!(edits[0].1.as_str(), "3.14159");
50 });
51
52 accept_completion(&mut cx);
53
54 cx.assert_editor_state("let pi = 3.14159ˇ;")
55}
56
57#[gpui::test]
58async fn test_edit_prediction_jump_button(cx: &mut gpui::TestAppContext) {
59 init_test(cx, |_| {});
60
61 let mut cx = EditorTestContext::new(cx).await;
62 let provider = cx.new(|_| FakeEditPredictionProvider::default());
63 assign_editor_completion_provider(provider.clone(), &mut cx);
64
65 // Cursor is 2+ lines above the proposed edit
66 cx.set_state(indoc! {"
67 line 0
68 line ˇ1
69 line 2
70 line 3
71 line
72 "});
73
74 propose_edits(
75 &provider,
76 vec![(Point::new(4, 3)..Point::new(4, 3), " 4")],
77 &mut cx,
78 );
79
80 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
81 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
82 assert_eq!(move_target.to_point(&snapshot), Point::new(4, 3));
83 });
84
85 // When accepting, cursor is moved to the proposed location
86 accept_completion(&mut cx);
87 cx.assert_editor_state(indoc! {"
88 line 0
89 line 1
90 line 2
91 line 3
92 linˇe
93 "});
94
95 // Cursor is 2+ lines below the proposed edit
96 cx.set_state(indoc! {"
97 line 0
98 line
99 line 2
100 line 3
101 line ˇ4
102 "});
103
104 propose_edits(
105 &provider,
106 vec![(Point::new(1, 3)..Point::new(1, 3), " 1")],
107 &mut cx,
108 );
109
110 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
111 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
112 assert_eq!(move_target.to_point(&snapshot), Point::new(1, 3));
113 });
114
115 // When accepting, cursor is moved to the proposed location
116 accept_completion(&mut cx);
117 cx.assert_editor_state(indoc! {"
118 line 0
119 linˇe
120 line 2
121 line 3
122 line 4
123 "});
124}
125
126#[gpui::test]
127async fn test_edit_prediction_invalidation_range(cx: &mut gpui::TestAppContext) {
128 init_test(cx, |_| {});
129
130 let mut cx = EditorTestContext::new(cx).await;
131 let provider = cx.new(|_| FakeEditPredictionProvider::default());
132 assign_editor_completion_provider(provider.clone(), &mut cx);
133
134 // Cursor is 3+ lines above the proposed edit
135 cx.set_state(indoc! {"
136 line 0
137 line ˇ1
138 line 2
139 line 3
140 line 4
141 line
142 "});
143 let edit_location = Point::new(5, 3);
144
145 propose_edits(
146 &provider,
147 vec![(edit_location..edit_location, " 5")],
148 &mut cx,
149 );
150
151 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
152 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
153 assert_eq!(move_target.to_point(&snapshot), edit_location);
154 });
155
156 // If we move *towards* the completion, it stays active
157 cx.set_selections_state(indoc! {"
158 line 0
159 line 1
160 line ˇ2
161 line 3
162 line 4
163 line
164 "});
165 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
166 assert_eq!(move_target.to_point(&snapshot), edit_location);
167 });
168
169 // If we move *away* from the completion, it is discarded
170 cx.set_selections_state(indoc! {"
171 line ˇ0
172 line 1
173 line 2
174 line 3
175 line 4
176 line
177 "});
178 cx.editor(|editor, _, _| {
179 assert!(editor.active_edit_prediction.is_none());
180 });
181
182 // Cursor is 3+ lines below the proposed edit
183 cx.set_state(indoc! {"
184 line
185 line 1
186 line 2
187 line 3
188 line ˇ4
189 line 5
190 "});
191 let edit_location = Point::new(0, 3);
192
193 propose_edits(
194 &provider,
195 vec![(edit_location..edit_location, " 0")],
196 &mut cx,
197 );
198
199 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
200 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
201 assert_eq!(move_target.to_point(&snapshot), edit_location);
202 });
203
204 // If we move *towards* the completion, it stays active
205 cx.set_selections_state(indoc! {"
206 line
207 line 1
208 line 2
209 line ˇ3
210 line 4
211 line 5
212 "});
213 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
214 assert_eq!(move_target.to_point(&snapshot), edit_location);
215 });
216
217 // If we move *away* from the completion, it is discarded
218 cx.set_selections_state(indoc! {"
219 line
220 line 1
221 line 2
222 line 3
223 line 4
224 line ˇ5
225 "});
226 cx.editor(|editor, _, _| {
227 assert!(editor.active_edit_prediction.is_none());
228 });
229}
230
231fn assert_editor_active_edit_completion(
232 cx: &mut EditorTestContext,
233 assert: impl FnOnce(MultiBufferSnapshot, &Vec<(Range<Anchor>, String)>),
234) {
235 cx.editor(|editor, _, cx| {
236 let completion_state = editor
237 .active_edit_prediction
238 .as_ref()
239 .expect("editor has no active completion");
240
241 if let EditPrediction::Edit { edits, .. } = &completion_state.completion {
242 assert(editor.buffer().read(cx).snapshot(cx), edits);
243 } else {
244 panic!("expected edit completion");
245 }
246 })
247}
248
249fn assert_editor_active_move_completion(
250 cx: &mut EditorTestContext,
251 assert: impl FnOnce(MultiBufferSnapshot, Anchor),
252) {
253 cx.editor(|editor, _, cx| {
254 let completion_state = editor
255 .active_edit_prediction
256 .as_ref()
257 .expect("editor has no active completion");
258
259 if let EditPrediction::Move { target, .. } = &completion_state.completion {
260 assert(editor.buffer().read(cx).snapshot(cx), *target);
261 } else {
262 panic!("expected move completion");
263 }
264 })
265}
266
267fn accept_completion(cx: &mut EditorTestContext) {
268 cx.update_editor(|editor, window, cx| {
269 editor.accept_edit_prediction(&crate::AcceptEditPrediction, window, cx)
270 })
271}
272
273fn propose_edits<T: ToOffset>(
274 provider: &Entity<FakeEditPredictionProvider>,
275 edits: Vec<(Range<T>, &str)>,
276 cx: &mut EditorTestContext,
277) {
278 let snapshot = cx.buffer_snapshot();
279 let edits = edits.into_iter().map(|(range, text)| {
280 let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
281 (range, text.into())
282 });
283
284 cx.update(|_, cx| {
285 provider.update(cx, |provider, _| {
286 provider.set_edit_prediction(Some(edit_prediction::EditPrediction {
287 id: None,
288 edits: edits.collect(),
289 edit_preview: None,
290 }))
291 })
292 });
293}
294
295fn assign_editor_completion_provider(
296 provider: Entity<FakeEditPredictionProvider>,
297 cx: &mut EditorTestContext,
298) {
299 cx.update_editor(|editor, window, cx| {
300 editor.set_edit_prediction_provider(Some(provider), window, cx);
301 })
302}
303
304#[derive(Default, Clone)]
305pub struct FakeEditPredictionProvider {
306 pub completion: Option<edit_prediction::EditPrediction>,
307}
308
309impl FakeEditPredictionProvider {
310 pub fn set_edit_prediction(&mut self, completion: Option<edit_prediction::EditPrediction>) {
311 self.completion = completion;
312 }
313}
314
315impl EditPredictionProvider for FakeEditPredictionProvider {
316 fn name() -> &'static str {
317 "fake-completion-provider"
318 }
319
320 fn display_name() -> &'static str {
321 "Fake Completion Provider"
322 }
323
324 fn show_completions_in_menu() -> bool {
325 false
326 }
327
328 fn is_enabled(
329 &self,
330 _buffer: &gpui::Entity<language::Buffer>,
331 _cursor_position: language::Anchor,
332 _cx: &gpui::App,
333 ) -> bool {
334 true
335 }
336
337 fn is_refreshing(&self) -> bool {
338 false
339 }
340
341 fn refresh(
342 &mut self,
343 _project: Option<Entity<Project>>,
344 _buffer: gpui::Entity<language::Buffer>,
345 _cursor_position: language::Anchor,
346 _debounce: bool,
347 _cx: &mut gpui::Context<Self>,
348 ) {
349 }
350
351 fn cycle(
352 &mut self,
353 _buffer: gpui::Entity<language::Buffer>,
354 _cursor_position: language::Anchor,
355 _direction: edit_prediction::Direction,
356 _cx: &mut gpui::Context<Self>,
357 ) {
358 }
359
360 fn accept(&mut self, _cx: &mut gpui::Context<Self>) {}
361
362 fn discard(&mut self, _cx: &mut gpui::Context<Self>) {}
363
364 fn suggest<'a>(
365 &mut self,
366 _buffer: &gpui::Entity<language::Buffer>,
367 _cursor_position: language::Anchor,
368 _cx: &mut gpui::Context<Self>,
369 ) -> Option<edit_prediction::EditPrediction> {
370 self.completion.clone()
371 }
372}