After streaming generation is over, show a regular, batch diff in the file altered (#16350)

Release Notes:

- N/A

Cargo.lock

@@ -398,6 +398,7 @@ dependencies = [
  "serde_json",
  "serde_json_lenient",
  "settings",
+ "similar",
  "smallvec",
  "smol",
  "telemetry_events",

crates/assistant/Cargo.toml

@@ -68,6 +68,7 @@ serde.workspace = true
 serde_json.workspace = true
 settings.workspace = true
 smallvec.workspace = true
+similar.workspace = true
 smol.workspace = true
 telemetry_events.workspace = true
 terminal.workspace = true

crates/assistant/src/inline_assistant.rs

@@ -19,6 +19,7 @@ use fs::Fs;
 use futures::{
     channel::mpsc,
     future::{BoxFuture, LocalBoxFuture},
+    join,
     stream::{self, BoxStream},
     SinkExt, Stream, StreamExt,
 };
@@ -2445,12 +2446,12 @@ impl Codegen {
         self.diff = Diff::default();
         self.status = CodegenStatus::Pending;
         let mut edit_start = edit_range.start.to_offset(&snapshot);
-        self.generation = cx.spawn(|this, mut cx| {
+        self.generation = cx.spawn(|codegen, mut cx| {
             async move {
                 let chunks = stream.await;
                 let generate = async {
                     let (mut diff_tx, mut diff_rx) = mpsc::channel(1);
-                    let diff: Task<anyhow::Result<()>> =
+                    let line_based_stream_diff: Task<anyhow::Result<()>> =
                         cx.background_executor().spawn(async move {
                             let mut response_latency = None;
                             let request_start = Instant::now();
@@ -2501,10 +2502,10 @@ impl Codegen {
                         });
 
                     while let Some((char_ops, line_diff)) = diff_rx.next().await {
-                        this.update(&mut cx, |this, cx| {
+                        codegen.update(&mut cx, |codegen, cx| {
-                            this.last_equal_ranges.clear();
+                            codegen.last_equal_ranges.clear();
 
-                            let transaction = this.buffer.update(cx, |buffer, cx| {
+                            let transaction = codegen.buffer.update(cx, |buffer, cx| {
                                 // Avoid grouping assistant edits with user edits.
                                 buffer.finalize_last_transaction(cx);
 
@@ -2529,23 +2530,24 @@ impl Codegen {
                                                 let edit_range = snapshot.anchor_after(edit_start)
                                                     ..snapshot.anchor_before(edit_end);
                                                 edit_start = edit_end;
-                                                this.last_equal_ranges.push(edit_range);
+                                                codegen.last_equal_ranges.push(edit_range);
                                                 None
                                             }
                                         }),
                                     None,
                                     cx,
                                 );
-                                this.edit_position = Some(snapshot.anchor_after(edit_start));
+                                codegen.edit_position = Some(snapshot.anchor_after(edit_start));
 
                                 buffer.end_transaction(cx)
                             });
 
                             if let Some(transaction) = transaction {
-                                if let Some(first_transaction) = this.transformation_transaction_id
+                                if let Some(first_transaction) =
+                                    codegen.transformation_transaction_id
                                 {
                                     // Group all assistant edits into the first transaction.
-                                    this.buffer.update(cx, |buffer, cx| {
+                                    codegen.buffer.update(cx, |buffer, cx| {
                                         buffer.merge_transactions(
                                             transaction,
                                             first_transaction,
@@ -2553,26 +2555,35 @@ impl Codegen {
                                         )
                                     });
                                 } else {
-                                    this.transformation_transaction_id = Some(transaction);
+                                    codegen.transformation_transaction_id = Some(transaction);
-                                    this.buffer.update(cx, |buffer, cx| {
+                                    codegen.buffer.update(cx, |buffer, cx| {
                                         buffer.finalize_last_transaction(cx)
                                     });
                                 }
                             }
 
-                            this.update_diff(edit_range.clone(), line_diff, cx);
+                            codegen.reapply_line_based_diff(edit_range.clone(), line_diff, cx);
 
                             cx.notify();
                         })?;
                     }
 
-                    diff.await?;
+                    // Streaming stopped and we have the new text in the buffer, and a line-based diff applied for the whole new buffer.
+                    // That diff is not what a regular diff is and might look unexpected, ergo apply a regular diff.
+                    // It's fine to apply even if the rest of the line diffing fails, as no more hunks are coming through `diff_rx`.
+                    let batch_diff_task = codegen.update(&mut cx, |codegen, cx| {
+                        codegen.reapply_batch_diff(edit_range.clone(), cx)
+                    })?;
+                    let (line_based_stream_diff, ()) =
+                        join!(line_based_stream_diff, batch_diff_task);
+                    line_based_stream_diff?;
 
                     anyhow::Ok(())
                 };
 
                 let result = generate.await;
-                this.update(&mut cx, |this, cx| {
+                codegen
+                    .update(&mut cx, |this, cx| {
                         this.last_equal_ranges.clear();
                         if let Err(error) = result {
                             this.status = CodegenStatus::Error(error);
@@ -2614,7 +2625,7 @@ impl Codegen {
         });
     }
 
-    fn update_diff(
+    fn reapply_line_based_diff(
         &mut self,
         edit_range: Range<Anchor>,
         line_operations: Vec<LineOperation>,
@@ -2673,6 +2684,99 @@ impl Codegen {
             cx.notify();
         }
     }
+
+    fn reapply_batch_diff(
+        &mut self,
+        edit_range: Range<Anchor>,
+        cx: &mut ModelContext<Self>,
+    ) -> Task<()> {
+        let old_snapshot = self.snapshot.clone();
+        let old_range = edit_range.to_point(&old_snapshot);
+        let new_snapshot = self.buffer.read(cx).snapshot(cx);
+        let new_range = edit_range.to_point(&new_snapshot);
+
+        cx.spawn(|codegen, mut cx| async move {
+            let (deleted_row_ranges, inserted_row_ranges) = cx
+                .background_executor()
+                .spawn(async move {
+                    let old_text = old_snapshot
+                        .text_for_range(
+                            Point::new(old_range.start.row, 0)
+                                ..Point::new(
+                                    old_range.end.row,
+                                    old_snapshot.line_len(MultiBufferRow(old_range.end.row)),
+                                ),
+                        )
+                        .collect::<String>();
+                    let new_text = new_snapshot
+                        .text_for_range(
+                            Point::new(new_range.start.row, 0)
+                                ..Point::new(
+                                    new_range.end.row,
+                                    new_snapshot.line_len(MultiBufferRow(new_range.end.row)),
+                                ),
+                        )
+                        .collect::<String>();
+
+                    let mut old_row = old_range.start.row;
+                    let mut new_row = new_range.start.row;
+                    let batch_diff =
+                        similar::TextDiff::from_lines(old_text.as_str(), new_text.as_str());
+
+                    let mut deleted_row_ranges: Vec<(Anchor, RangeInclusive<u32>)> = Vec::new();
+                    let mut inserted_row_ranges = Vec::new();
+                    for change in batch_diff.iter_all_changes() {
+                        let line_count = change.value().lines().count() as u32;
+                        match change.tag() {
+                            similar::ChangeTag::Equal => {
+                                old_row += line_count;
+                                new_row += line_count;
+                            }
+                            similar::ChangeTag::Delete => {
+                                let old_end_row = old_row + line_count - 1;
+                                let new_row = new_snapshot.anchor_before(Point::new(new_row, 0));
+
+                                if let Some((_, last_deleted_row_range)) =
+                                    deleted_row_ranges.last_mut()
+                                {
+                                    if *last_deleted_row_range.end() + 1 == old_row {
+                                        *last_deleted_row_range =
+                                            *last_deleted_row_range.start()..=old_end_row;
+                                    } else {
+                                        deleted_row_ranges.push((new_row, old_row..=old_end_row));
+                                    }
+                                } else {
+                                    deleted_row_ranges.push((new_row, old_row..=old_end_row));
+                                }
+
+                                old_row += line_count;
+                            }
+                            similar::ChangeTag::Insert => {
+                                let new_end_row = new_row + line_count - 1;
+                                let start = new_snapshot.anchor_before(Point::new(new_row, 0));
+                                let end = new_snapshot.anchor_before(Point::new(
+                                    new_end_row,
+                                    new_snapshot.line_len(MultiBufferRow(new_end_row)),
+                                ));
+                                inserted_row_ranges.push(start..=end);
+                                new_row += line_count;
+                            }
+                        }
+                    }
+
+                    (deleted_row_ranges, inserted_row_ranges)
+                })
+                .await;
+
+            codegen
+                .update(&mut cx, |codegen, cx| {
+                    codegen.diff.deleted_row_ranges = deleted_row_ranges;
+                    codegen.diff.inserted_row_ranges = inserted_row_ranges;
+                    cx.notify();
+                })
+                .ok();
+        })
+    }
 }
 
 struct StripInvalidSpans<T> {