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