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: String,
 37    },
 38    Create {
 39        new_text: String,
 40        description: String,
 41    },
 42    InsertBefore {
 43        old_text: String,
 44        new_text: String,
 45        description: String,
 46    },
 47    InsertAfter {
 48        old_text: String,
 49        new_text: String,
 50        description: 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// A measure of the currently quality of an in-progress fuzzy search.
 90//
 91// Uses 60 bits to store a numeric cost, and 4 bits to store the preceding
 92// operation in the search.
 93#[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
 94struct SearchState {
 95    score: u32,
 96    direction: SearchDirection,
 97}
 98
 99impl SearchState {
100    fn new(score: u32, direction: SearchDirection) -> Self {
101        Self { score, direction }
102    }
103}
104
105impl ResolvedPatch {
106    pub fn apply(&self, editor: &ProposedChangesEditor, cx: &mut AppContext) {
107        for (buffer, groups) in &self.edit_groups {
108            let branch = editor.branch_buffer_for_base(buffer).unwrap();
109            Self::apply_edit_groups(groups, &branch, cx);
110        }
111        editor.recalculate_all_buffer_diffs();
112    }
113
114    fn apply_edit_groups(
115        groups: &Vec<ResolvedEditGroup>,
116        buffer: &Model<Buffer>,
117        cx: &mut AppContext,
118    ) {
119        let mut edits = Vec::new();
120        for group in groups {
121            for suggestion in &group.edits {
122                edits.push((suggestion.range.clone(), suggestion.new_text.clone()));
123            }
124        }
125        buffer.update(cx, |buffer, cx| {
126            buffer.edit(
127                edits,
128                Some(AutoindentMode::Block {
129                    original_indent_columns: Vec::new(),
130                }),
131                cx,
132            );
133        });
134    }
135}
136
137impl ResolvedEdit {
138    pub fn try_merge(&mut self, other: &Self, buffer: &text::BufferSnapshot) -> bool {
139        let range = &self.range;
140        let other_range = &other.range;
141
142        // Don't merge if we don't contain the other suggestion.
143        if range.start.cmp(&other_range.start, buffer).is_gt()
144            || range.end.cmp(&other_range.end, buffer).is_lt()
145        {
146            return false;
147        }
148
149        let other_offset_range = other_range.to_offset(buffer);
150        let offset_range = range.to_offset(buffer);
151
152        // If the other range is empty at the start of this edit's range, combine the new text
153        if other_offset_range.is_empty() && other_offset_range.start == offset_range.start {
154            self.new_text = format!("{}\n{}", other.new_text, self.new_text);
155            self.range.start = other_range.start;
156
157            if let Some((description, other_description)) =
158                self.description.as_mut().zip(other.description.as_ref())
159            {
160                *description = format!("{}\n{}", other_description, description)
161            }
162        } else {
163            if let Some((description, other_description)) =
164                self.description.as_mut().zip(other.description.as_ref())
165            {
166                description.push('\n');
167                description.push_str(other_description);
168            }
169        }
170
171        true
172    }
173}
174
175impl AssistantEdit {
176    pub fn new(
177        path: Option<String>,
178        operation: Option<String>,
179        old_text: Option<String>,
180        new_text: Option<String>,
181        description: Option<String>,
182    ) -> Result<Self> {
183        let path = path.ok_or_else(|| anyhow!("missing path"))?;
184        let operation = operation.ok_or_else(|| anyhow!("missing operation"))?;
185
186        let kind = match operation.as_str() {
187            "update" => AssistantEditKind::Update {
188                old_text: old_text.ok_or_else(|| anyhow!("missing old_text"))?,
189                new_text: new_text.ok_or_else(|| anyhow!("missing new_text"))?,
190                description: description.ok_or_else(|| anyhow!("missing description"))?,
191            },
192            "insert_before" => AssistantEditKind::InsertBefore {
193                old_text: old_text.ok_or_else(|| anyhow!("missing old_text"))?,
194                new_text: new_text.ok_or_else(|| anyhow!("missing new_text"))?,
195                description: description.ok_or_else(|| anyhow!("missing description"))?,
196            },
197            "insert_after" => AssistantEditKind::InsertAfter {
198                old_text: old_text.ok_or_else(|| anyhow!("missing old_text"))?,
199                new_text: new_text.ok_or_else(|| anyhow!("missing new_text"))?,
200                description: description.ok_or_else(|| anyhow!("missing description"))?,
201            },
202            "delete" => AssistantEditKind::Delete {
203                old_text: old_text.ok_or_else(|| anyhow!("missing old_text"))?,
204            },
205            "create" => AssistantEditKind::Create {
206                description: description.ok_or_else(|| anyhow!("missing description"))?,
207                new_text: new_text.ok_or_else(|| anyhow!("missing new_text"))?,
208            },
209            _ => Err(anyhow!("unknown operation {operation:?}"))?,
210        };
211
212        Ok(Self { path, kind })
213    }
214
215    pub async fn resolve(
216        &self,
217        project: Model<Project>,
218        mut cx: AsyncAppContext,
219    ) -> Result<(Model<Buffer>, ResolvedEdit)> {
220        let path = self.path.clone();
221        let kind = self.kind.clone();
222        let buffer = project
223            .update(&mut cx, |project, cx| {
224                let project_path = project
225                    .find_project_path(Path::new(&path), cx)
226                    .or_else(|| {
227                        // If we couldn't find a project path for it, put it in the active worktree
228                        // so that when we create the buffer, it can be saved.
229                        let worktree = project
230                            .active_entry()
231                            .and_then(|entry_id| project.worktree_for_entry(entry_id, cx))
232                            .or_else(|| project.worktrees(cx).next())?;
233                        let worktree = worktree.read(cx);
234
235                        Some(ProjectPath {
236                            worktree_id: worktree.id(),
237                            path: Arc::from(Path::new(&path)),
238                        })
239                    })
240                    .with_context(|| format!("worktree not found for {:?}", path))?;
241                anyhow::Ok(project.open_buffer(project_path, cx))
242            })??
243            .await?;
244
245        let snapshot = buffer.update(&mut cx, |buffer, _| buffer.snapshot())?;
246        let suggestion = cx
247            .background_executor()
248            .spawn(async move { kind.resolve(&snapshot) })
249            .await;
250
251        Ok((buffer, suggestion))
252    }
253}
254
255impl AssistantEditKind {
256    fn resolve(self, snapshot: &BufferSnapshot) -> ResolvedEdit {
257        match self {
258            Self::Update {
259                old_text,
260                new_text,
261                description,
262            } => {
263                let range = Self::resolve_location(&snapshot, &old_text);
264                ResolvedEdit {
265                    range,
266                    new_text,
267                    description: Some(description),
268                }
269            }
270            Self::Create {
271                new_text,
272                description,
273            } => ResolvedEdit {
274                range: text::Anchor::MIN..text::Anchor::MAX,
275                description: Some(description),
276                new_text,
277            },
278            Self::InsertBefore {
279                old_text,
280                mut new_text,
281                description,
282            } => {
283                let range = Self::resolve_location(&snapshot, &old_text);
284                new_text.push('\n');
285                ResolvedEdit {
286                    range: range.start..range.start,
287                    new_text,
288                    description: Some(description),
289                }
290            }
291            Self::InsertAfter {
292                old_text,
293                mut new_text,
294                description,
295            } => {
296                let range = Self::resolve_location(&snapshot, &old_text);
297                new_text.insert(0, '\n');
298                ResolvedEdit {
299                    range: range.end..range.end,
300                    new_text,
301                    description: Some(description),
302                }
303            }
304            Self::Delete { old_text } => {
305                let range = Self::resolve_location(&snapshot, &old_text);
306                ResolvedEdit {
307                    range,
308                    new_text: String::new(),
309                    description: None,
310                }
311            }
312        }
313    }
314
315    fn resolve_location(buffer: &text::BufferSnapshot, search_query: &str) -> Range<text::Anchor> {
316        const INSERTION_COST: u32 = 3;
317        const WHITESPACE_INSERTION_COST: u32 = 1;
318        const DELETION_COST: u32 = 3;
319        const WHITESPACE_DELETION_COST: u32 = 1;
320        const EQUALITY_BONUS: u32 = 5;
321
322        struct Matrix {
323            cols: usize,
324            data: Vec<SearchState>,
325        }
326
327        impl Matrix {
328            fn new(rows: usize, cols: usize) -> Self {
329                Matrix {
330                    cols,
331                    data: vec![SearchState::new(0, SearchDirection::Diagonal); rows * cols],
332                }
333            }
334
335            fn get(&self, row: usize, col: usize) -> SearchState {
336                self.data[row * self.cols + col]
337            }
338
339            fn set(&mut self, row: usize, col: usize, cost: SearchState) {
340                self.data[row * self.cols + col] = cost;
341            }
342        }
343
344        let buffer_len = buffer.len();
345        let query_len = search_query.len();
346        let mut matrix = Matrix::new(query_len + 1, buffer_len + 1);
347
348        for (row, query_byte) in search_query.bytes().enumerate() {
349            for (col, buffer_byte) in buffer.bytes_in_range(0..buffer.len()).flatten().enumerate() {
350                let deletion_cost = if query_byte.is_ascii_whitespace() {
351                    WHITESPACE_DELETION_COST
352                } else {
353                    DELETION_COST
354                };
355                let insertion_cost = if buffer_byte.is_ascii_whitespace() {
356                    WHITESPACE_INSERTION_COST
357                } else {
358                    INSERTION_COST
359                };
360
361                let up = SearchState::new(
362                    matrix.get(row, col + 1).score.saturating_sub(deletion_cost),
363                    SearchDirection::Up,
364                );
365                let left = SearchState::new(
366                    matrix
367                        .get(row + 1, col)
368                        .score
369                        .saturating_sub(insertion_cost),
370                    SearchDirection::Left,
371                );
372                let diagonal = SearchState::new(
373                    if query_byte == *buffer_byte {
374                        matrix.get(row, col).score.saturating_add(EQUALITY_BONUS)
375                    } else {
376                        matrix
377                            .get(row, col)
378                            .score
379                            .saturating_sub(deletion_cost + insertion_cost)
380                    },
381                    SearchDirection::Diagonal,
382                );
383                matrix.set(row + 1, col + 1, up.max(left).max(diagonal));
384            }
385        }
386
387        // Traceback to find the best match
388        let mut best_buffer_end = buffer_len;
389        let mut best_score = 0;
390        for col in 1..=buffer_len {
391            let score = matrix.get(query_len, col).score;
392            if score > best_score {
393                best_score = score;
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 text::{OffsetRangeExt, Point};
564    use ui::BorrowAppContext;
565    use unindent::Unindent as _;
566
567    #[gpui::test]
568    fn test_resolve_location(cx: &mut AppContext) {
569        {
570            let buffer = cx.new_model(|cx| {
571                Buffer::local(
572                    concat!(
573                        "    Lorem\n",
574                        "    ipsum\n",
575                        "    dolor sit amet\n",
576                        "    consecteur",
577                    ),
578                    cx,
579                )
580            });
581            let snapshot = buffer.read(cx).snapshot();
582            assert_eq!(
583                AssistantEditKind::resolve_location(&snapshot, "ipsum\ndolor").to_point(&snapshot),
584                Point::new(1, 0)..Point::new(2, 18)
585            );
586        }
587
588        {
589            let buffer = cx.new_model(|cx| {
590                Buffer::local(
591                    concat!(
592                        "fn foo1(a: usize) -> usize {\n",
593                        "    40\n",
594                        "}\n",
595                        "\n",
596                        "fn foo2(b: usize) -> usize {\n",
597                        "    42\n",
598                        "}\n",
599                    ),
600                    cx,
601                )
602            });
603            let snapshot = buffer.read(cx).snapshot();
604            assert_eq!(
605                AssistantEditKind::resolve_location(&snapshot, "fn foo1(b: usize) {\n40\n}")
606                    .to_point(&snapshot),
607                Point::new(0, 0)..Point::new(2, 1)
608            );
609        }
610
611        {
612            let buffer = cx.new_model(|cx| {
613                Buffer::local(
614                    concat!(
615                        "fn main() {\n",
616                        "    Foo\n",
617                        "        .bar()\n",
618                        "        .baz()\n",
619                        "        .qux()\n",
620                        "}\n",
621                        "\n",
622                        "fn foo2(b: usize) -> usize {\n",
623                        "    42\n",
624                        "}\n",
625                    ),
626                    cx,
627                )
628            });
629            let snapshot = buffer.read(cx).snapshot();
630            assert_eq!(
631                AssistantEditKind::resolve_location(&snapshot, "Foo.bar.baz.qux()")
632                    .to_point(&snapshot),
633                Point::new(1, 0)..Point::new(4, 14)
634            );
635        }
636    }
637
638    #[gpui::test]
639    fn test_resolve_edits(cx: &mut AppContext) {
640        let settings_store = SettingsStore::test(cx);
641        cx.set_global(settings_store);
642        language::init(cx);
643        cx.update_global::<SettingsStore, _>(|settings, cx| {
644            settings.update_user_settings::<AllLanguageSettings>(cx, |_| {});
645        });
646
647        assert_edits(
648            "
649                /// A person
650                struct Person {
651                    name: String,
652                    age: usize,
653                }
654
655                /// A dog
656                struct Dog {
657                    weight: f32,
658                }
659
660                impl Person {
661                    fn name(&self) -> &str {
662                        &self.name
663                    }
664                }
665            "
666            .unindent(),
667            vec![
668                AssistantEditKind::Update {
669                    old_text: "
670                        name: String,
671                    "
672                    .unindent(),
673                    new_text: "
674                        first_name: String,
675                        last_name: String,
676                    "
677                    .unindent(),
678                    description: "".into(),
679                },
680                AssistantEditKind::Update {
681                    old_text: "
682                        fn name(&self) -> &str {
683                            &self.name
684                        }
685                    "
686                    .unindent(),
687                    new_text: "
688                        fn name(&self) -> String {
689                            format!(\"{} {}\", self.first_name, self.last_name)
690                        }
691                    "
692                    .unindent(),
693                    description: "".into(),
694                },
695            ],
696            "
697                /// A person
698                struct Person {
699                    first_name: String,
700                    last_name: String,
701                    age: usize,
702                }
703
704                /// A dog
705                struct Dog {
706                    weight: f32,
707                }
708
709                impl Person {
710                    fn name(&self) -> String {
711                        format!(\"{} {}\", self.first_name, self.last_name)
712                    }
713                }
714            "
715            .unindent(),
716            cx,
717        );
718
719        // Ensure InsertBefore merges correctly with Update of the same text
720        assert_edits(
721            "
722                fn foo() {
723
724                }
725            "
726            .unindent(),
727            vec![
728                AssistantEditKind::InsertBefore {
729                    old_text: "
730                        fn foo() {"
731                        .unindent(),
732                    new_text: "
733                        fn bar() {
734                            qux();
735                        }"
736                    .unindent(),
737                    description: "implement bar".into(),
738                },
739                AssistantEditKind::Update {
740                    old_text: "
741                        fn foo() {
742
743                        }"
744                    .unindent(),
745                    new_text: "
746                        fn foo() {
747                            bar();
748                        }"
749                    .unindent(),
750                    description: "call bar in foo".into(),
751                },
752                AssistantEditKind::InsertAfter {
753                    old_text: "
754                        fn foo() {
755
756                        }
757                    "
758                    .unindent(),
759                    new_text: "
760                        fn qux() {
761                            // todo
762                        }
763                    "
764                    .unindent(),
765                    description: "implement qux".into(),
766                },
767            ],
768            "
769                fn bar() {
770                    qux();
771                }
772
773                fn foo() {
774                    bar();
775                }
776
777                fn qux() {
778                    // todo
779                }
780            "
781            .unindent(),
782            cx,
783        );
784
785        // Correctly indent new text when replacing multiple adjacent indented blocks.
786        assert_edits(
787            "
788            impl Numbers {
789                fn one() {
790                    1
791                }
792
793                fn two() {
794                    2
795                }
796
797                fn three() {
798                    3
799                }
800            }
801            "
802            .unindent(),
803            vec![
804                AssistantEditKind::Update {
805                    old_text: "
806                        fn one() {
807                            1
808                        }
809                    "
810                    .unindent(),
811                    new_text: "
812                        fn one() {
813                            101
814                        }
815                    "
816                    .unindent(),
817                    description: "pick better number".into(),
818                },
819                AssistantEditKind::Update {
820                    old_text: "
821                        fn two() {
822                            2
823                        }
824                    "
825                    .unindent(),
826                    new_text: "
827                        fn two() {
828                            102
829                        }
830                    "
831                    .unindent(),
832                    description: "pick better number".into(),
833                },
834                AssistantEditKind::Update {
835                    old_text: "
836                        fn three() {
837                            3
838                        }
839                    "
840                    .unindent(),
841                    new_text: "
842                        fn three() {
843                            103
844                        }
845                    "
846                    .unindent(),
847                    description: "pick better number".into(),
848                },
849            ],
850            "
851                impl Numbers {
852                    fn one() {
853                        101
854                    }
855
856                    fn two() {
857                        102
858                    }
859
860                    fn three() {
861                        103
862                    }
863                }
864            "
865            .unindent(),
866            cx,
867        );
868    }
869
870    #[track_caller]
871    fn assert_edits(
872        old_text: String,
873        edits: Vec<AssistantEditKind>,
874        new_text: String,
875        cx: &mut AppContext,
876    ) {
877        let buffer =
878            cx.new_model(|cx| Buffer::local(old_text, cx).with_language(Arc::new(rust_lang()), cx));
879        let snapshot = buffer.read(cx).snapshot();
880        let resolved_edits = edits
881            .into_iter()
882            .map(|kind| kind.resolve(&snapshot))
883            .collect();
884        let edit_groups = AssistantPatch::group_edits(resolved_edits, &snapshot);
885        ResolvedPatch::apply_edit_groups(&edit_groups, &buffer, cx);
886        let actual_new_text = buffer.read(cx).text();
887        pretty_assertions::assert_eq!(actual_new_text, new_text);
888    }
889
890    fn rust_lang() -> Language {
891        Language::new(
892            LanguageConfig {
893                name: "Rust".into(),
894                matcher: LanguageMatcher {
895                    path_suffixes: vec!["rs".to_string()],
896                    ..Default::default()
897                },
898                ..Default::default()
899            },
900            Some(language::tree_sitter_rust::LANGUAGE.into()),
901        )
902        .with_indents_query(
903            r#"
904            (call_expression) @indent
905            (field_expression) @indent
906            (_ "(" ")" @end) @indent
907            (_ "{" "}" @end) @indent
908            "#,
909        )
910        .unwrap()
911    }
912}