1use edit_prediction_types::EditPredictionDelegate;
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};
7
8use crate::{
9 AcceptEditPrediction, EditPrediction, MenuEditPredictionsPolicy, editor_tests::init_test,
10 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(|_| FakeEditPredictionDelegate::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_ref(), "-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(|_| FakeEditPredictionDelegate::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_ref(), "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(|_| FakeEditPredictionDelegate::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(|_| FakeEditPredictionDelegate::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(|_| FakeNonZedEditPredictionDelegate::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::MoveWithin { .. } | EditPrediction::MoveOutside { .. } => {
265 panic!(
266 "Non-Zed providers should not show Move predictions (jump functionality)"
267 );
268 }
269 }
270 }
271 });
272}
273
274#[gpui::test]
275async fn test_edit_prediction_preview_cleanup_on_toggle_off(cx: &mut gpui::TestAppContext) {
276 init_test(cx, |_| {});
277
278 // Bind `ctrl-shift-a` to accept the provided edit prediction. The actual key
279 // binding here doesn't matter, we simply need to confirm that holding the
280 // binding's modifiers triggers the edit prediction preview.
281 cx.update(|cx| cx.bind_keys([KeyBinding::new("ctrl-shift-a", AcceptEditPrediction, None)]));
282
283 let mut cx = EditorTestContext::new(cx).await;
284 let provider = cx.new(|_| FakeEditPredictionDelegate::default());
285 assign_editor_completion_provider(provider.clone(), &mut cx);
286 cx.set_state("let x = ˇ;");
287
288 propose_edits(&provider, vec![(8..8, "42")], &mut cx);
289 cx.update_editor(|editor, window, cx| {
290 editor.set_menu_edit_predictions_policy(MenuEditPredictionsPolicy::ByProvider);
291 editor.update_visible_edit_prediction(window, cx)
292 });
293
294 cx.editor(|editor, _, _| {
295 assert!(editor.has_active_edit_prediction());
296 });
297
298 // Simulate pressing the modifiers for `AcceptEditPrediction`, namely
299 // `ctrl-shift`, so that we can confirm that the edit prediction preview is
300 // activated.
301 let modifiers = Modifiers::control_shift();
302 cx.simulate_modifiers_change(modifiers);
303 cx.run_until_parked();
304
305 cx.editor(|editor, _, _| {
306 assert!(editor.edit_prediction_preview_is_active());
307 });
308
309 // Disable showing edit predictions without issuing a new modifiers changed
310 // event, to confirm that the edit prediction preview is still active.
311 cx.update_editor(|editor, window, cx| {
312 editor.set_show_edit_predictions(Some(false), window, cx);
313 });
314
315 cx.editor(|editor, _, _| {
316 assert!(!editor.has_active_edit_prediction());
317 assert!(editor.edit_prediction_preview_is_active());
318 });
319
320 // Now release the modifiers
321 // Simulate releasing all modifiers, ensuring that even with edit prediction
322 // disabled, the edit prediction preview is cleaned up.
323 cx.simulate_modifiers_change(Modifiers::none());
324 cx.run_until_parked();
325
326 cx.editor(|editor, _, _| {
327 assert!(!editor.edit_prediction_preview_is_active());
328 });
329}
330
331fn assert_editor_active_edit_completion(
332 cx: &mut EditorTestContext,
333 assert: impl FnOnce(MultiBufferSnapshot, &Vec<(Range<Anchor>, Arc<str>)>),
334) {
335 cx.editor(|editor, _, cx| {
336 let completion_state = editor
337 .active_edit_prediction
338 .as_ref()
339 .expect("editor has no active completion");
340
341 if let EditPrediction::Edit { edits, .. } = &completion_state.completion {
342 assert(editor.buffer().read(cx).snapshot(cx), edits);
343 } else {
344 panic!("expected edit completion");
345 }
346 })
347}
348
349fn assert_editor_active_move_completion(
350 cx: &mut EditorTestContext,
351 assert: impl FnOnce(MultiBufferSnapshot, Anchor),
352) {
353 cx.editor(|editor, _, cx| {
354 let completion_state = editor
355 .active_edit_prediction
356 .as_ref()
357 .expect("editor has no active completion");
358
359 if let EditPrediction::MoveWithin { target, .. } = &completion_state.completion {
360 assert(editor.buffer().read(cx).snapshot(cx), *target);
361 } else {
362 panic!("expected move completion");
363 }
364 })
365}
366
367fn accept_completion(cx: &mut EditorTestContext) {
368 cx.update_editor(|editor, window, cx| {
369 editor.accept_edit_prediction(&crate::AcceptEditPrediction, window, cx)
370 })
371}
372
373fn propose_edits<T: ToOffset>(
374 provider: &Entity<FakeEditPredictionDelegate>,
375 edits: Vec<(Range<T>, &str)>,
376 cx: &mut EditorTestContext,
377) {
378 let snapshot = cx.buffer_snapshot();
379 let edits = edits.into_iter().map(|(range, text)| {
380 let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
381 (range, text.into())
382 });
383
384 cx.update(|_, cx| {
385 provider.update(cx, |provider, _| {
386 provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
387 id: None,
388 edits: edits.collect(),
389 edit_preview: None,
390 }))
391 })
392 });
393}
394
395fn assign_editor_completion_provider(
396 provider: Entity<FakeEditPredictionDelegate>,
397 cx: &mut EditorTestContext,
398) {
399 cx.update_editor(|editor, window, cx| {
400 editor.set_edit_prediction_provider(Some(provider), window, cx);
401 })
402}
403
404fn propose_edits_non_zed<T: ToOffset>(
405 provider: &Entity<FakeNonZedEditPredictionDelegate>,
406 edits: Vec<(Range<T>, &str)>,
407 cx: &mut EditorTestContext,
408) {
409 let snapshot = cx.buffer_snapshot();
410 let edits = edits.into_iter().map(|(range, text)| {
411 let range = snapshot.anchor_after(range.start)..snapshot.anchor_before(range.end);
412 (range, text.into())
413 });
414
415 cx.update(|_, cx| {
416 provider.update(cx, |provider, _| {
417 provider.set_edit_prediction(Some(edit_prediction_types::EditPrediction::Local {
418 id: None,
419 edits: edits.collect(),
420 edit_preview: None,
421 }))
422 })
423 });
424}
425
426fn assign_editor_completion_provider_non_zed(
427 provider: Entity<FakeNonZedEditPredictionDelegate>,
428 cx: &mut EditorTestContext,
429) {
430 cx.update_editor(|editor, window, cx| {
431 editor.set_edit_prediction_provider(Some(provider), window, cx);
432 })
433}
434
435#[derive(Default, Clone)]
436pub struct FakeEditPredictionDelegate {
437 pub completion: Option<edit_prediction_types::EditPrediction>,
438}
439
440impl FakeEditPredictionDelegate {
441 pub fn set_edit_prediction(
442 &mut self,
443 completion: Option<edit_prediction_types::EditPrediction>,
444 ) {
445 self.completion = completion;
446 }
447}
448
449impl EditPredictionDelegate for FakeEditPredictionDelegate {
450 fn name() -> &'static str {
451 "fake-completion-provider"
452 }
453
454 fn display_name() -> &'static str {
455 "Fake Completion Provider"
456 }
457
458 fn show_predictions_in_menu() -> bool {
459 true
460 }
461
462 fn supports_jump_to_edit() -> bool {
463 true
464 }
465
466 fn is_enabled(
467 &self,
468 _buffer: &gpui::Entity<language::Buffer>,
469 _cursor_position: language::Anchor,
470 _cx: &gpui::App,
471 ) -> bool {
472 true
473 }
474
475 fn is_refreshing(&self, _cx: &gpui::App) -> bool {
476 false
477 }
478
479 fn refresh(
480 &mut self,
481 _buffer: gpui::Entity<language::Buffer>,
482 _cursor_position: language::Anchor,
483 _debounce: bool,
484 _cx: &mut gpui::Context<Self>,
485 ) {
486 }
487
488 fn cycle(
489 &mut self,
490 _buffer: gpui::Entity<language::Buffer>,
491 _cursor_position: language::Anchor,
492 _direction: edit_prediction_types::Direction,
493 _cx: &mut gpui::Context<Self>,
494 ) {
495 }
496
497 fn accept(&mut self, _cx: &mut gpui::Context<Self>) {}
498
499 fn discard(&mut self, _cx: &mut gpui::Context<Self>) {}
500
501 fn suggest<'a>(
502 &mut self,
503 _buffer: &gpui::Entity<language::Buffer>,
504 _cursor_position: language::Anchor,
505 _cx: &mut gpui::Context<Self>,
506 ) -> Option<edit_prediction_types::EditPrediction> {
507 self.completion.clone()
508 }
509}
510
511#[derive(Default, Clone)]
512pub struct FakeNonZedEditPredictionDelegate {
513 pub completion: Option<edit_prediction_types::EditPrediction>,
514}
515
516impl FakeNonZedEditPredictionDelegate {
517 pub fn set_edit_prediction(
518 &mut self,
519 completion: Option<edit_prediction_types::EditPrediction>,
520 ) {
521 self.completion = completion;
522 }
523}
524
525impl EditPredictionDelegate for FakeNonZedEditPredictionDelegate {
526 fn name() -> &'static str {
527 "fake-non-zed-provider"
528 }
529
530 fn display_name() -> &'static str {
531 "Fake Non-Zed Provider"
532 }
533
534 fn show_predictions_in_menu() -> bool {
535 false
536 }
537
538 fn supports_jump_to_edit() -> bool {
539 false
540 }
541
542 fn is_enabled(
543 &self,
544 _buffer: &gpui::Entity<language::Buffer>,
545 _cursor_position: language::Anchor,
546 _cx: &gpui::App,
547 ) -> bool {
548 true
549 }
550
551 fn is_refreshing(&self, _cx: &gpui::App) -> bool {
552 false
553 }
554
555 fn refresh(
556 &mut self,
557 _buffer: gpui::Entity<language::Buffer>,
558 _cursor_position: language::Anchor,
559 _debounce: bool,
560 _cx: &mut gpui::Context<Self>,
561 ) {
562 }
563
564 fn cycle(
565 &mut self,
566 _buffer: gpui::Entity<language::Buffer>,
567 _cursor_position: language::Anchor,
568 _direction: edit_prediction_types::Direction,
569 _cx: &mut gpui::Context<Self>,
570 ) {
571 }
572
573 fn accept(&mut self, _cx: &mut gpui::Context<Self>) {}
574
575 fn discard(&mut self, _cx: &mut gpui::Context<Self>) {}
576
577 fn suggest<'a>(
578 &mut self,
579 _buffer: &gpui::Entity<language::Buffer>,
580 _cursor_position: language::Anchor,
581 _cx: &mut gpui::Context<Self>,
582 ) -> Option<edit_prediction_types::EditPrediction> {
583 self.completion.clone()
584 }
585}