patch.rs

  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(crate) 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(crate) enum AssistantPatchStatus {
 21    Pending,
 22    Ready,
 23}
 24
 25#[derive(Clone, Debug, PartialEq, Eq)]
 26pub(crate) 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(crate) 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(crate) 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(crate) 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
584585
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}