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