tenant.openagents/omega

1use anyhow::Result;
2use buffer_diff::BufferDiff;
3use gpui::{App, AppContext, AsyncApp, Context, Entity, Subscription, Task};
4use itertools::Itertools;
5use language::{
6    Anchor, Buffer, Capability, LanguageRegistry, OffsetRangeExt as _, Point, TextBuffer,
7};
8use multi_buffer::{MultiBuffer, PathKey, excerpt_context_lines};
9use std::{cmp::Reverse, ops::Range, path::Path, sync::Arc};
10use util::ResultExt;
11
12pub enum Diff {
13    Pending(PendingDiff),
14    Finalized(FinalizedDiff),
15}
16
17impl Diff {
18    pub fn finalized(
19        path: String,
20        old_text: Option<String>,
21        new_text: String,
22        language_registry: Arc<LanguageRegistry>,
23        cx: &mut Context<Self>,
24    ) -> Self {
25        let multibuffer = cx.new(|_cx| MultiBuffer::without_headers(Capability::ReadOnly));
26        let new_buffer = cx.new(|cx| Buffer::local(new_text, cx));
27        let base_text_exists = old_text.is_some();
28        let base_text = old_text.clone().unwrap_or(String::new()).into();
29        let task = cx.spawn({
30            let multibuffer = multibuffer.clone();
31            let path = path.clone();
32            let buffer = new_buffer.clone();
33            async move |_, cx| {
34                let language = language_registry
35                    .load_language_for_file_path(Path::new(&path))
36                    .await
37                    .log_err();
38
39                buffer.update(cx, |buffer, cx| buffer.set_language(language.clone(), cx));
40                buffer.update(cx, |buffer, _| buffer.parsing_idle()).await;
41
42                let diff = build_buffer_diff(
43                    old_text.unwrap_or("".into()).into(),
44                    base_text_exists,
45                    &buffer,
46                    cx,
47                )
48                .await?;
49
50                multibuffer.update(cx, |multibuffer, cx| {
51                    let hunk_ranges = {
52                        let buffer = buffer.read(cx);
53                        diff.read(cx)
54                            .snapshot(cx)
55                            .hunks_intersecting_range(
56                                Anchor::min_for_buffer(buffer.remote_id())
57                                    ..Anchor::max_for_buffer(buffer.remote_id()),
58                                buffer,
59                            )
60                            .map(|diff_hunk| diff_hunk.buffer_range.to_point(buffer))
61                            .collect::<Vec<_>>()
62                    };
63
64                    multibuffer.set_excerpts_for_path(
65                        PathKey::for_buffer(&buffer, cx),
66                        buffer.clone(),
67                        hunk_ranges,
68                        excerpt_context_lines(cx),
69                        cx,
70                    );
71                    multibuffer.add_diff(diff, cx);
72                });
73
74                anyhow::Ok(())
75            }
76        });
77
78        Self::Finalized(FinalizedDiff {
79            multibuffer,
80            path,
81            base_text,
82            new_buffer,
83            _update_diff: task,
84        })
85    }
86
87    pub fn new(buffer: Entity<Buffer>, cx: &mut Context<Self>) -> Self {
88        let buffer_text_snapshot = buffer.read(cx).text_snapshot();
89        let language = buffer.read(cx).language().cloned();
90        let language_registry = buffer.read(cx).language_registry();
91        let buffer_diff = cx.new(|cx| {
92            BufferDiff::new_unchanged(&buffer_text_snapshot, language, language_registry, cx)
93        });
94
95        let multibuffer = cx.new(|cx| {
96            let mut multibuffer = MultiBuffer::without_headers(Capability::ReadOnly);
97            multibuffer.add_diff(buffer_diff.clone(), cx);
98            multibuffer
99        });
100
101        Self::Pending(PendingDiff {
102            multibuffer,
103            base_text: Arc::from(buffer_text_snapshot.text().as_str()),
104            _subscription: cx.observe(&buffer, |this, _, cx| {
105                if let Diff::Pending(diff) = this {
106                    diff.update(cx);
107                }
108            }),
109            new_buffer: buffer,
110            diff: buffer_diff,
111            revealed_ranges: Vec::new(),
112            update_diff: Task::ready(Ok(())),
113        })
114    }
115
116    pub fn reveal_range(&mut self, range: Range<Anchor>, cx: &mut Context<Self>) {
117        if let Self::Pending(diff) = self {
118            diff.reveal_range(range, cx);
119        }
120    }
121
122    pub fn finalize(&mut self, cx: &mut Context<Self>) {
123        if let Self::Pending(diff) = self {
124            *self = Self::Finalized(diff.finalize(cx));
125        }
126    }
127
128    /// Returns the original text before any edits were applied.
129    pub fn base_text(&self) -> &Arc<str> {
130        match self {
131            Self::Pending(PendingDiff { base_text, .. }) => base_text,
132            Self::Finalized(FinalizedDiff { base_text, .. }) => base_text,
133        }
134    }
135
136    /// Returns the buffer being edited (for pending diffs) or the snapshot buffer (for finalized diffs).
137    pub fn buffer(&self) -> &Entity<Buffer> {
138        match self {
139            Self::Pending(PendingDiff { new_buffer, .. }) => new_buffer,
140            Self::Finalized(FinalizedDiff { new_buffer, .. }) => new_buffer,
141        }
142    }
143
144    pub fn file_path(&self, cx: &App) -> Option<String> {
145        match self {
146            Self::Pending(PendingDiff { new_buffer, .. }) => new_buffer
147                .read(cx)
148                .file()
149                .map(|file| file.full_path(cx).to_string_lossy().into_owned()),
150            Self::Finalized(FinalizedDiff { path, .. }) => Some(path.clone()),
151        }
152    }
153
154    pub fn multibuffer(&self) -> &Entity<MultiBuffer> {
155        match self {
156            Self::Pending(PendingDiff { multibuffer, .. }) => multibuffer,
157            Self::Finalized(FinalizedDiff { multibuffer, .. }) => multibuffer,
158        }
159    }
160
161    pub fn to_markdown(&self, cx: &App) -> String {
162        let buffer_text = self
163            .multibuffer()
164            .read(cx)
165            .all_buffers()
166            .iter()
167            .map(|buffer| buffer.read(cx).text())
168            .join("\n");
169        let path = match self {
170            Diff::Pending(PendingDiff {
171                new_buffer: buffer, ..
172            }) => buffer
173                .read(cx)
174                .file()
175                .map(|file| file.path().display(file.path_style(cx))),
176            Diff::Finalized(FinalizedDiff { path, .. }) => Some(path.as_str().into()),
177        };
178        format!(
179            "Diff: {}\n```\n{}\n```\n",
180            path.unwrap_or(MultiBuffer::DEFAULT_TITLE.into()),
181            buffer_text
182        )
183    }
184
185    pub fn has_revealed_range(&self, cx: &App) -> bool {
186        !self.multibuffer().read(cx).is_empty()
187    }
188
189    pub fn needs_update(&self, old_text: &str, new_text: &str, cx: &App) -> bool {
190        match self {
191            Diff::Pending(PendingDiff {
192                base_text,
193                new_buffer,
194                ..
195            }) => {
196                base_text.as_ref() != old_text
197                    || !new_buffer.read(cx).as_rope().chunks().equals_str(new_text)
198            }
199            Diff::Finalized(FinalizedDiff {
200                base_text,
201                new_buffer,
202                ..
203            }) => {
204                base_text.as_ref() != old_text
205                    || !new_buffer.read(cx).as_rope().chunks().equals_str(new_text)
206            }
207        }
208    }
209}
210
211pub struct PendingDiff {
212    multibuffer: Entity<MultiBuffer>,
213    base_text: Arc<str>,
214    new_buffer: Entity<Buffer>,
215    diff: Entity<BufferDiff>,
216    revealed_ranges: Vec<Range<Anchor>>,
217    _subscription: Subscription,
218    update_diff: Task<Result<()>>,
219}
220
221impl PendingDiff {
222    pub fn update(&mut self, cx: &mut Context<Diff>) {
223        let buffer = self.new_buffer.clone();
224        let buffer_diff = self.diff.clone();
225        let base_text = self.base_text.clone();
226        self.update_diff = cx.spawn(async move |diff, cx| {
227            let text_snapshot = buffer.read_with(cx, |buffer, _| buffer.text_snapshot());
228            let base_text_snapshot = buffer_diff.read_with(cx, |diff, cx| diff.base_text(cx));
229            let update = buffer_diff
230                .update(cx, |diff, cx| {
231                    diff.update_diff(
232                        text_snapshot.clone(),
233                        &base_text_snapshot,
234                        Some(base_text.clone()),
235                        cx,
236                    )
237                })
238                .await;
239            buffer_diff.update(cx, |diff, cx| {
240                diff.set_snapshot(update.clone(), cx);
241            });
242            diff.update(cx, |diff, cx| {
243                if let Diff::Pending(diff) = diff {
244                    diff.update_visible_ranges(cx);
245                }
246            })
247        });
248    }
249
250    pub fn reveal_range(&mut self, range: Range<Anchor>, cx: &mut Context<Diff>) {
251        self.revealed_ranges.push(range);
252        self.update_visible_ranges(cx);
253    }
254
255    fn finalize(&self, cx: &mut Context<Diff>) -> FinalizedDiff {
256        let ranges = self.excerpt_ranges(cx);
257        let base_text = self.base_text.clone();
258        let new_buffer = self.new_buffer.read(cx);
259
260        let path = new_buffer
261            .file()
262            .map(|file| file.path().display(file.path_style(cx)))
263            .unwrap_or(MultiBuffer::DEFAULT_TITLE.into())
264            .into();
265        let replica_id = new_buffer.replica_id();
266
267        // Replace the buffer in the multibuffer with the snapshot
268        let buffer = cx.new(|cx| {
269            let language = self.new_buffer.read(cx).language().cloned();
270            let buffer = TextBuffer::new_normalized(
271                replica_id,
272                cx.entity_id().as_non_zero_u64().into(),
273                self.new_buffer.read(cx).line_ending(),
274                self.new_buffer.read(cx).as_rope().clone(),
275            );
276            let mut buffer = Buffer::build(buffer, None, Capability::ReadWrite);
277            buffer.set_language(language, cx);
278            buffer
279        });
280
281        let buffer_diff = cx.spawn({
282            let buffer = buffer.clone();
283            async move |_this, cx| {
284                buffer.update(cx, |buffer, _| buffer.parsing_idle()).await;
285                build_buffer_diff(base_text, true, &buffer, cx).await
286            }
287        });
288
289        let update_diff = cx.spawn(async move |this, cx| {
290            let buffer_diff = buffer_diff.await?;
291            this.update(cx, |this, cx| {
292                this.multibuffer().update(cx, |multibuffer, cx| {
293                    let path_key = PathKey::for_buffer(&buffer, cx);
294                    multibuffer.clear(cx);
295                    multibuffer.set_excerpts_for_path(
296                        path_key,
297                        buffer,
298                        ranges,
299                        excerpt_context_lines(cx),
300                        cx,
301                    );
302                    multibuffer.add_diff(buffer_diff.clone(), cx);
303                });
304
305                cx.notify();
306            })
307        });
308
309        FinalizedDiff {
310            path,
311            base_text: self.base_text.clone(),
312            multibuffer: self.multibuffer.clone(),
313            new_buffer: self.new_buffer.clone(),
314            _update_diff: update_diff,
315        }
316    }
317
318    fn update_visible_ranges(&mut self, cx: &mut Context<Diff>) {
319        let ranges = self.excerpt_ranges(cx);
320        self.multibuffer.update(cx, |multibuffer, cx| {
321            multibuffer.set_excerpts_for_path(
322                PathKey::for_buffer(&self.new_buffer, cx),
323                self.new_buffer.clone(),
324                ranges,
325                excerpt_context_lines(cx),
326                cx,
327            );
328            let end = multibuffer.len(cx);
329            Some(multibuffer.snapshot(cx).offset_to_point(end).row + 1)
330        });
331        cx.notify();
332    }
333
334    fn excerpt_ranges(&self, cx: &App) -> Vec<Range<Point>> {
335        let buffer = self.new_buffer.read(cx);
336        let mut ranges = self
337            .diff
338            .read(cx)
339            .snapshot(cx)
340            .hunks_intersecting_range(
341                Anchor::min_for_buffer(buffer.remote_id())
342                    ..Anchor::max_for_buffer(buffer.remote_id()),
343                buffer,
344            )
345            .map(|diff_hunk| diff_hunk.buffer_range.to_point(buffer))
346            .collect::<Vec<_>>();
347        ranges.extend(
348            self.revealed_ranges
349                .iter()
350                .map(|range| range.to_point(buffer)),
351        );
352        ranges.sort_unstable_by_key(|range| (range.start, Reverse(range.end)));
353
354        // Merge adjacent ranges
355        let mut ranges = ranges.into_iter().peekable();
356        let mut merged_ranges = Vec::new();
357        while let Some(mut range) = ranges.next() {
358            while let Some(next_range) = ranges.peek() {
359                if range.end >= next_range.start {
360                    range.end = range.end.max(next_range.end);
361                    ranges.next();
362                } else {
363                    break;
364                }
365            }
366
367            merged_ranges.push(range);
368        }
369        merged_ranges
370    }
371}
372
373pub struct FinalizedDiff {
374    path: String,
375    base_text: Arc<str>,
376    new_buffer: Entity<Buffer>,
377    multibuffer: Entity<MultiBuffer>,
378    _update_diff: Task<Result<()>>,
379}
380
381async fn build_buffer_diff(
382    old_text: Arc<str>,
383    base_text_exists: bool,
384    buffer: &Entity<Buffer>,
385    cx: &mut AsyncApp,
386) -> Result<Entity<BufferDiff>> {
387    let language = cx.update(|cx| buffer.read(cx).language().cloned());
388    let language_registry = cx.update(|cx| buffer.read(cx).language_registry());
389    let buffer = cx.update(|cx| buffer.read(cx).snapshot());
390    let base_text = base_text_exists.then(|| old_text);
391
392    let diff = cx.new(|cx| BufferDiff::new(&buffer, language, language_registry, cx));
393    diff.update(cx, |diff, cx| {
394        diff.set_base_text(base_text, buffer.text, cx)
395    })
396    .await;
397    Ok(diff)
398}
399
400#[cfg(test)]
401mod tests {
402    use gpui::{AppContext as _, TestAppContext};
403    use language::Buffer;
404
405    use crate::Diff;
406
407    #[gpui::test]
408    async fn test_pending_diff(cx: &mut TestAppContext) {
409        let buffer = cx.new(|cx| Buffer::local("hello!", cx));
410        let _diff = cx.new(|cx| Diff::new(buffer.clone(), cx));
411        buffer.update(cx, |buffer, cx| {
412            buffer.set_text("HELLO!", cx);
413        });
414        cx.run_until_parked();
415    }
416}
417