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
231#[gpui::test]
232async fn test_edit_prediction_jump_disabled_for_non_zed_providers(cx: &mut gpui::TestAppContext) {
233 init_test(cx, |_| {});
234
235 let mut cx = EditorTestContext::new(cx).await;
236 let provider = cx.new(|_| FakeNonZedEditPredictionProvider::default());
237 assign_editor_completion_provider_non_zed(provider.clone(), &mut cx);
238
239 // Cursor is 2+ lines above the proposed edit
240 cx.set_state(indoc! {"
241 line 0
242 line ˇ1
243 line 2
244 line 3
245 line
246 "});
247
248 propose_edits_non_zed(
249 &provider,
250 vec![(Point::new(4, 3)..Point::new(4, 3), " 4")],
251 &mut cx,
252 );
253
254 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
255
256 // For non-Zed providers, there should be no move completion (jump functionality disabled)
257 cx.editor(|editor, _, _| {
258 if let Some(completion_state) = &editor.active_edit_prediction {
259 // Should be an Edit prediction, not a Move prediction
260 match &completion_state.completion {
261 EditPrediction::Edit { .. } => {
262 // This is expected for non-Zed providers
263 }
264 EditPrediction::Move { .. } => {
265 panic!(
266 "Non-Zed providers should not show Move predictions (jump functionality)"
267 );
268 }
269 }
270 }
271 });
272}
273
274fn assert_editor_active_edit_completion(
275 cx: &mut EditorTestContext,
276 assert: impl FnOnce(MultiBufferSnapshot, &Vec<(Range<Anchor>, String)>),
277) {
278 cx.editor(|editor, _, cx| {
279 let completion_state = editor
280 .active_edit_prediction
281 .as_ref()
282 .expect("editor has no active completion");
283
284 if let EditPrediction::Edit { edits, .. } = &completion_state.completion {
285 assert(editor.buffer().read(cx).snapshot(cx), edits);
286 } else {
287 panic!("expected edit completion");
288 }
289 })
290}
291
292fn assert_editor_active_move_completion(
293 cx: &mut EditorTestContext,
294 assert: impl FnOnce(MultiBufferSnapshot, Anchor),
295) {
296 cx.editor(|editor, _, cx| {
297 let completion_state = editor
298 .active_edit_prediction
299 .as_ref()
300 .expect("editor has no active completion");
301
302 if let EditPrediction::Move { target, .. } = &completion_state.completion {
303 assert(editor.buffer().read(cx).snapshot(cx), *target);
304 } else {
305 panic!("expected move completion");
306 }
307 })
308}
309
310fn accept_completion(cx: &mut EditorTestContext) {
311 cx.update_editor(|editor, window, cx| {
312 editor.accept_edit_prediction(&crate::AcceptEditPrediction, window, cx)
313 })
314}
315
316fn propose_edits<T: ToOffset>(
317 provider: &Entity<FakeEditPredictionProvider>,
318 edits: Vec<(Range<T>, &str)>,
319 cx: &mut EditorTestContext,
320) {
321 let snapshot = cx.buffer_snapshot();
322 let edits = edits.into_iter().map(|(range, text)| {
323 let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
324 (range, text.into())
325 });
326
327 cx.update(|_, cx| {
328 provider.update(cx, |provider, _| {
329 provider.set_edit_prediction(Some(edit_prediction::EditPrediction {
330 id: None,
331 edits: edits.collect(),
332 edit_preview: None,
333 }))
334 })
335 });
336}
337
338fn assign_editor_completion_provider(
339 provider: Entity<FakeEditPredictionProvider>,
340 cx: &mut EditorTestContext,
341) {
342 cx.update_editor(|editor, window, cx| {
343 editor.set_edit_prediction_provider(Some(provider), window, cx);
344 })
345}
346
347fn propose_edits_non_zed<T: ToOffset>(
348 provider: &Entity<FakeNonZedEditPredictionProvider>,
349 edits: Vec<(Range<T>, &str)>,
350 cx: &mut EditorTestContext,
351) {
352 let snapshot = cx.buffer_snapshot();
353 let edits = edits.into_iter().map(|(range, text)| {
354 let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
355 (range, text.into())
356 });
357
358 cx.update(|_, cx| {
359 provider.update(cx, |provider, _| {
360 provider.set_edit_prediction(Some(edit_prediction::EditPrediction {
361 id: None,
362 edits: edits.collect(),
363 edit_preview: None,
364 }))
365 })
366 });
367}
368
369fn assign_editor_completion_provider_non_zed(
370 provider: Entity<FakeNonZedEditPredictionProvider>,
371 cx: &mut EditorTestContext,
372) {
373 cx.update_editor(|editor, window, cx| {
374 editor.set_edit_prediction_provider(Some(provider), window, cx);
375 })
376}
377
378#[derive(Default, Clone)]
379pub struct FakeEditPredictionProvider {
380 pub completion: Option<edit_prediction::EditPrediction>,
381}
382
383impl FakeEditPredictionProvider {
384 pub fn set_edit_prediction(&mut self, completion: Option<edit_prediction::EditPrediction>) {
385 self.completion = completion;
386 }
387}
388
389impl EditPredictionProvider for FakeEditPredictionProvider {
390 fn name() -> &'static str {
391 "fake-completion-provider"
392 }
393
394 fn display_name() -> &'static str {
395 "Fake Completion Provider"
396 }
397
398 fn show_completions_in_menu() -> bool {
399 false
400 }
401
402 fn supports_jump_to_edit() -> bool {
403 true
404 }
405
406 fn is_enabled(
407 &self,
408 _buffer: &gpui::Entity<language::Buffer>,
409 _cursor_position: language::Anchor,
410 _cx: &gpui::App,
411 ) -> bool {
412 true
413 }
414
415 fn is_refreshing(&self) -> bool {
416 false
417 }
418
419 fn refresh(
420 &mut self,
421 _project: Option<Entity<Project>>,
422 _buffer: gpui::Entity<language::Buffer>,
423 _cursor_position: language::Anchor,
424 _debounce: bool,
425 _cx: &mut gpui::Context<Self>,
426 ) {
427 }
428
429 fn cycle(
430 &mut self,
431 _buffer: gpui::Entity<language::Buffer>,
432 _cursor_position: language::Anchor,
433 _direction: edit_prediction::Direction,
434 _cx: &mut gpui::Context<Self>,
435 ) {
436 }
437
438 fn accept(&mut self, _cx: &mut gpui::Context<Self>) {}
439
440 fn discard(&mut self, _cx: &mut gpui::Context<Self>) {}
441
442 fn suggest<'a>(
443 &mut self,
444 _buffer: &gpui::Entity<language::Buffer>,
445 _cursor_position: language::Anchor,
446 _cx: &mut gpui::Context<Self>,
447 ) -> Option<edit_prediction::EditPrediction> {
448 self.completion.clone()
449 }
450}
451
452#[derive(Default, Clone)]
453pub struct FakeNonZedEditPredictionProvider {
454 pub completion: Option<edit_prediction::EditPrediction>,
455}
456
457impl FakeNonZedEditPredictionProvider {
458 pub fn set_edit_prediction(&mut self, completion: Option<edit_prediction::EditPrediction>) {
459 self.completion = completion;
460 }
461}
462
463impl EditPredictionProvider for FakeNonZedEditPredictionProvider {
464 fn name() -> &'static str {
465 "fake-non-zed-provider"
466 }
467
468 fn display_name() -> &'static str {
469 "Fake Non-Zed Provider"
470 }
471
472 fn show_completions_in_menu() -> bool {
473 false
474 }
475
476 fn supports_jump_to_edit() -> bool {
477 false
478 }
479
480 fn is_enabled(
481 &self,
482 _buffer: &gpui::Entity<language::Buffer>,
483 _cursor_position: language::Anchor,
484 _cx: &gpui::App,
485 ) -> bool {
486 true
487 }
488
489 fn is_refreshing(&self) -> bool {
490 false
491 }
492
493 fn refresh(
494 &mut self,
495 _project: Option<Entity<Project>>,
496 _buffer: gpui::Entity<language::Buffer>,
497 _cursor_position: language::Anchor,
498 _debounce: bool,
499 _cx: &mut gpui::Context<Self>,
500 ) {
501 }
502
503 fn cycle(
504 &mut self,
505 _buffer: gpui::Entity<language::Buffer>,
506 _cursor_position: language::Anchor,
507 _direction: edit_prediction::Direction,
508 _cx: &mut gpui::Context<Self>,
509 ) {
510 }
511
512 fn accept(&mut self, _cx: &mut gpui::Context<Self>) {}
513
514 fn discard(&mut self, _cx: &mut gpui::Context<Self>) {}
515
516 fn suggest<'a>(
517 &mut self,
518 _buffer: &gpui::Entity<language::Buffer>,
519 _cursor_position: language::Anchor,
520 _cx: &mut gpui::Context<Self>,
521 ) -> Option<edit_prediction::EditPrediction> {
522 self.completion.clone()
523 }
524}