From 658d616b9661f086681b25e33b5e76620fd18b1d Mon Sep 17 00:00:00 2001 From: Antonio Scandurra Date: Sat, 26 Aug 2023 11:55:03 +0200 Subject: [PATCH] Allow multiple inline assistant highlights at once --- crates/ai/src/assistant.rs | 216 +++++++++++++++++++++++------------- crates/editor/src/editor.rs | 1 + todo.md | 4 +- 3 files changed, 143 insertions(+), 78 deletions(-) diff --git a/crates/ai/src/assistant.rs b/crates/ai/src/assistant.rs index b99d4b4fac..c5bf027fcc 100644 --- a/crates/ai/src/assistant.rs +++ b/crates/ai/src/assistant.rs @@ -13,13 +13,14 @@ use editor::{ BlockContext, BlockDisposition, BlockId, BlockProperties, BlockStyle, ToDisplayPoint, }, scroll::autoscroll::{Autoscroll, AutoscrollStrategy}, - Anchor, Editor, ToOffset, ToPoint, + Anchor, Editor, MultiBufferSnapshot, ToOffset, ToPoint, }; use fs::Fs; use futures::{channel::mpsc, SinkExt, StreamExt}; use gpui::{ actions, elements::*, + fonts::HighlightStyle, geometry::vector::{vec2f, Vector2F}, platform::{CursorStyle, MouseButton}, Action, AppContext, AsyncAppContext, ClipboardItem, Entity, ModelContext, ModelHandle, @@ -279,11 +280,6 @@ impl AssistantPanel { editor.change_selections(None, cx, |selections| { selections.select_anchor_ranges([selection.head()..selection.head()]) }); - editor.highlight_background::( - vec![range.clone()], - |theme| theme.assistant.inline.pending_edit_background, - cx, - ); editor.insert_blocks( [BlockProperties { style: BlockStyle::Flex, @@ -318,6 +314,7 @@ impl AssistantPanel { kind: assist_kind, editor: editor.downgrade(), range, + highlighted_ranges: Default::default(), inline_assistant_block_id: Some(block_id), code_generation: Task::ready(None), transaction_id: None, @@ -342,6 +339,7 @@ impl AssistantPanel { .entry(editor.downgrade()) .or_default() .push(id); + self.update_highlights_for_editor(&editor, cx); } fn handle_inline_assistant_event( @@ -416,10 +414,7 @@ impl AssistantPanel { } if let Some(editor) = pending_assist.editor.upgrade(cx) { - editor.update(cx, |editor, cx| { - editor.clear_background_highlights::(cx); - editor.clear_text_highlights::(cx); - }); + self.update_highlights_for_editor(&editor, cx); if cancel { if let Some(transaction_id) = pending_assist.transaction_id { @@ -741,78 +736,75 @@ impl AssistantPanel { }); while let Some(hunks) = hunks_rx.next().await { - let this = this + let editor = editor .upgrade(&cx) - .ok_or_else(|| anyhow!("assistant was dropped"))?; - editor.update(&mut cx, |editor, cx| { - let mut highlights = Vec::new(); + .ok_or_else(|| anyhow!("editor was dropped"))?; - let transaction = editor.buffer().update(cx, |buffer, cx| { - // Avoid grouping assistant edits with user edits. - buffer.finalize_last_transaction(cx); + this.update(&mut cx, |this, cx| { + let pending_assist = if let Some(pending_assist) = + this.pending_inline_assists.get_mut(&inline_assist_id) + { + pending_assist + } else { + return; + }; - buffer.start_transaction(cx); - buffer.edit( - hunks.into_iter().filter_map(|hunk| match hunk { - Hunk::Insert { text } => { - let edit_start = snapshot.anchor_after(edit_start); - Some((edit_start..edit_start, text)) - } - Hunk::Remove { len } => { - let edit_end = edit_start + len; - let edit_range = snapshot.anchor_after(edit_start) - ..snapshot.anchor_before(edit_end); - edit_start = edit_end; - Some((edit_range, String::new())) - } - Hunk::Keep { len } => { - let edit_end = edit_start + len; - let edit_range = snapshot.anchor_after(edit_start) - ..snapshot.anchor_before(edit_end); - edit_start += len; - highlights.push(edit_range); - None - } - }), - None, - cx, - ); + pending_assist.highlighted_ranges.clear(); + editor.update(cx, |editor, cx| { + let transaction = editor.buffer().update(cx, |buffer, cx| { + // Avoid grouping assistant edits with user edits. + buffer.finalize_last_transaction(cx); - buffer.end_transaction(cx) + buffer.start_transaction(cx); + buffer.edit( + hunks.into_iter().filter_map(|hunk| match hunk { + Hunk::Insert { text } => { + let edit_start = snapshot.anchor_after(edit_start); + Some((edit_start..edit_start, text)) + } + Hunk::Remove { len } => { + let edit_end = edit_start + len; + let edit_range = snapshot.anchor_after(edit_start) + ..snapshot.anchor_before(edit_end); + edit_start = edit_end; + Some((edit_range, String::new())) + } + Hunk::Keep { len } => { + let edit_end = edit_start + len; + let edit_range = snapshot.anchor_after(edit_start) + ..snapshot.anchor_before(edit_end); + edit_start += len; + pending_assist.highlighted_ranges.push(edit_range); + None + } + }), + None, + cx, + ); + + buffer.end_transaction(cx) + }); + + if let Some(transaction) = transaction { + if let Some(first_transaction) = pending_assist.transaction_id { + // Group all assistant edits into the first transaction. + editor.buffer().update(cx, |buffer, cx| { + buffer.merge_transactions( + transaction, + first_transaction, + cx, + ) + }); + } else { + pending_assist.transaction_id = Some(transaction); + editor.buffer().update(cx, |buffer, cx| { + buffer.finalize_last_transaction(cx) + }); + } + } }); - if let Some(transaction) = transaction { - this.update(cx, |this, cx| { - if let Some(pending_assist) = - this.pending_inline_assists.get_mut(&inline_assist_id) - { - if let Some(first_transaction) = pending_assist.transaction_id { - // Group all assistant edits into the first transaction. - editor.buffer().update(cx, |buffer, cx| { - buffer.merge_transactions( - transaction, - first_transaction, - cx, - ) - }); - } else { - pending_assist.transaction_id = Some(transaction); - editor.buffer().update(cx, |buffer, cx| { - buffer.finalize_last_transaction(cx) - }); - } - } - }); - } - - editor.highlight_text::( - highlights, - gpui::fonts::HighlightStyle { - fade_out: Some(0.6), - ..Default::default() - }, - cx, - ); + this.update_highlights_for_editor(&editor, cx); })?; } diff.await?; @@ -823,6 +815,55 @@ impl AssistantPanel { }); } + fn update_highlights_for_editor( + &self, + editor: &ViewHandle, + cx: &mut ViewContext, + ) { + let mut background_ranges = Vec::new(); + let mut foreground_ranges = Vec::new(); + let empty_inline_assist_ids = Vec::new(); + let inline_assist_ids = self + .pending_inline_assist_ids_by_editor + .get(&editor.downgrade()) + .unwrap_or(&empty_inline_assist_ids); + + for inline_assist_id in inline_assist_ids { + if let Some(pending_assist) = self.pending_inline_assists.get(inline_assist_id) { + background_ranges.push(pending_assist.range.clone()); + foreground_ranges.extend(pending_assist.highlighted_ranges.iter().cloned()); + } + } + + let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx); + merge_ranges(&mut background_ranges, &snapshot); + merge_ranges(&mut foreground_ranges, &snapshot); + editor.update(cx, |editor, cx| { + if background_ranges.is_empty() { + editor.clear_background_highlights::(cx); + } else { + editor.highlight_background::( + background_ranges, + |theme| theme.assistant.inline.pending_edit_background, + cx, + ); + } + + if foreground_ranges.is_empty() { + editor.clear_text_highlights::(cx); + } else { + editor.highlight_text::( + foreground_ranges, + HighlightStyle { + fade_out: Some(0.6), + ..Default::default() + }, + cx, + ); + } + }); + } + fn new_conversation(&mut self, cx: &mut ViewContext) -> ViewHandle { let editor = cx.add_view(|cx| { ConversationEditor::new( @@ -2842,12 +2883,35 @@ struct PendingInlineAssist { kind: InlineAssistKind, editor: WeakViewHandle, range: Range, + highlighted_ranges: Vec>, inline_assistant_block_id: Option, code_generation: Task>, transaction_id: Option, _subscriptions: Vec, } +fn merge_ranges(ranges: &mut Vec>, buffer: &MultiBufferSnapshot) { + ranges.sort_unstable_by(|a, b| { + a.start + .cmp(&b.start, buffer) + .then_with(|| b.end.cmp(&a.end, buffer)) + }); + + let mut ix = 0; + while ix + 1 < ranges.len() { + let b = ranges[ix + 1].clone(); + let a = &mut ranges[ix]; + if a.end.cmp(&b.start, buffer).is_gt() { + if a.end.cmp(&b.end, buffer).is_lt() { + a.end = b.end; + } + ranges.remove(ix + 1); + } else { + ix += 1; + } + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/editor/src/editor.rs b/crates/editor/src/editor.rs index b206f2ec8b..d5141a7f7d 100644 --- a/crates/editor/src/editor.rs +++ b/crates/editor/src/editor.rs @@ -7592,6 +7592,7 @@ impl Editor { } results } + pub fn background_highlights_in_range_for( &self, search_range: Range, diff --git a/todo.md b/todo.md index 71ca5a7c7b..0507708502 100644 --- a/todo.md +++ b/todo.md @@ -11,8 +11,8 @@ - When cursor is inside a prompt - [x] Escape cancels/undoes - [x] Enter confirms -- [ ] Selection is cleared and cursor is moved to prompt input -- [ ] Ability to highlight background multiple times for the same type +- [x] Selection is cleared and cursor is moved to prompt input +- [x] Ability to highlight background multiple times for the same type - [x] Basic Styling - [ ] Look into why insert prompts have a weird indentation sometimes