1use edit_prediction_types::{EditPredictionDelegate, EditPredictionIconSet};
2use gpui::{Entity, KeyBinding, Modifiers, prelude::*};
3use indoc::indoc;
4use multi_buffer::{Anchor, MultiBufferSnapshot, ToPoint};
5use std::{ops::Range, sync::Arc};
6use text::{Point, ToOffset};
7use ui::prelude::*;
8
9use crate::{
10 AcceptEditPrediction, EditPrediction, MenuEditPredictionsPolicy, editor_tests::init_test,
11 test::editor_test_context::EditorTestContext,
12};
13
14#[gpui::test]
15async fn test_edit_prediction_insert(cx: &mut gpui::TestAppContext) {
16 init_test(cx, |_| {});
17
18 let mut cx = EditorTestContext::new(cx).await;
19 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
20 assign_editor_completion_provider(provider.clone(), &mut cx);
21 cx.set_state("let absolute_zero_celsius = ˇ;");
22
23 propose_edits(&provider, vec![(28..28, "-273.15")], &mut cx);
24 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
25
26 assert_editor_active_edit_completion(&mut cx, |_, edits| {
27 assert_eq!(edits.len(), 1);
28 assert_eq!(edits[0].1.as_ref(), "-273.15");
29 });
30
31 accept_completion(&mut cx);
32
33 cx.assert_editor_state("let absolute_zero_celsius = -273.15ˇ;")
34}
35
36#[gpui::test]
37async fn test_edit_prediction_modification(cx: &mut gpui::TestAppContext) {
38 init_test(cx, |_| {});
39
40 let mut cx = EditorTestContext::new(cx).await;
41 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
42 assign_editor_completion_provider(provider.clone(), &mut cx);
43 cx.set_state("let pi = ˇ\"foo\";");
44
45 propose_edits(&provider, vec![(9..14, "3.14159")], &mut cx);
46 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
47
48 assert_editor_active_edit_completion(&mut cx, |_, edits| {
49 assert_eq!(edits.len(), 1);
50 assert_eq!(edits[0].1.as_ref(), "3.14159");
51 });
52
53 accept_completion(&mut cx);
54
55 cx.assert_editor_state("let pi = 3.14159ˇ;")
56}
57
58#[gpui::test]
59async fn test_edit_prediction_jump_button(cx: &mut gpui::TestAppContext) {
60 init_test(cx, |_| {});
61
62 let mut cx = EditorTestContext::new(cx).await;
63 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
64 assign_editor_completion_provider(provider.clone(), &mut cx);
65
66 // Cursor is 2+ lines above the proposed edit
67 cx.set_state(indoc! {"
68 line 0
69 line ˇ1
70 line 2
71 line 3
72 line
73 "});
74
75 propose_edits(
76 &provider,
77 vec![(Point::new(4, 3)..Point::new(4, 3), " 4")],
78 &mut cx,
79 );
80
81 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
82 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
83 assert_eq!(move_target.to_point(&snapshot), Point::new(4, 3));
84 });
85
86 // When accepting, cursor is moved to the proposed location
87 accept_completion(&mut cx);
88 cx.assert_editor_state(indoc! {"
89 line 0
90 line 1
91 line 2
92 line 3
93 linˇe
94 "});
95
96 // Cursor is 2+ lines below the proposed edit
97 cx.set_state(indoc! {"
98 line 0
99 line
100 line 2
101 line 3
102 line ˇ4
103 "});
104
105 propose_edits(
106 &provider,
107 vec![(Point::new(1, 3)..Point::new(1, 3), " 1")],
108 &mut cx,
109 );
110
111 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
112 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
113 assert_eq!(move_target.to_point(&snapshot), Point::new(1, 3));
114 });
115
116 // When accepting, cursor is moved to the proposed location
117 accept_completion(&mut cx);
118 cx.assert_editor_state(indoc! {"
119 line 0
120 linˇe
121 line 2
122 line 3
123 line 4
124 "});
125}
126
127#[gpui::test]
128async fn test_edit_prediction_invalidation_range(cx: &mut gpui::TestAppContext) {
129 init_test(cx, |_| {});
130
131 let mut cx = EditorTestContext::new(cx).await;
132 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
133 assign_editor_completion_provider(provider.clone(), &mut cx);
134
135 // Cursor is 3+ lines above the proposed edit
136 cx.set_state(indoc! {"
137 line 0
138 line ˇ1
139 line 2
140 line 3
141 line 4
142 line
143 "});
144 let edit_location = Point::new(5, 3);
145
146 propose_edits(
147 &provider,
148 vec![(edit_location..edit_location, " 5")],
149 &mut cx,
150 );
151
152 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
153 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
154 assert_eq!(move_target.to_point(&snapshot), edit_location);
155 });
156
157 // If we move *towards* the completion, it stays active
158 cx.set_selections_state(indoc! {"
159 line 0
160 line 1
161 line ˇ2
162 line 3
163 line 4
164 line
165 "});
166 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
167 assert_eq!(move_target.to_point(&snapshot), edit_location);
168 });
169
170 // If we move *away* from the completion, it is discarded
171 cx.set_selections_state(indoc! {"
172 line ˇ0
173 line 1
174 line 2
175 line 3
176 line 4
177 line
178 "});
179 cx.editor(|editor, _, _| {
180 assert!(editor.active_edit_prediction.is_none());
181 });
182
183 // Cursor is 3+ lines below the proposed edit
184 cx.set_state(indoc! {"
185 line
186 line 1
187 line 2
188 line 3
189 line ˇ4
190 line 5
191 "});
192 let edit_location = Point::new(0, 3);
193
194 propose_edits(
195 &provider,
196 vec![(edit_location..edit_location, " 0")],
197 &mut cx,
198 );
199
200 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
201 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
202 assert_eq!(move_target.to_point(&snapshot), edit_location);
203 });
204
205 // If we move *towards* the completion, it stays active
206 cx.set_selections_state(indoc! {"
207 line
208 line 1
209 line 2
210 line ˇ3
211 line 4
212 line 5
213 "});
214 assert_editor_active_move_completion(&mut cx, |snapshot, move_target| {
215 assert_eq!(move_target.to_point(&snapshot), edit_location);
216 });
217
218 // If we move *away* from the completion, it is discarded
219 cx.set_selections_state(indoc! {"
220 line
221 line 1
222 line 2
223 line 3
224 line 4
225 line ˇ5
226 "});
227 cx.editor(|editor, _, _| {
228 assert!(editor.active_edit_prediction.is_none());
229 });
230}
231
232#[gpui::test]
233async fn test_edit_prediction_jump_disabled_for_non_zed_providers(cx: &mut gpui::TestAppContext) {
234 init_test(cx, |_| {});
235
236 let mut cx = EditorTestContext::new(cx).await;
237 let provider = cx.new(|_| FakeNonZedEditPredictionDelegate::default());
238 assign_editor_completion_provider_non_zed(provider.clone(), &mut cx);
239
240 // Cursor is 2+ lines above the proposed edit
241 cx.set_state(indoc! {"
242 line 0
243 line ˇ1
244 line 2
245 line 3
246 line
247 "});
248
249 propose_edits_non_zed(
250 &provider,
251 vec![(Point::new(4, 3)..Point::new(4, 3), " 4")],
252 &mut cx,
253 );
254
255 cx.update_editor(|editor, window, cx| editor.update_visible_edit_prediction(window, cx));
256
257 // For non-Zed providers, there should be no move completion (jump functionality disabled)
258 cx.editor(|editor, _, _| {
259 if let Some(completion_state) = &editor.active_edit_prediction {
260 // Should be an Edit prediction, not a Move prediction
261 match &completion_state.completion {
262 EditPrediction::Edit { .. } => {
263 // This is expected for non-Zed providers
264 }
265 EditPrediction::MoveWithin { .. } | EditPrediction::MoveOutside { .. } => {
266 panic!(
267 "Non-Zed providers should not show Move predictions (jump functionality)"
268 );
269 }
270 }
271 }
272 });
273}
274
275#[gpui::test]
276async fn test_edit_prediction_preview_cleanup_on_toggle_off(cx: &mut gpui::TestAppContext) {
277 init_test(cx, |_| {});
278
279 // Bind `ctrl-shift-a` to accept the provided edit prediction. The actual key
280 // binding here doesn't matter, we simply need to confirm that holding the
281 // binding's modifiers triggers the edit prediction preview.
282 cx.update(|cx| cx.bind_keys([KeyBinding::new("ctrl-shift-a", AcceptEditPrediction, None)]));
283
284 let mut cx = EditorTestContext::new(cx).await;
285 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
286 assign_editor_completion_provider(provider.clone(), &mut cx);
287 cx.set_state("let x = ˇ;");
288
289 propose_edits(&provider, vec![(8..8, "42")], &mut cx);
290 cx.update_editor(|editor, window, cx| {
291 editor.set_menu_edit_predictions_policy(MenuEditPredictionsPolicy::ByProvider);
292 editor.update_visible_edit_prediction(window, cx)
293 });
294
295 cx.editor(|editor, _, _| {
296 assert!(editor.has_active_edit_prediction());
297 });
298
299 // Simulate pressing the modifiers for `AcceptEditPrediction`, namely
300 // `ctrl-shift`, so that we can confirm that the edit prediction preview is
301 // activated.
302 let modifiers = Modifiers::control_shift();
303 cx.simulate_modifiers_change(modifiers);
304 cx.run_until_parked();
305
306 cx.editor(|editor, _, _| {
307 assert!(editor.edit_prediction_preview_is_active());
308 });
309
310 // Disable showing edit predictions without issuing a new modifiers changed
311 // event, to confirm that the edit prediction preview is still active.
312 cx.update_editor(|editor, window, cx| {
313 editor.set_show_edit_predictions(Some(false), window, cx);
314 });
315
316 cx.editor(|editor, _, _| {
317 assert!(!editor.has_active_edit_prediction());
318 assert!(editor.edit_prediction_preview_is_active());
319 });
320
321 // Now release the modifiers
322 // Simulate releasing all modifiers, ensuring that even with edit prediction
323 // disabled, the edit prediction preview is cleaned up.
324 cx.simulate_modifiers_change(Modifiers::none());
325 cx.run_until_parked();
326
327 cx.editor(|editor, _, _| {
328 assert!(!editor.edit_prediction_preview_is_active());
329 });
330}
331
332fn assert_editor_active_edit_completion(
333 cx: &mut EditorTestContext,
334 assert: impl FnOnce(MultiBufferSnapshot, &Vec<(Range<Anchor>, Arc<str>)>),
335) {
336 cx.editor(|editor, _, cx| {
337 let completion_state = editor
338 .active_edit_prediction
339 .as_ref()
340 .expect("editor has no active completion");
341
342 if let EditPrediction::Edit { edits, .. } = &completion_state.completion {
343 assert(editor.buffer().read(cx).snapshot(cx), edits);
344 } else {
345 panic!("expected edit completion");
346 }
347 })
348}
349
350fn assert_editor_active_move_completion(
351 cx: &mut EditorTestContext,
352 assert: impl FnOnce(MultiBufferSnapshot, Anchor),
353) {
354 cx.editor(|editor, _, cx| {
355 let completion_state = editor
356 .active_edit_prediction
357 .as_ref()
358 .expect("editor has no active completion");
359
360 if let EditPrediction::MoveWithin { target, .. } = &completion_state.completion {
361 assert(editor.buffer().read(cx).snapshot(cx), *target);
362 } else {
363 panic!("expected move completion");
364 }
365 })
366}
367
368fn accept_completion(cx: &mut EditorTestContext) {
369 cx.update_editor(|editor, window, cx| {
370 editor.accept_edit_prediction(&crate::AcceptEditPrediction, window, cx)
371 })
372}
373
374fn propose_edits<T: ToOffset>(
375 provider: &Entity<FakeEditPredictionDelegate>,
376 edits: Vec<(Range<T>, &str)>,
377 cx: &mut EditorTestContext,
378) {
379 let snapshot = cx.buffer_snapshot();
380 let edits = edits.into_iter().map(|(range, text)| {
381 let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
382 (range, text.into())
383 });
384
385 cx.update(|_, cx| {
386 provider.update(cx, |provider, _| {
387 provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
388 id: None,
389 edits: edits.collect(),
390 edit_preview: None,
391 }))
392 })
393 });
394}
395
396fn assign_editor_completion_provider(
397 provider: Entity<FakeEditPredictionDelegate>,
398 cx: &mut EditorTestContext,
399) {
400 cx.update_editor(|editor, window, cx| {
401 editor.set_edit_prediction_provider(Some(provider), window, cx);
402 })
403}
404
405fn propose_edits_non_zed<T: ToOffset>(
406 provider: &Entity<FakeNonZedEditPredictionDelegate>,
407 edits: Vec<(Range<T>, &str)>,
408 cx: &mut EditorTestContext,
409) {
410 let snapshot = cx.buffer_snapshot();
411 let edits = edits.into_iter().map(|(range, text)| {
412 let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
413 (range, text.into())
414 });
415
416 cx.update(|_, cx| {
417 provider.update(cx, |provider, _| {
418 provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
419 id: None,
420 edits: edits.collect(),
421 edit_preview: None,
422 }))
423 })
424 });
425}
426
427fn assign_editor_completion_provider_non_zed(
428 provider: Entity<FakeNonZedEditPredictionDelegate>,
429 cx: &mut EditorTestContext,
430) {
431 cx.update_editor(|editor, window, cx| {
432 editor.set_edit_prediction_provider(Some(provider), window, cx);
433 })
434}
435
436#[derive(Default, Clone)]
437pub struct FakeEditPredictionDelegate {
438 pub completion: Option<edit_prediction_types::EditPrediction>,
439}
440
441impl FakeEditPredictionDelegate {
442 pub fn set_edit_prediction(
443 &mut self,
444 completion: Option<edit_prediction_types::EditPrediction>,
445 ) {
446 self.completion = completion;
447 }
448}
449
450impl EditPredictionDelegate for FakeEditPredictionDelegate {
451 fn name() -> &'static str {
452 "fake-completion-provider"
453 }
454
455 fn display_name() -> &'static str {
456 "Fake Completion Provider"
457 }
458
459 fn show_predictions_in_menu() -> bool {
460 true
461 }
462
463 fn supports_jump_to_edit() -> bool {
464 true
465 }
466
467 fn icons(&self, _cx: &gpui::App) -> EditPredictionIconSet {
468 EditPredictionIconSet::new(IconName::ZedPredict)
469 }
470
471 fn is_enabled(
472 &self,
473 _buffer: &gpui::Entity<language::Buffer>,
474 _cursor_position: language::Anchor,
475 _cx: &gpui::App,
476 ) -> bool {
477 true
478 }
479
480 fn is_refreshing(&self, _cx: &gpui::App) -> bool {
481 false
482 }
483
484 fn refresh(
485 &mut self,
486 _buffer: gpui::Entity<language::Buffer>,
487 _cursor_position: language::Anchor,
488 _debounce: bool,
489 _cx: &mut gpui::Context<Self>,
490 ) {
491 }
492
493 fn accept(&mut self, _cx: &mut gpui::Context<Self>) {}
494
495 fn discard(
496 &mut self,
497 _reason: edit_prediction_types::EditPredictionDismissReason,
498 _cx: &mut gpui::Context<Self>,
499 ) {
500 }
501
502 fn suggest<'a>(
503 &mut self,
504 _buffer: &gpui::Entity<language::Buffer>,
505 _cursor_position: language::Anchor,
506 _cx: &mut gpui::Context<Self>,
507 ) -> Option<edit_prediction_types::EditPrediction> {
508 self.completion.clone()
509 }
510}
511
512#[derive(Default, Clone)]
513pub struct FakeNonZedEditPredictionDelegate {
514 pub completion: Option<edit_prediction_types::EditPrediction>,
515}
516
517impl FakeNonZedEditPredictionDelegate {
518 pub fn set_edit_prediction(
519 &mut self,
520 completion: Option<edit_prediction_types::EditPrediction>,
521 ) {
522 self.completion = completion;
523 }
524}
525
526impl EditPredictionDelegate for FakeNonZedEditPredictionDelegate {
527 fn name() -> &'static str {
528 "fake-non-zed-provider"
529 }
530
531 fn display_name() -> &'static str {
532 "Fake Non-Zed Provider"
533 }
534
535 fn show_predictions_in_menu() -> bool {
536 false
537 }
538
539 fn supports_jump_to_edit() -> bool {
540 false
541 }
542
543 fn icons(&self, _cx: &gpui::App) -> EditPredictionIconSet {
544 EditPredictionIconSet::new(IconName::ZedPredict)
545 }
546
547 fn is_enabled(
548 &self,
549 _buffer: &gpui::Entity<language::Buffer>,
550 _cursor_position: language::Anchor,
551 _cx: &gpui::App,
552 ) -> bool {
553 true
554 }
555
556 fn is_refreshing(&self, _cx: &gpui::App) -> bool {
557 false
558 }
559
560 fn refresh(
561 &mut self,
562 _buffer: gpui::Entity<language::Buffer>,
563 _cursor_position: language::Anchor,
564 _debounce: bool,
565 _cx: &mut gpui::Context<Self>,
566 ) {
567 }
568
569 fn accept(&mut self, _cx: &mut gpui::Context<Self>) {}
570
571 fn discard(
572 &mut self,
573 _reason: edit_prediction_types::EditPredictionDismissReason,
574 _cx: &mut gpui::Context<Self>,
575 ) {
576 }
577
578 fn suggest<'a>(
579 &mut self,
580 _buffer: &gpui::Entity<language::Buffer>,
581 _cursor_position: language::Anchor,
582 _cx: &mut gpui::Context<Self>,
583 ) -> Option<edit_prediction_types::EditPrediction> {
584 self.completion.clone()
585 }
586}