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