Cargo.lock 🔗
@@ -398,6 +398,7 @@ dependencies = [
"serde_json",
"serde_json_lenient",
"settings",
+ "similar",
"smallvec",
"smol",
"telemetry_events",
After streaming generation is over, show a regular, batch diff in the file altered (#16350)
Kirill Bulatov created
Release Notes:
- N/AChange summary
Cargo.lock | 1
crates/assistant/Cargo.toml | 1
crates/assistant/src/inline_assistant.rs | 154 +++++++++++++++++++++----
3 files changed, 131 insertions(+), 25 deletions(-)
Detailed changes
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| {
- this.last_equal_ranges.clear();
+ codegen.update(&mut cx, |codegen, cx| {
+ 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,36 +2555,45 @@ impl Codegen {
)
});
} else {
- this.transformation_transaction_id = Some(transaction);
- this.buffer.update(cx, |buffer, cx| {
+ codegen.transformation_transaction_id = Some(transaction);
+ 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| {
- this.last_equal_ranges.clear();
- if let Err(error) = result {
- this.status = CodegenStatus::Error(error);
- } else {
- this.status = CodegenStatus::Done;
- }
- cx.emit(CodegenEvent::Finished);
- cx.notify();
- })
- .ok();
+ codegen
+ .update(&mut cx, |this, cx| {
+ this.last_equal_ranges.clear();
+ if let Err(error) = result {
+ this.status = CodegenStatus::Error(error);
+ } else {
+ this.status = CodegenStatus::Done;
+ }
+ cx.emit(CodegenEvent::Finished);
+ cx.notify();
+ })
+ .ok();
}
});
cx.notify();
@@ -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> {