1use anyhow::{anyhow, Context as _, Result};
2use collections::HashMap;
3use editor::ProposedChangesEditor;
4use futures::{future, TryFutureExt as _};
5use gpui::{AppContext, AsyncAppContext, Model, SharedString};
6use language::{AutoindentMode, Buffer, BufferSnapshot};
7use project::{Project, ProjectPath};
8use std::{cmp, ops::Range, path::Path, sync::Arc};
9use text::{AnchorRangeExt as _, Bias, OffsetRangeExt as _, Point};
10
11#[derive(Clone, Debug)]
12pub struct AssistantPatch {
13 pub range: Range<language::Anchor>,
14 pub title: SharedString,
15 pub edits: Arc<[Result<AssistantEdit>]>,
16 pub status: AssistantPatchStatus,
17}
18
19#[derive(Copy, Clone, Debug, PartialEq, Eq)]
20pub enum AssistantPatchStatus {
21 Pending,
22 Ready,
23}
24
25#[derive(Clone, Debug, PartialEq, Eq)]
26pub struct AssistantEdit {
27 pub path: String,
28 pub kind: AssistantEditKind,
29}
30
31#[derive(Clone, Debug, PartialEq, Eq)]
32pub enum AssistantEditKind {
33 Update {
34 old_text: String,
35 new_text: String,
36 description: Option<String>,
37 },
38 Create {
39 new_text: String,
40 description: Option<String>,
41 },
42 InsertBefore {
43 old_text: String,
44 new_text: String,
45 description: Option<String>,
46 },
47 InsertAfter {
48 old_text: String,
49 new_text: String,
50 description: Option<String>,
51 },
52 Delete {
53 old_text: String,
54 },
55}
56
57#[derive(Clone, Debug, Eq, PartialEq)]
58pub struct ResolvedPatch {
59 pub edit_groups: HashMap<Model<Buffer>, Vec<ResolvedEditGroup>>,
60 pub errors: Vec<AssistantPatchResolutionError>,
61}
62
63#[derive(Clone, Debug, Eq, PartialEq)]
64pub struct ResolvedEditGroup {
65 pub context_range: Range<language::Anchor>,
66 pub edits: Vec<ResolvedEdit>,
67}
68
69#[derive(Clone, Debug, Eq, PartialEq)]
70pub struct ResolvedEdit {
71 range: Range<language::Anchor>,
72 new_text: String,
73 description: Option<String>,
74}
75
76#[derive(Clone, Debug, Eq, PartialEq)]
77pub struct AssistantPatchResolutionError {
78 pub edit_ix: usize,
79 pub message: String,
80}
81
82#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord)]
83enum SearchDirection {
84 Up,
85 Left,
86 Diagonal,
87}
88
89#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
90struct SearchState {
91 cost: u32,
92 direction: SearchDirection,
93}
94
95impl SearchState {
96 fn new(cost: u32, direction: SearchDirection) -> Self {
97 Self { cost, direction }
98 }
99}
100
101struct SearchMatrix {
102 cols: usize,
103 data: Vec<SearchState>,
104}
105
106impl SearchMatrix {
107 fn new(rows: usize, cols: usize) -> Self {
108 SearchMatrix {
109 cols,
110 data: vec![SearchState::new(0, SearchDirection::Diagonal); rows * cols],
111 }
112 }
113
114 fn get(&self, row: usize, col: usize) -> SearchState {
115 self.data[row * self.cols + col]
116 }
117
118 fn set(&mut self, row: usize, col: usize, cost: SearchState) {
119 self.data[row * self.cols + col] = cost;
120 }
121}
122
123impl ResolvedPatch {
124 pub fn apply(&self, editor: &ProposedChangesEditor, cx: &mut AppContext) {
125 for (buffer, groups) in &self.edit_groups {
126 let branch = editor.branch_buffer_for_base(buffer).unwrap();
127 Self::apply_edit_groups(groups, &branch, cx);
128 }
129 editor.recalculate_all_buffer_diffs();
130 }
131
132 fn apply_edit_groups(
133 groups: &Vec<ResolvedEditGroup>,
134 buffer: &Model<Buffer>,
135 cx: &mut AppContext,
136 ) {
137 let mut edits = Vec::new();
138 for group in groups {
139 for suggestion in &group.edits {
140 edits.push((suggestion.range.clone(), suggestion.new_text.clone()));
141 }
142 }
143 buffer.update(cx, |buffer, cx| {
144 buffer.edit(
145 edits,
146 Some(AutoindentMode::Block {
147 original_indent_columns: Vec::new(),
148 }),
149 cx,
150 );
151 });
152 }
153}
154
155impl ResolvedEdit {
156 pub fn try_merge(&mut self, other: &Self, buffer: &text::BufferSnapshot) -> bool {
157 let range = &self.range;
158 let other_range = &other.range;
159
160 // Don't merge if we don't contain the other suggestion.
161 if range.start.cmp(&other_range.start, buffer).is_gt()
162 || range.end.cmp(&other_range.end, buffer).is_lt()
163 {
164 return false;
165 }
166
167 let other_offset_range = other_range.to_offset(buffer);
168 let offset_range = range.to_offset(buffer);
169
170 // If the other range is empty at the start of this edit's range, combine the new text
171 if other_offset_range.is_empty() && other_offset_range.start == offset_range.start {
172 self.new_text = format!("{}\n{}", other.new_text, self.new_text);
173 self.range.start = other_range.start;
174
175 if let Some((description, other_description)) =
176 self.description.as_mut().zip(other.description.as_ref())
177 {
178 *description = format!("{}\n{}", other_description, description)
179 }
180 } else {
181 if let Some((description, other_description)) =
182 self.description.as_mut().zip(other.description.as_ref())
183 {
184 description.push('\n');
185 description.push_str(other_description);
186 }
187 }
188
189 true
190 }
191}
192
193impl AssistantEdit {
194 pub fn new(
195 path: Option<String>,
196 operation: Option<String>,
197 old_text: Option<String>,
198 new_text: Option<String>,
199 description: Option<String>,
200 ) -> Result<Self> {
201 let path = path.ok_or_else(|| anyhow!("missing path"))?;
202 let operation = operation.ok_or_else(|| anyhow!("missing operation"))?;
203
204 let kind = match operation.as_str() {
205 "update" => AssistantEditKind::Update {
206 old_text: old_text.ok_or_else(|| anyhow!("missing old_text"))?,
207 new_text: new_text.ok_or_else(|| anyhow!("missing new_text"))?,
208 description,
209 },
210 "insert_before" => AssistantEditKind::InsertBefore {
211 old_text: old_text.ok_or_else(|| anyhow!("missing old_text"))?,
212 new_text: new_text.ok_or_else(|| anyhow!("missing new_text"))?,
213 description,
214 },
215 "insert_after" => AssistantEditKind::InsertAfter {
216 old_text: old_text.ok_or_else(|| anyhow!("missing old_text"))?,
217 new_text: new_text.ok_or_else(|| anyhow!("missing new_text"))?,
218 description,
219 },
220 "delete" => AssistantEditKind::Delete {
221 old_text: old_text.ok_or_else(|| anyhow!("missing old_text"))?,
222 },
223 "create" => AssistantEditKind::Create {
224 description,
225 new_text: new_text.ok_or_else(|| anyhow!("missing new_text"))?,
226 },
227 _ => Err(anyhow!("unknown operation {operation:?}"))?,
228 };
229
230 Ok(Self { path, kind })
231 }
232
233 pub async fn resolve(
234 &self,
235 project: Model<Project>,
236 mut cx: AsyncAppContext,
237 ) -> Result<(Model<Buffer>, ResolvedEdit)> {
238 let path = self.path.clone();
239 let kind = self.kind.clone();
240 let buffer = project
241 .update(&mut cx, |project, cx| {
242 let project_path = project
243 .find_project_path(Path::new(&path), cx)
244 .or_else(|| {
245 // If we couldn't find a project path for it, put it in the active worktree
246 // so that when we create the buffer, it can be saved.
247 let worktree = project
248 .active_entry()
249 .and_then(|entry_id| project.worktree_for_entry(entry_id, cx))
250 .or_else(|| project.worktrees(cx).next())?;
251 let worktree = worktree.read(cx);
252
253 Some(ProjectPath {
254 worktree_id: worktree.id(),
255 path: Arc::from(Path::new(&path)),
256 })
257 })
258 .with_context(|| format!("worktree not found for {:?}", path))?;
259 anyhow::Ok(project.open_buffer(project_path, cx))
260 })??
261 .await?;
262
263 let snapshot = buffer.update(&mut cx, |buffer, _| buffer.snapshot())?;
264 let suggestion = cx
265 .background_executor()
266 .spawn(async move { kind.resolve(&snapshot) })
267 .await;
268
269 Ok((buffer, suggestion))
270 }
271}
272
273impl AssistantEditKind {
274 fn resolve(self, snapshot: &BufferSnapshot) -> ResolvedEdit {
275 match self {
276 Self::Update {
277 old_text,
278 new_text,
279 description,
280 } => {
281 let range = Self::resolve_location(&snapshot, &old_text);
282 ResolvedEdit {
283 range,
284 new_text,
285 description,
286 }
287 }
288 Self::Create {
289 new_text,
290 description,
291 } => ResolvedEdit {
292 range: text::Anchor::MIN..text::Anchor::MAX,
293 description,
294 new_text,
295 },
296 Self::InsertBefore {
297 old_text,
298 mut new_text,
299 description,
300 } => {
301 let range = Self::resolve_location(&snapshot, &old_text);
302 new_text.push('\n');
303 ResolvedEdit {
304 range: range.start..range.start,
305 new_text,
306 description,
307 }
308 }
309 Self::InsertAfter {
310 old_text,
311 mut new_text,
312 description,
313 } => {
314 let range = Self::resolve_location(&snapshot, &old_text);
315 new_text.insert(0, '\n');
316 ResolvedEdit {
317 range: range.end..range.end,
318 new_text,
319 description,
320 }
321 }
322 Self::Delete { old_text } => {
323 let range = Self::resolve_location(&snapshot, &old_text);
324 ResolvedEdit {
325 range,
326 new_text: String::new(),
327 description: None,
328 }
329 }
330 }
331 }
332
333 fn resolve_location(buffer: &text::BufferSnapshot, search_query: &str) -> Range<text::Anchor> {
334 const INSERTION_COST: u32 = 3;
335 const DELETION_COST: u32 = 10;
336 const WHITESPACE_INSERTION_COST: u32 = 1;
337 const WHITESPACE_DELETION_COST: u32 = 1;
338
339 let buffer_len = buffer.len();
340 let query_len = search_query.len();
341 let mut matrix = SearchMatrix::new(query_len + 1, buffer_len + 1);
342 let mut leading_deletion_cost = 0_u32;
343 for (row, query_byte) in search_query.bytes().enumerate() {
344 let deletion_cost = if query_byte.is_ascii_whitespace() {
345 WHITESPACE_DELETION_COST
346 } else {
347 DELETION_COST
348 };
349
350 leading_deletion_cost = leading_deletion_cost.saturating_add(deletion_cost);
351 matrix.set(
352 row + 1,
353 0,
354 SearchState::new(leading_deletion_cost, SearchDirection::Diagonal),
355 );
356
357 for (col, buffer_byte) in buffer.bytes_in_range(0..buffer.len()).flatten().enumerate() {
358 let insertion_cost = if buffer_byte.is_ascii_whitespace() {
359 WHITESPACE_INSERTION_COST
360 } else {
361 INSERTION_COST
362 };
363
364 let up = SearchState::new(
365 matrix.get(row, col + 1).cost.saturating_add(deletion_cost),
366 SearchDirection::Up,
367 );
368 let left = SearchState::new(
369 matrix.get(row + 1, col).cost.saturating_add(insertion_cost),
370 SearchDirection::Left,
371 );
372 let diagonal = SearchState::new(
373 if query_byte == *buffer_byte {
374 matrix.get(row, col).cost
375 } else {
376 matrix
377 .get(row, col)
378 .cost
379 .saturating_add(deletion_cost + insertion_cost)
380 },
381 SearchDirection::Diagonal,
382 );
383 matrix.set(row + 1, col + 1, up.min(left).min(diagonal));
384 }
385 }
386
387 // Traceback to find the best match
388 let mut best_buffer_end = buffer_len;
389 let mut best_cost = u32::MAX;
390 for col in 1..=buffer_len {
391 let cost = matrix.get(query_len, col).cost;
392 if cost < best_cost {
393 best_cost = cost;
394 best_buffer_end = col;
395 }
396 }
397
398 let mut query_ix = query_len;
399 let mut buffer_ix = best_buffer_end;
400 while query_ix > 0 && buffer_ix > 0 {
401 let current = matrix.get(query_ix, buffer_ix);
402 match current.direction {
403 SearchDirection::Diagonal => {
404 query_ix -= 1;
405 buffer_ix -= 1;
406 }
407 SearchDirection::Up => {
408 query_ix -= 1;
409 }
410 SearchDirection::Left => {
411 buffer_ix -= 1;
412 }
413 }
414 }
415
416 let mut start = buffer.offset_to_point(buffer.clip_offset(buffer_ix, Bias::Left));
417 start.column = 0;
418 let mut end = buffer.offset_to_point(buffer.clip_offset(best_buffer_end, Bias::Right));
419 if end.column > 0 {
420 end.column = buffer.line_len(end.row);
421 }
422
423 buffer.anchor_after(start)..buffer.anchor_before(end)
424 }
425}
426
427impl AssistantPatch {
428 pub async fn resolve(
429 &self,
430 project: Model<Project>,
431 cx: &mut AsyncAppContext,
432 ) -> ResolvedPatch {
433 let mut resolve_tasks = Vec::new();
434 for (ix, edit) in self.edits.iter().enumerate() {
435 if let Ok(edit) = edit.as_ref() {
436 resolve_tasks.push(
437 edit.resolve(project.clone(), cx.clone())
438 .map_err(move |error| (ix, error)),
439 );
440 }
441 }
442
443 let edits = future::join_all(resolve_tasks).await;
444 let mut errors = Vec::new();
445 let mut edits_by_buffer = HashMap::default();
446 for entry in edits {
447 match entry {
448 Ok((buffer, edit)) => {
449 edits_by_buffer
450 .entry(buffer)
451 .or_insert_with(Vec::new)
452 .push(edit);
453 }
454 Err((edit_ix, error)) => errors.push(AssistantPatchResolutionError {
455 edit_ix,
456 message: error.to_string(),
457 }),
458 }
459 }
460
461 // Expand the context ranges of each edit and group edits with overlapping context ranges.
462 let mut edit_groups_by_buffer = HashMap::default();
463 for (buffer, edits) in edits_by_buffer {
464 if let Ok(snapshot) = buffer.update(cx, |buffer, _| buffer.text_snapshot()) {
465 edit_groups_by_buffer.insert(buffer, Self::group_edits(edits, &snapshot));
466 }
467 }
468
469 ResolvedPatch {
470 edit_groups: edit_groups_by_buffer,
471 errors,
472 }
473 }
474
475 fn group_edits(
476 mut edits: Vec<ResolvedEdit>,
477 snapshot: &text::BufferSnapshot,
478 ) -> Vec<ResolvedEditGroup> {
479 let mut edit_groups = Vec::<ResolvedEditGroup>::new();
480 // Sort edits by their range so that earlier, larger ranges come first
481 edits.sort_by(|a, b| a.range.cmp(&b.range, &snapshot));
482
483 // Merge overlapping edits
484 edits.dedup_by(|a, b| b.try_merge(a, &snapshot));
485
486 // Create context ranges for each edit
487 for edit in edits {
488 let context_range = {
489 let edit_point_range = edit.range.to_point(&snapshot);
490 let start_row = edit_point_range.start.row.saturating_sub(5);
491 let end_row = cmp::min(edit_point_range.end.row + 5, snapshot.max_point().row);
492 let start = snapshot.anchor_before(Point::new(start_row, 0));
493 let end = snapshot.anchor_after(Point::new(end_row, snapshot.line_len(end_row)));
494 start..end
495 };
496
497 if let Some(last_group) = edit_groups.last_mut() {
498 if last_group
499 .context_range
500 .end
501 .cmp(&context_range.start, &snapshot)
502 .is_ge()
503 {
504 // Merge with the previous group if context ranges overlap
505 last_group.context_range.end = context_range.end;
506 last_group.edits.push(edit);
507 } else {
508 // Create a new group
509 edit_groups.push(ResolvedEditGroup {
510 context_range,
511 edits: vec![edit],
512 });
513 }
514 } else {
515 // Create the first group
516 edit_groups.push(ResolvedEditGroup {
517 context_range,
518 edits: vec![edit],
519 });
520 }
521 }
522
523 edit_groups
524 }
525
526 pub fn path_count(&self) -> usize {
527 self.paths().count()
528 }
529
530 pub fn paths(&self) -> impl '_ + Iterator<Item = &str> {
531 let mut prev_path = None;
532 self.edits.iter().filter_map(move |edit| {
533 if let Ok(edit) = edit {
534 let path = Some(edit.path.as_str());
535 if path != prev_path {
536 prev_path = path;
537 return path;
538 }
539 }
540 None
541 })
542 }
543}
544
545impl PartialEq for AssistantPatch {
546 fn eq(&self, other: &Self) -> bool {
547 self.range == other.range
548 && self.title == other.title
549 && Arc::ptr_eq(&self.edits, &other.edits)
550 }
551}
552
553impl Eq for AssistantPatch {}
554
555#[cfg(test)]
556mod tests {
557 use super::*;
558 use gpui::{AppContext, Context};
559 use language::{
560 language_settings::AllLanguageSettings, Language, LanguageConfig, LanguageMatcher,
561 };
562 use settings::SettingsStore;
563 use ui::BorrowAppContext;
564 use unindent::Unindent as _;
565 use util::test::{generate_marked_text, marked_text_ranges};
566
567 #[gpui::test]
568 fn test_resolve_location(cx: &mut AppContext) {
569 assert_location_resolution(
570 concat!(
571 " Lorem\n",
572 "« ipsum\n",
573 " dolor sit amet»\n",
574 " consecteur",
575 ),
576 "ipsum\ndolor",
577 cx,
578 );
579
580 assert_location_resolution(
581 &"
582 «fn foo1(a: usize) -> usize {
583 40
584 }»
585
586 fn foo2(b: usize) -> usize {
587 42
588 }
589 "
590 .unindent(),
591 "fn foo1(b: usize) {\n40\n}",
592 cx,
593 );
594
595 assert_location_resolution(
596 &"
597 fn main() {
598 « Foo
599 .bar()
600 .baz()
601 .qux()»
602 }
603
604 fn foo2(b: usize) -> usize {
605 42
606 }
607 "
608 .unindent(),
609 "Foo.bar.baz.qux()",
610 cx,
611 );
612
613 assert_location_resolution(
614 &"
615 class Something {
616 one() { return 1; }
617 « two() { return 2222; }
618 three() { return 333; }
619 four() { return 4444; }
620 five() { return 5555; }
621 six() { return 6666; }
622 » seven() { return 7; }
623 eight() { return 8; }
624 }
625 "
626 .unindent(),
627 &"
628 two() { return 2222; }
629 four() { return 4444; }
630 five() { return 5555; }
631 six() { return 6666; }
632 "
633 .unindent(),
634 cx,
635 );
636 }
637
638 #[gpui::test]
639 fn test_resolve_edits(cx: &mut AppContext) {
640 init_test(cx);
641
642 assert_edits(
643 "
644 /// A person
645 struct Person {
646 name: String,
647 age: usize,
648 }
649
650 /// A dog
651 struct Dog {
652 weight: f32,
653 }
654
655 impl Person {
656 fn name(&self) -> &str {
657 &self.name
658 }
659 }
660 "
661 .unindent(),
662 vec![
663 AssistantEditKind::Update {
664 old_text: "
665 name: String,
666 "
667 .unindent(),
668 new_text: "
669 first_name: String,
670 last_name: String,
671 "
672 .unindent(),
673 description: None,
674 },
675 AssistantEditKind::Update {
676 old_text: "
677 fn name(&self) -> &str {
678 &self.name
679 }
680 "
681 .unindent(),
682 new_text: "
683 fn name(&self) -> String {
684 format!(\"{} {}\", self.first_name, self.last_name)
685 }
686 "
687 .unindent(),
688 description: None,
689 },
690 ],
691 "
692 /// A person
693 struct Person {
694 first_name: String,
695 last_name: String,
696 age: usize,
697 }
698
699 /// A dog
700 struct Dog {
701 weight: f32,
702 }
703
704 impl Person {
705 fn name(&self) -> String {
706 format!(\"{} {}\", self.first_name, self.last_name)
707 }
708 }
709 "
710 .unindent(),
711 cx,
712 );
713
714 // Ensure InsertBefore merges correctly with Update of the same text
715 assert_edits(
716 "
717 fn foo() {
718
719 }
720 "
721 .unindent(),
722 vec![
723 AssistantEditKind::InsertBefore {
724 old_text: "
725 fn foo() {"
726 .unindent(),
727 new_text: "
728 fn bar() {
729 qux();
730 }"
731 .unindent(),
732 description: Some("implement bar".into()),
733 },
734 AssistantEditKind::Update {
735 old_text: "
736 fn foo() {
737
738 }"
739 .unindent(),
740 new_text: "
741 fn foo() {
742 bar();
743 }"
744 .unindent(),
745 description: Some("call bar in foo".into()),
746 },
747 AssistantEditKind::InsertAfter {
748 old_text: "
749 fn foo() {
750
751 }
752 "
753 .unindent(),
754 new_text: "
755 fn qux() {
756 // todo
757 }
758 "
759 .unindent(),
760 description: Some("implement qux".into()),
761 },
762 ],
763 "
764 fn bar() {
765 qux();
766 }
767
768 fn foo() {
769 bar();
770 }
771
772 fn qux() {
773 // todo
774 }
775 "
776 .unindent(),
777 cx,
778 );
779
780 // Correctly indent new text when replacing multiple adjacent indented blocks.
781 assert_edits(
782 "
783 impl Numbers {
784 fn one() {
785 1
786 }
787
788 fn two() {
789 2
790 }
791
792 fn three() {
793 3
794 }
795 }
796 "
797 .unindent(),
798 vec![
799 AssistantEditKind::Update {
800 old_text: "
801 fn one() {
802 1
803 }
804 "
805 .unindent(),
806 new_text: "
807 fn one() {
808 101
809 }
810 "
811 .unindent(),
812 description: None,
813 },
814 AssistantEditKind::Update {
815 old_text: "
816 fn two() {
817 2
818 }
819 "
820 .unindent(),
821 new_text: "
822 fn two() {
823 102
824 }
825 "
826 .unindent(),
827 description: None,
828 },
829 AssistantEditKind::Update {
830 old_text: "
831 fn three() {
832 3
833 }
834 "
835 .unindent(),
836 new_text: "
837 fn three() {
838 103
839 }
840 "
841 .unindent(),
842 description: None,
843 },
844 ],
845 "
846 impl Numbers {
847 fn one() {
848 101
849 }
850
851 fn two() {
852 102
853 }
854
855 fn three() {
856 103
857 }
858 }
859 "
860 .unindent(),
861 cx,
862 );
863
864 assert_edits(
865 "
866 impl Person {
867 fn set_name(&mut self, name: String) {
868 self.name = name;
869 }
870
871 fn name(&self) -> String {
872 return self.name;
873 }
874 }
875 "
876 .unindent(),
877 vec![
878 AssistantEditKind::Update {
879 old_text: "self.name = name;".unindent(),
880 new_text: "self._name = name;".unindent(),
881 description: None,
882 },
883 AssistantEditKind::Update {
884 old_text: "return self.name;\n".unindent(),
885 new_text: "return self._name;\n".unindent(),
886 description: None,
887 },
888 ],
889 "
890 impl Person {
891 fn set_name(&mut self, name: String) {
892 self._name = name;
893 }
894
895 fn name(&self) -> String {
896 return self._name;
897 }
898 }
899 "
900 .unindent(),
901 cx,
902 );
903 }
904
905 fn init_test(cx: &mut AppContext) {
906 let settings_store = SettingsStore::test(cx);
907 cx.set_global(settings_store);
908 language::init(cx);
909 cx.update_global::<SettingsStore, _>(|settings, cx| {
910 settings.update_user_settings::<AllLanguageSettings>(cx, |_| {});
911 });
912 }
913
914 #[track_caller]
915 fn assert_location_resolution(
916 text_with_expected_range: &str,
917 query: &str,
918 cx: &mut AppContext,
919 ) {
920 let (text, _) = marked_text_ranges(text_with_expected_range, false);
921 let buffer = cx.new_model(|cx| Buffer::local(text.clone(), cx));
922 let snapshot = buffer.read(cx).snapshot();
923 let range = AssistantEditKind::resolve_location(&snapshot, query).to_offset(&snapshot);
924 let text_with_actual_range = generate_marked_text(&text, &[range], false);
925 pretty_assertions::assert_eq!(text_with_actual_range, text_with_expected_range);
926 }
927
928 #[track_caller]
929 fn assert_edits(
930 old_text: String,
931 edits: Vec<AssistantEditKind>,
932 new_text: String,
933 cx: &mut AppContext,
934 ) {
935 let buffer =
936 cx.new_model(|cx| Buffer::local(old_text, cx).with_language(Arc::new(rust_lang()), cx));
937 let snapshot = buffer.read(cx).snapshot();
938 let resolved_edits = edits
939 .into_iter()
940 .map(|kind| kind.resolve(&snapshot))
941 .collect();
942 let edit_groups = AssistantPatch::group_edits(resolved_edits, &snapshot);
943 ResolvedPatch::apply_edit_groups(&edit_groups, &buffer, cx);
944 let actual_new_text = buffer.read(cx).text();
945 pretty_assertions::assert_eq!(actual_new_text, new_text);
946 }
947
948 fn rust_lang() -> Language {
949 Language::new(
950 LanguageConfig {
951 name: "Rust".into(),
952 matcher: LanguageMatcher {
953 path_suffixes: vec!["rs".to_string()],
954 ..Default::default()
955 },
956 ..Default::default()
957 },
958 Some(language::tree_sitter_rust::LANGUAGE.into()),
959 )
960 .with_indents_query(
961 r#"
962 (call_expression) @indent
963 (field_expression) @indent
964 (_ "(" ")" @end) @indent
965 (_ "{" "}" @end) @indent
966 "#,
967 )
968 .unwrap()
969 }
970}