From d9c08de58a7a164df43614791e5318d5d9824a05 Mon Sep 17 00:00:00 2001 From: Antonio Scandurra Date: Fri, 29 Sep 2023 17:15:26 +0200 Subject: [PATCH 1/8] Revert "Revert "leverage file outline and selection as opposed to entire file"" --- crates/assistant/src/assistant.rs | 1 + crates/assistant/src/assistant_panel.rs | 117 ++------ crates/assistant/src/prompts.rs | 382 ++++++++++++++++++++++++ 3 files changed, 411 insertions(+), 89 deletions(-) create mode 100644 crates/assistant/src/prompts.rs diff --git a/crates/assistant/src/assistant.rs b/crates/assistant/src/assistant.rs index 258684db47096bac2d3df33d0289462dbc841214..6c9b14333e34cbf5fd49d8299ba7bd891b607526 100644 --- a/crates/assistant/src/assistant.rs +++ b/crates/assistant/src/assistant.rs @@ -1,6 +1,7 @@ pub mod assistant_panel; mod assistant_settings; mod codegen; +mod prompts; mod streaming_diff; use ai::completion::Role; diff --git a/crates/assistant/src/assistant_panel.rs b/crates/assistant/src/assistant_panel.rs index 42e5fb78979a6b8136c5c60d29e38e064df3435d..37d0d729fe64e0def057402fb9a24b796ebdf317 100644 --- a/crates/assistant/src/assistant_panel.rs +++ b/crates/assistant/src/assistant_panel.rs @@ -1,6 +1,7 @@ use crate::{ assistant_settings::{AssistantDockPosition, AssistantSettings, OpenAIModel}, codegen::{self, Codegen, CodegenKind}, + prompts::generate_content_prompt, MessageId, MessageMetadata, MessageStatus, Role, SavedConversation, SavedConversationMetadata, SavedMessage, }; @@ -541,11 +542,25 @@ impl AssistantPanel { self.inline_prompt_history.pop_front(); } - let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx); + let multi_buffer = editor.read(cx).buffer().read(cx); + let multi_buffer_snapshot = multi_buffer.snapshot(cx); + let snapshot = if multi_buffer.is_singleton() { + multi_buffer.as_singleton().unwrap().read(cx).snapshot() + } else { + return; + }; + let range = pending_assist.codegen.read(cx).range(); - let selected_text = snapshot.text_for_range(range.clone()).collect::(); + let language_range = snapshot.anchor_at( + range.start.to_offset(&multi_buffer_snapshot), + language::Bias::Left, + ) + ..snapshot.anchor_at( + range.end.to_offset(&multi_buffer_snapshot), + language::Bias::Right, + ); - let language = snapshot.language_at(range.start); + let language = snapshot.language_at(language_range.start); let language_name = if let Some(language) = language.as_ref() { if Arc::ptr_eq(language, &language::PLAIN_TEXT) { None @@ -557,93 +572,17 @@ impl AssistantPanel { }; let language_name = language_name.as_deref(); - let mut prompt = String::new(); - if let Some(language_name) = language_name { - writeln!(prompt, "You're an expert {language_name} engineer.").unwrap(); - } - match pending_assist.codegen.read(cx).kind() { - CodegenKind::Transform { .. } => { - writeln!( - prompt, - "You're currently working inside an editor on this file:" - ) - .unwrap(); - if let Some(language_name) = language_name { - writeln!(prompt, "```{language_name}").unwrap(); - } else { - writeln!(prompt, "```").unwrap(); - } - for chunk in snapshot.text_for_range(Anchor::min()..Anchor::max()) { - write!(prompt, "{chunk}").unwrap(); - } - writeln!(prompt, "```").unwrap(); + let codegen_kind = pending_assist.codegen.read(cx).kind().clone(); + let prompt = generate_content_prompt( + user_prompt.to_string(), + language_name, + &snapshot, + language_range, + cx, + codegen_kind, + ); - writeln!( - prompt, - "In particular, the user has selected the following text:" - ) - .unwrap(); - if let Some(language_name) = language_name { - writeln!(prompt, "```{language_name}").unwrap(); - } else { - writeln!(prompt, "```").unwrap(); - } - writeln!(prompt, "{selected_text}").unwrap(); - writeln!(prompt, "```").unwrap(); - writeln!(prompt).unwrap(); - writeln!( - prompt, - "Modify the selected text given the user prompt: {user_prompt}" - ) - .unwrap(); - writeln!( - prompt, - "You MUST reply only with the edited selected text, not the entire file." - ) - .unwrap(); - } - CodegenKind::Generate { .. } => { - writeln!( - prompt, - "You're currently working inside an editor on this file:" - ) - .unwrap(); - if let Some(language_name) = language_name { - writeln!(prompt, "```{language_name}").unwrap(); - } else { - writeln!(prompt, "```").unwrap(); - } - for chunk in snapshot.text_for_range(Anchor::min()..range.start) { - write!(prompt, "{chunk}").unwrap(); - } - write!(prompt, "<|>").unwrap(); - for chunk in snapshot.text_for_range(range.start..Anchor::max()) { - write!(prompt, "{chunk}").unwrap(); - } - writeln!(prompt).unwrap(); - writeln!(prompt, "```").unwrap(); - writeln!( - prompt, - "Assume the cursor is located where the `<|>` marker is." - ) - .unwrap(); - writeln!( - prompt, - "Text can't be replaced, so assume your answer will be inserted at the cursor." - ) - .unwrap(); - writeln!( - prompt, - "Complete the text given the user prompt: {user_prompt}" - ) - .unwrap(); - } - } - if let Some(language_name) = language_name { - writeln!(prompt, "Your answer MUST always be valid {language_name}.").unwrap(); - } - writeln!(prompt, "Always wrap your response in a Markdown codeblock.").unwrap(); - writeln!(prompt, "Never make remarks about the output.").unwrap(); + dbg!(&prompt); let mut messages = Vec::new(); let mut model = settings::get::(cx) diff --git a/crates/assistant/src/prompts.rs b/crates/assistant/src/prompts.rs new file mode 100644 index 0000000000000000000000000000000000000000..ee58090d04dbff9d7da058a905c9e52b8fe1b0cd --- /dev/null +++ b/crates/assistant/src/prompts.rs @@ -0,0 +1,382 @@ +use gpui::AppContext; +use language::{BufferSnapshot, OffsetRangeExt, ToOffset}; +use std::cmp; +use std::ops::Range; +use std::{fmt::Write, iter}; + +use crate::codegen::CodegenKind; + +fn outline_for_prompt( + buffer: &BufferSnapshot, + range: Range, + cx: &AppContext, +) -> Option { + let indent = buffer + .language_indent_size_at(0, cx) + .chars() + .collect::(); + let outline = buffer.outline(None)?; + let range = range.to_offset(buffer); + + let mut text = String::new(); + let mut items = outline.items.into_iter().peekable(); + + let mut intersected = false; + let mut intersection_indent = 0; + let mut extended_range = range.clone(); + + while let Some(item) = items.next() { + let item_range = item.range.to_offset(buffer); + if item_range.end < range.start || item_range.start > range.end { + text.extend(iter::repeat(indent.as_str()).take(item.depth)); + text.push_str(&item.text); + text.push('\n'); + } else { + intersected = true; + let is_terminal = items + .peek() + .map_or(true, |next_item| next_item.depth <= item.depth); + if is_terminal { + if item_range.start <= extended_range.start { + extended_range.start = item_range.start; + intersection_indent = item.depth; + } + extended_range.end = cmp::max(extended_range.end, item_range.end); + } else { + let name_start = item_range.start + item.name_ranges.first().unwrap().start; + let name_end = item_range.start + item.name_ranges.last().unwrap().end; + + if range.start > name_end { + text.extend(iter::repeat(indent.as_str()).take(item.depth)); + text.push_str(&item.text); + text.push('\n'); + } else { + if name_start <= extended_range.start { + extended_range.start = item_range.start; + intersection_indent = item.depth; + } + extended_range.end = cmp::max(extended_range.end, name_end); + } + } + } + + if intersected + && items.peek().map_or(true, |next_item| { + next_item.range.start.to_offset(buffer) > range.end + }) + { + intersected = false; + text.extend(iter::repeat(indent.as_str()).take(intersection_indent)); + text.extend(buffer.text_for_range(extended_range.start..range.start)); + text.push_str("<|START|"); + text.extend(buffer.text_for_range(range.clone())); + if range.start != range.end { + text.push_str("|END|>"); + } else { + text.push_str(">"); + } + text.extend(buffer.text_for_range(range.end..extended_range.end)); + text.push('\n'); + } + } + + Some(text) +} + +pub fn generate_content_prompt( + user_prompt: String, + language_name: Option<&str>, + buffer: &BufferSnapshot, + range: Range, + cx: &AppContext, + kind: CodegenKind, +) -> String { + let mut prompt = String::new(); + + // General Preamble + if let Some(language_name) = language_name { + writeln!(prompt, "You're an expert {language_name} engineer.\n").unwrap(); + } else { + writeln!(prompt, "You're an expert engineer.\n").unwrap(); + } + + let outline = outline_for_prompt(buffer, range.clone(), cx); + if let Some(outline) = outline { + writeln!( + prompt, + "The file you are currently working on has the following outline:" + ) + .unwrap(); + if let Some(language_name) = language_name { + let language_name = language_name.to_lowercase(); + writeln!(prompt, "```{language_name}\n{outline}\n```").unwrap(); + } else { + writeln!(prompt, "```\n{outline}\n```").unwrap(); + } + } + + // Assume for now that we are just generating + if range.clone().start == range.end { + writeln!(prompt, "In particular, the user's cursor is current on the '<|START|>' span in the above outline, with no text selected.").unwrap(); + } else { + writeln!(prompt, "In particular, the user has selected a section of the text between the '<|START|' and '|END|>' spans.").unwrap(); + } + + match kind { + CodegenKind::Generate { position: _ } => { + writeln!( + prompt, + "Assume the cursor is located where the `<|START|` marker is." + ) + .unwrap(); + writeln!( + prompt, + "Text can't be replaced, so assume your answer will be inserted at the cursor." + ) + .unwrap(); + writeln!( + prompt, + "Generate text based on the users prompt: {user_prompt}" + ) + .unwrap(); + } + CodegenKind::Transform { range: _ } => { + writeln!( + prompt, + "Modify the users code selected text based upon the users prompt: {user_prompt}" + ) + .unwrap(); + writeln!( + prompt, + "You MUST reply with only the adjusted code (within the '<|START|' and '|END|>' spans), not the entire file." + ) + .unwrap(); + } + } + + if let Some(language_name) = language_name { + writeln!(prompt, "Your answer MUST always be valid {language_name}").unwrap(); + } + writeln!(prompt, "Always wrap your response in a Markdown codeblock").unwrap(); + writeln!(prompt, "Never make remarks about the output.").unwrap(); + + prompt +} + +#[cfg(test)] +pub(crate) mod tests { + + use super::*; + use std::sync::Arc; + + use gpui::AppContext; + use indoc::indoc; + use language::{language_settings, tree_sitter_rust, Buffer, Language, LanguageConfig, Point}; + use settings::SettingsStore; + + pub(crate) fn rust_lang() -> Language { + Language::new( + LanguageConfig { + name: "Rust".into(), + path_suffixes: vec!["rs".to_string()], + ..Default::default() + }, + Some(tree_sitter_rust::language()), + ) + .with_indents_query( + r#" + (call_expression) @indent + (field_expression) @indent + (_ "(" ")" @end) @indent + (_ "{" "}" @end) @indent + "#, + ) + .unwrap() + .with_outline_query( + r#" + (struct_item + "struct" @context + name: (_) @name) @item + (enum_item + "enum" @context + name: (_) @name) @item + (enum_variant + name: (_) @name) @item + (field_declaration + name: (_) @name) @item + (impl_item + "impl" @context + trait: (_)? @name + "for"? @context + type: (_) @name) @item + (function_item + "fn" @context + name: (_) @name) @item + (mod_item + "mod" @context + name: (_) @name) @item + "#, + ) + .unwrap() + } + + #[gpui::test] + fn test_outline_for_prompt(cx: &mut AppContext) { + cx.set_global(SettingsStore::test(cx)); + language_settings::init(cx); + let text = indoc! {" + struct X { + a: usize, + b: usize, + } + + impl X { + + fn new() -> Self { + let a = 1; + let b = 2; + Self { a, b } + } + + pub fn a(&self, param: bool) -> usize { + self.a + } + + pub fn b(&self) -> usize { + self.b + } + } + "}; + let buffer = + cx.add_model(|cx| Buffer::new(0, 0, text).with_language(Arc::new(rust_lang()), cx)); + let snapshot = buffer.read(cx).snapshot(); + + let outline = outline_for_prompt( + &snapshot, + snapshot.anchor_before(Point::new(1, 4))..snapshot.anchor_before(Point::new(1, 4)), + cx, + ); + assert_eq!( + outline.as_deref(), + Some(indoc! {" + struct X + <|START|>a: usize + b + impl X + fn new + fn a + fn b + "}) + ); + + let outline = outline_for_prompt( + &snapshot, + snapshot.anchor_before(Point::new(8, 12))..snapshot.anchor_before(Point::new(8, 14)), + cx, + ); + assert_eq!( + outline.as_deref(), + Some(indoc! {" + struct X + a + b + impl X + fn new() -> Self { + let <|START|a |END|>= 1; + let b = 2; + Self { a, b } + } + fn a + fn b + "}) + ); + + let outline = outline_for_prompt( + &snapshot, + snapshot.anchor_before(Point::new(6, 0))..snapshot.anchor_before(Point::new(6, 0)), + cx, + ); + assert_eq!( + outline.as_deref(), + Some(indoc! {" + struct X + a + b + impl X + <|START|> + fn new + fn a + fn b + "}) + ); + + let outline = outline_for_prompt( + &snapshot, + snapshot.anchor_before(Point::new(8, 12))..snapshot.anchor_before(Point::new(13, 9)), + cx, + ); + assert_eq!( + outline.as_deref(), + Some(indoc! {" + struct X + a + b + impl X + fn new() -> Self { + let <|START|a = 1; + let b = 2; + Self { a, b } + } + + pub f|END|>n a(&self, param: bool) -> usize { + self.a + } + fn b + "}) + ); + + let outline = outline_for_prompt( + &snapshot, + snapshot.anchor_before(Point::new(5, 6))..snapshot.anchor_before(Point::new(12, 0)), + cx, + ); + assert_eq!( + outline.as_deref(), + Some(indoc! {" + struct X + a + b + impl X<|START| { + + fn new() -> Self { + let a = 1; + let b = 2; + Self { a, b } + } + |END|> + fn a + fn b + "}) + ); + + let outline = outline_for_prompt( + &snapshot, + snapshot.anchor_before(Point::new(18, 8))..snapshot.anchor_before(Point::new(18, 8)), + cx, + ); + assert_eq!( + outline.as_deref(), + Some(indoc! {" + struct X + a + b + impl X + fn new + fn a + pub fn b(&self) -> usize { + <|START|>self.b + } + "}) + ); + } +} From 53c25690f940e396b99d004ecefd353c15e1aa8f Mon Sep 17 00:00:00 2001 From: Antonio Scandurra Date: Fri, 29 Sep 2023 20:37:07 +0200 Subject: [PATCH 2/8] WIP: Use a different approach to codegen outline --- crates/zed/src/languages/rust/summary.scm | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 crates/zed/src/languages/rust/summary.scm diff --git a/crates/zed/src/languages/rust/summary.scm b/crates/zed/src/languages/rust/summary.scm new file mode 100644 index 0000000000000000000000000000000000000000..7174eec3c384336d9c2d647428d60d0eb82dea2a --- /dev/null +++ b/crates/zed/src/languages/rust/summary.scm @@ -0,0 +1,6 @@ +(function_item + body: (block + "{" @keep + "}" @keep) @collapse) + +(use_declaration) @collapse From 64a55681e615d71ae5146d16b1d6f28d2237819c Mon Sep 17 00:00:00 2001 From: Antonio Scandurra Date: Mon, 2 Oct 2023 14:32:13 +0200 Subject: [PATCH 3/8] Summarize the contents of a file using the embedding query --- crates/assistant/src/assistant_panel.rs | 1 - crates/assistant/src/prompts.rs | 458 ++++++++++++---------- crates/language/src/buffer.rs | 12 +- crates/zed/src/languages/rust/summary.scm | 6 - 4 files changed, 253 insertions(+), 224 deletions(-) delete mode 100644 crates/zed/src/languages/rust/summary.scm diff --git a/crates/assistant/src/assistant_panel.rs b/crates/assistant/src/assistant_panel.rs index 37d0d729fe64e0def057402fb9a24b796ebdf317..816047e325c9d5377a52c380d3b1c92d3b3a983f 100644 --- a/crates/assistant/src/assistant_panel.rs +++ b/crates/assistant/src/assistant_panel.rs @@ -578,7 +578,6 @@ impl AssistantPanel { language_name, &snapshot, language_range, - cx, codegen_kind, ); diff --git a/crates/assistant/src/prompts.rs b/crates/assistant/src/prompts.rs index ee58090d04dbff9d7da058a905c9e52b8fe1b0cd..8699c77cd13b31791a10b1fb54fcf322ed25240f 100644 --- a/crates/assistant/src/prompts.rs +++ b/crates/assistant/src/prompts.rs @@ -1,86 +1,118 @@ -use gpui::AppContext; +use crate::codegen::CodegenKind; use language::{BufferSnapshot, OffsetRangeExt, ToOffset}; use std::cmp; use std::ops::Range; use std::{fmt::Write, iter}; -use crate::codegen::CodegenKind; +fn summarize(buffer: &BufferSnapshot, selected_range: Range) -> String { + #[derive(Debug)] + struct Match { + collapse: Range, + keep: Vec>, + } -fn outline_for_prompt( - buffer: &BufferSnapshot, - range: Range, - cx: &AppContext, -) -> Option { - let indent = buffer - .language_indent_size_at(0, cx) - .chars() - .collect::(); - let outline = buffer.outline(None)?; - let range = range.to_offset(buffer); - - let mut text = String::new(); - let mut items = outline.items.into_iter().peekable(); - - let mut intersected = false; - let mut intersection_indent = 0; - let mut extended_range = range.clone(); - - while let Some(item) = items.next() { - let item_range = item.range.to_offset(buffer); - if item_range.end < range.start || item_range.start > range.end { - text.extend(iter::repeat(indent.as_str()).take(item.depth)); - text.push_str(&item.text); - text.push('\n'); - } else { - intersected = true; - let is_terminal = items - .peek() - .map_or(true, |next_item| next_item.depth <= item.depth); - if is_terminal { - if item_range.start <= extended_range.start { - extended_range.start = item_range.start; - intersection_indent = item.depth; + let selected_range = selected_range.to_offset(buffer); + let mut matches = buffer.matches(0..buffer.len(), |grammar| { + Some(&grammar.embedding_config.as_ref()?.query) + }); + let configs = matches + .grammars() + .iter() + .map(|g| g.embedding_config.as_ref().unwrap()) + .collect::>(); + let mut matches = iter::from_fn(move || { + while let Some(mat) = matches.peek() { + let config = &configs[mat.grammar_index]; + if let Some(collapse) = mat.captures.iter().find_map(|cap| { + if Some(cap.index) == config.collapse_capture_ix { + Some(cap.node.byte_range()) + } else { + None } - extended_range.end = cmp::max(extended_range.end, item_range.end); + }) { + let mut keep = Vec::new(); + for capture in mat.captures.iter() { + if Some(capture.index) == config.keep_capture_ix { + keep.push(capture.node.byte_range()); + } else { + continue; + } + } + matches.advance(); + return Some(Match { collapse, keep }); + } else { + matches.advance(); + } + } + None + }) + .peekable(); + + let mut summary = String::new(); + let mut offset = 0; + let mut flushed_selection = false; + while let Some(mut mat) = matches.next() { + // Keep extending the collapsed range if the next match surrounds + // the current one. + while let Some(next_mat) = matches.peek() { + if next_mat.collapse.start <= mat.collapse.start + && next_mat.collapse.end >= mat.collapse.end + { + mat = matches.next().unwrap(); } else { - let name_start = item_range.start + item.name_ranges.first().unwrap().start; - let name_end = item_range.start + item.name_ranges.last().unwrap().end; + break; + } + } + + if offset >= mat.collapse.start { + // Skip collapsed nodes that have already been summarized. + offset = cmp::max(offset, mat.collapse.end); + continue; + } - if range.start > name_end { - text.extend(iter::repeat(indent.as_str()).take(item.depth)); - text.push_str(&item.text); - text.push('\n'); + if offset <= selected_range.start && selected_range.start <= mat.collapse.end { + if !flushed_selection { + // The collapsed node ends after the selection starts, so we'll flush the selection first. + summary.extend(buffer.text_for_range(offset..selected_range.start)); + summary.push_str("<|START|"); + if selected_range.end == selected_range.start { + summary.push_str(">"); } else { - if name_start <= extended_range.start { - extended_range.start = item_range.start; - intersection_indent = item.depth; - } - extended_range.end = cmp::max(extended_range.end, name_end); + summary.extend(buffer.text_for_range(selected_range.clone())); + summary.push_str("|END|>"); } + offset = selected_range.end; + flushed_selection = true; } - } - if intersected - && items.peek().map_or(true, |next_item| { - next_item.range.start.to_offset(buffer) > range.end - }) - { - intersected = false; - text.extend(iter::repeat(indent.as_str()).take(intersection_indent)); - text.extend(buffer.text_for_range(extended_range.start..range.start)); - text.push_str("<|START|"); - text.extend(buffer.text_for_range(range.clone())); - if range.start != range.end { - text.push_str("|END|>"); - } else { - text.push_str(">"); + // If the selection intersects the collapsed node, we won't collapse it. + if selected_range.end >= mat.collapse.start { + continue; } - text.extend(buffer.text_for_range(range.end..extended_range.end)); - text.push('\n'); } + + summary.extend(buffer.text_for_range(offset..mat.collapse.start)); + for keep in mat.keep { + summary.extend(buffer.text_for_range(keep)); + } + offset = mat.collapse.end; + } + + // Flush selection if we haven't already done so. + if !flushed_selection && offset <= selected_range.start { + summary.extend(buffer.text_for_range(offset..selected_range.start)); + summary.push_str("<|START|"); + if selected_range.end == selected_range.start { + summary.push_str(">"); + } else { + summary.extend(buffer.text_for_range(selected_range.clone())); + summary.push_str("|END|>"); + } + offset = selected_range.end; } - Some(text) + summary.extend(buffer.text_for_range(offset..buffer.len())); + summary } pub fn generate_content_prompt( @@ -88,7 +120,6 @@ pub fn generate_content_prompt( language_name: Option<&str>, buffer: &BufferSnapshot, range: Range, - cx: &AppContext, kind: CodegenKind, ) -> String { let mut prompt = String::new(); @@ -100,19 +131,17 @@ pub fn generate_content_prompt( writeln!(prompt, "You're an expert engineer.\n").unwrap(); } - let outline = outline_for_prompt(buffer, range.clone(), cx); - if let Some(outline) = outline { - writeln!( - prompt, - "The file you are currently working on has the following outline:" - ) - .unwrap(); - if let Some(language_name) = language_name { - let language_name = language_name.to_lowercase(); - writeln!(prompt, "```{language_name}\n{outline}\n```").unwrap(); - } else { - writeln!(prompt, "```\n{outline}\n```").unwrap(); - } + let outline = summarize(buffer, range.clone()); + writeln!( + prompt, + "The file you are currently working on has the following outline:" + ) + .unwrap(); + if let Some(language_name) = language_name { + let language_name = language_name.to_lowercase(); + writeln!(prompt, "```{language_name}\n{outline}\n```").unwrap(); + } else { + writeln!(prompt, "```\n{outline}\n```").unwrap(); } // Assume for now that we are just generating @@ -183,39 +212,37 @@ pub(crate) mod tests { }, Some(tree_sitter_rust::language()), ) - .with_indents_query( + .with_embedding_query( r#" - (call_expression) @indent - (field_expression) @indent - (_ "(" ")" @end) @indent - (_ "{" "}" @end) @indent - "#, - ) - .unwrap() - .with_outline_query( - r#" - (struct_item - "struct" @context - name: (_) @name) @item - (enum_item - "enum" @context - name: (_) @name) @item - (enum_variant - name: (_) @name) @item - (field_declaration - name: (_) @name) @item - (impl_item - "impl" @context - trait: (_)? @name - "for"? @context - type: (_) @name) @item - (function_item - "fn" @context - name: (_) @name) @item - (mod_item - "mod" @context - name: (_) @name) @item - "#, + ( + [(line_comment) (attribute_item)]* @context + . + [ + (struct_item + name: (_) @name) + + (enum_item + name: (_) @name) + + (impl_item + trait: (_)? @name + "for"? @name + type: (_) @name) + + (trait_item + name: (_) @name) + + (function_item + name: (_) @name + body: (block + "{" @keep + "}" @keep) @collapse) + + (macro_definition + name: (_) @name) + ] @item + ) + "#, ) .unwrap() } @@ -251,132 +278,133 @@ pub(crate) mod tests { cx.add_model(|cx| Buffer::new(0, 0, text).with_language(Arc::new(rust_lang()), cx)); let snapshot = buffer.read(cx).snapshot(); - let outline = outline_for_prompt( - &snapshot, - snapshot.anchor_before(Point::new(1, 4))..snapshot.anchor_before(Point::new(1, 4)), - cx, - ); assert_eq!( - outline.as_deref(), - Some(indoc! {" - struct X - <|START|>a: usize - b - impl X - fn new - fn a - fn b - "}) - ); + summarize(&snapshot, Point::new(1, 4)..Point::new(1, 4)), + indoc! {" + struct X { + <|START|>a: usize, + b: usize, + } + + impl X { + + fn new() -> Self {} - let outline = outline_for_prompt( - &snapshot, - snapshot.anchor_before(Point::new(8, 12))..snapshot.anchor_before(Point::new(8, 14)), - cx, + pub fn a(&self, param: bool) -> usize {} + + pub fn b(&self) -> usize {} + } + "} ); + assert_eq!( - outline.as_deref(), - Some(indoc! {" - struct X - a - b - impl X + summarize(&snapshot, Point::new(8, 12)..Point::new(8, 14)), + indoc! {" + struct X { + a: usize, + b: usize, + } + + impl X { + fn new() -> Self { let <|START|a |END|>= 1; let b = 2; Self { a, b } } - fn a - fn b - "}) - ); - let outline = outline_for_prompt( - &snapshot, - snapshot.anchor_before(Point::new(6, 0))..snapshot.anchor_before(Point::new(6, 0)), - cx, + pub fn a(&self, param: bool) -> usize {} + + pub fn b(&self) -> usize {} + } + "} ); + assert_eq!( - outline.as_deref(), - Some(indoc! {" - struct X - a - b - impl X + summarize(&snapshot, Point::new(6, 0)..Point::new(6, 0)), + indoc! {" + struct X { + a: usize, + b: usize, + } + + impl X { <|START|> - fn new - fn a - fn b - "}) - ); + fn new() -> Self {} + + pub fn a(&self, param: bool) -> usize {} - let outline = outline_for_prompt( - &snapshot, - snapshot.anchor_before(Point::new(8, 12))..snapshot.anchor_before(Point::new(13, 9)), - cx, + pub fn b(&self) -> usize {} + } + "} ); + assert_eq!( - outline.as_deref(), - Some(indoc! {" - struct X - a - b - impl X - fn new() -> Self { - let <|START|a = 1; - let b = 2; - Self { a, b } - } + summarize(&snapshot, Point::new(21, 0)..Point::new(21, 0)), + indoc! {" + struct X { + a: usize, + b: usize, + } - pub f|END|>n a(&self, param: bool) -> usize { - self.a - } - fn b - "}) - ); + impl X { + + fn new() -> Self {} + + pub fn a(&self, param: bool) -> usize {} - let outline = outline_for_prompt( - &snapshot, - snapshot.anchor_before(Point::new(5, 6))..snapshot.anchor_before(Point::new(12, 0)), - cx, + pub fn b(&self) -> usize {} + } + <|START|>"} ); - assert_eq!( - outline.as_deref(), - Some(indoc! {" - struct X - a - b - impl X<|START| { - fn new() -> Self { - let a = 1; - let b = 2; - Self { a, b } + // Ensure nested functions get collapsed properly. + let text = indoc! {" + struct X { + a: usize, + b: usize, + } + + impl X { + + fn new() -> Self { + let a = 1; + let b = 2; + Self { a, b } + } + + pub fn a(&self, param: bool) -> usize { + let a = 30; + fn nested() -> usize { + 3 } - |END|> - fn a - fn b - "}) - ); + self.a + nested() + } - let outline = outline_for_prompt( - &snapshot, - snapshot.anchor_before(Point::new(18, 8))..snapshot.anchor_before(Point::new(18, 8)), - cx, - ); + pub fn b(&self) -> usize { + self.b + } + } + "}; + buffer.update(cx, |buffer, cx| buffer.set_text(text, cx)); + let snapshot = buffer.read(cx).snapshot(); assert_eq!( - outline.as_deref(), - Some(indoc! {" - struct X - a - b - impl X - fn new - fn a - pub fn b(&self) -> usize { - <|START|>self.b - } - "}) + summarize(&snapshot, Point::new(0, 0)..Point::new(0, 0)), + indoc! {" + <|START|>struct X { + a: usize, + b: usize, + } + + impl X { + + fn new() -> Self {} + + pub fn a(&self, param: bool) -> usize {} + + pub fn b(&self) -> usize {} + } + "} ); } } diff --git a/crates/language/src/buffer.rs b/crates/language/src/buffer.rs index 38b2842c127f2a6ade0d43787a24a0c76ff13374..27b01543e1e3f04f9914b1da5c530ddd26a555c1 100644 --- a/crates/language/src/buffer.rs +++ b/crates/language/src/buffer.rs @@ -8,8 +8,8 @@ use crate::{ language_settings::{language_settings, LanguageSettings}, outline::OutlineItem, syntax_map::{ - SyntaxLayerInfo, SyntaxMap, SyntaxMapCapture, SyntaxMapCaptures, SyntaxSnapshot, - ToTreeSitterPoint, + SyntaxLayerInfo, SyntaxMap, SyntaxMapCapture, SyntaxMapCaptures, SyntaxMapMatches, + SyntaxSnapshot, ToTreeSitterPoint, }, CodeLabel, LanguageScope, Outline, }; @@ -2467,6 +2467,14 @@ impl BufferSnapshot { Some(items) } + pub fn matches( + &self, + range: Range, + query: fn(&Grammar) -> Option<&tree_sitter::Query>, + ) -> SyntaxMapMatches { + self.syntax.matches(range, self, query) + } + /// Returns bracket range pairs overlapping or adjacent to `range` pub fn bracket_ranges<'a, T: ToOffset>( &'a self, diff --git a/crates/zed/src/languages/rust/summary.scm b/crates/zed/src/languages/rust/summary.scm deleted file mode 100644 index 7174eec3c384336d9c2d647428d60d0eb82dea2a..0000000000000000000000000000000000000000 --- a/crates/zed/src/languages/rust/summary.scm +++ /dev/null @@ -1,6 +0,0 @@ -(function_item - body: (block - "{" @keep - "}" @keep) @collapse) - -(use_declaration) @collapse From df7ac9b815423c8e9d8afdf5830a55177318a20a Mon Sep 17 00:00:00 2001 From: Antonio Scandurra Date: Mon, 2 Oct 2023 14:36:16 +0200 Subject: [PATCH 4/8] :lipstick: --- crates/assistant/src/assistant_panel.rs | 2 -- crates/assistant/src/prompts.rs | 9 ++------- 2 files changed, 2 insertions(+), 9 deletions(-) diff --git a/crates/assistant/src/assistant_panel.rs b/crates/assistant/src/assistant_panel.rs index 816047e325c9d5377a52c380d3b1c92d3b3a983f..4a4dc087904aee4314fa05f796ec9c5d17b64b1d 100644 --- a/crates/assistant/src/assistant_panel.rs +++ b/crates/assistant/src/assistant_panel.rs @@ -581,8 +581,6 @@ impl AssistantPanel { codegen_kind, ); - dbg!(&prompt); - let mut messages = Vec::new(); let mut model = settings::get::(cx) .default_open_ai_model diff --git a/crates/assistant/src/prompts.rs b/crates/assistant/src/prompts.rs index 8699c77cd13b31791a10b1fb54fcf322ed25240f..0646534e011dd5a457ba9765c343a44ad5ffd782 100644 --- a/crates/assistant/src/prompts.rs +++ b/crates/assistant/src/prompts.rs @@ -144,15 +144,9 @@ pub fn generate_content_prompt( writeln!(prompt, "```\n{outline}\n```").unwrap(); } - // Assume for now that we are just generating - if range.clone().start == range.end { - writeln!(prompt, "In particular, the user's cursor is current on the '<|START|>' span in the above outline, with no text selected.").unwrap(); - } else { - writeln!(prompt, "In particular, the user has selected a section of the text between the '<|START|' and '|END|>' spans.").unwrap(); - } - match kind { CodegenKind::Generate { position: _ } => { + writeln!(prompt, "In particular, the user's cursor is current on the '<|START|>' span in the above outline, with no text selected.").unwrap(); writeln!( prompt, "Assume the cursor is located where the `<|START|` marker is." @@ -170,6 +164,7 @@ pub fn generate_content_prompt( .unwrap(); } CodegenKind::Transform { range: _ } => { + writeln!(prompt, "In particular, the user has selected a section of the text between the '<|START|' and '|END|>' spans.").unwrap(); writeln!( prompt, "Modify the users code selected text based upon the users prompt: {user_prompt}" From f52200a340649eac2e65895f01f83891e32891a4 Mon Sep 17 00:00:00 2001 From: Antonio Scandurra Date: Mon, 2 Oct 2023 15:21:58 +0200 Subject: [PATCH 5/8] Prevent deploying the inline assistant when selection spans multiple excerpts --- crates/assistant/src/assistant_panel.rs | 40 ++++++++++++++----------- crates/assistant/src/prompts.rs | 4 +-- 2 files changed, 24 insertions(+), 20 deletions(-) diff --git a/crates/assistant/src/assistant_panel.rs b/crates/assistant/src/assistant_panel.rs index 4a4dc087904aee4314fa05f796ec9c5d17b64b1d..0d9f69011e3d0d53dbeca52d3902728fcf582079 100644 --- a/crates/assistant/src/assistant_panel.rs +++ b/crates/assistant/src/assistant_panel.rs @@ -274,13 +274,17 @@ impl AssistantPanel { return; }; + let selection = editor.read(cx).selections.newest_anchor().clone(); + if selection.start.excerpt_id() != selection.end.excerpt_id() { + return; + } + let inline_assist_id = post_inc(&mut self.next_inline_assist_id); let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx); let provider = Arc::new(OpenAICompletionProvider::new( api_key, cx.background().clone(), )); - let selection = editor.read(cx).selections.newest_anchor().clone(); let codegen_kind = if editor.read(cx).selections.newest::(cx).is_empty() { CodegenKind::Generate { position: selection.start, @@ -542,25 +546,25 @@ impl AssistantPanel { self.inline_prompt_history.pop_front(); } - let multi_buffer = editor.read(cx).buffer().read(cx); - let multi_buffer_snapshot = multi_buffer.snapshot(cx); - let snapshot = if multi_buffer.is_singleton() { - multi_buffer.as_singleton().unwrap().read(cx).snapshot() + let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx); + let range = pending_assist.codegen.read(cx).range(); + let start = snapshot.point_to_buffer_offset(range.start); + let end = snapshot.point_to_buffer_offset(range.end); + let (buffer, range) = if let Some((start, end)) = start.zip(end) { + let (start_buffer, start_buffer_offset) = start; + let (end_buffer, end_buffer_offset) = end; + if start_buffer.remote_id() == end_buffer.remote_id() { + (start_buffer, start_buffer_offset..end_buffer_offset) + } else { + self.finish_inline_assist(inline_assist_id, false, cx); + return; + } } else { + self.finish_inline_assist(inline_assist_id, false, cx); return; }; - let range = pending_assist.codegen.read(cx).range(); - let language_range = snapshot.anchor_at( - range.start.to_offset(&multi_buffer_snapshot), - language::Bias::Left, - ) - ..snapshot.anchor_at( - range.end.to_offset(&multi_buffer_snapshot), - language::Bias::Right, - ); - - let language = snapshot.language_at(language_range.start); + let language = buffer.language_at(range.start); let language_name = if let Some(language) = language.as_ref() { if Arc::ptr_eq(language, &language::PLAIN_TEXT) { None @@ -576,8 +580,8 @@ impl AssistantPanel { let prompt = generate_content_prompt( user_prompt.to_string(), language_name, - &snapshot, - language_range, + &buffer, + range, codegen_kind, ); diff --git a/crates/assistant/src/prompts.rs b/crates/assistant/src/prompts.rs index 0646534e011dd5a457ba9765c343a44ad5ffd782..2451369a184b18c312efa1a829b9c37c40e06579 100644 --- a/crates/assistant/src/prompts.rs +++ b/crates/assistant/src/prompts.rs @@ -119,7 +119,7 @@ pub fn generate_content_prompt( user_prompt: String, language_name: Option<&str>, buffer: &BufferSnapshot, - range: Range, + range: Range, kind: CodegenKind, ) -> String { let mut prompt = String::new(); @@ -131,7 +131,7 @@ pub fn generate_content_prompt( writeln!(prompt, "You're an expert engineer.\n").unwrap(); } - let outline = summarize(buffer, range.clone()); + let outline = summarize(buffer, range); writeln!( prompt, "The file you are currently working on has the following outline:" From d70014cfd065dcc65823ab8c0c465edf5dff5d08 Mon Sep 17 00:00:00 2001 From: Antonio Scandurra Date: Mon, 2 Oct 2023 15:36:10 +0200 Subject: [PATCH 6/8] Summarize file in the background --- crates/assistant/src/assistant_panel.rs | 48 ++++++++++++------------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/crates/assistant/src/assistant_panel.rs b/crates/assistant/src/assistant_panel.rs index 0d9f69011e3d0d53dbeca52d3902728fcf582079..b69c12a2a328ed8643315f091be11d764dcdc00d 100644 --- a/crates/assistant/src/assistant_panel.rs +++ b/crates/assistant/src/assistant_panel.rs @@ -546,15 +546,16 @@ impl AssistantPanel { self.inline_prompt_history.pop_front(); } + let codegen = pending_assist.codegen.clone(); let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx); - let range = pending_assist.codegen.read(cx).range(); + let range = codegen.read(cx).range(); let start = snapshot.point_to_buffer_offset(range.start); let end = snapshot.point_to_buffer_offset(range.end); let (buffer, range) = if let Some((start, end)) = start.zip(end) { let (start_buffer, start_buffer_offset) = start; let (end_buffer, end_buffer_offset) = end; if start_buffer.remote_id() == end_buffer.remote_id() { - (start_buffer, start_buffer_offset..end_buffer_offset) + (start_buffer.clone(), start_buffer_offset..end_buffer_offset) } else { self.finish_inline_assist(inline_assist_id, false, cx); return; @@ -574,17 +575,13 @@ impl AssistantPanel { } else { None }; - let language_name = language_name.as_deref(); - - let codegen_kind = pending_assist.codegen.read(cx).kind().clone(); - let prompt = generate_content_prompt( - user_prompt.to_string(), - language_name, - &buffer, - range, - codegen_kind, - ); + let codegen_kind = codegen.read(cx).kind().clone(); + let user_prompt = user_prompt.to_string(); + let prompt = cx.background().spawn(async move { + let language_name = language_name.as_deref(); + generate_content_prompt(user_prompt, language_name, &buffer, range, codegen_kind) + }); let mut messages = Vec::new(); let mut model = settings::get::(cx) .default_open_ai_model @@ -600,18 +597,21 @@ impl AssistantPanel { model = conversation.model.clone(); } - messages.push(RequestMessage { - role: Role::User, - content: prompt, - }); - let request = OpenAIRequest { - model: model.full_name().into(), - messages, - stream: true, - }; - pending_assist - .codegen - .update(cx, |codegen, cx| codegen.start(request, cx)); + cx.spawn(|_, mut cx| async move { + let prompt = prompt.await; + + messages.push(RequestMessage { + role: Role::User, + content: prompt, + }); + let request = OpenAIRequest { + model: model.full_name().into(), + messages, + stream: true, + }; + codegen.update(&mut cx, |codegen, cx| codegen.start(request, cx)); + }) + .detach(); } fn update_highlights_for_editor( From bf5d9e32240e5752630988fc99df5f7c82031660 Mon Sep 17 00:00:00 2001 From: Antonio Scandurra Date: Mon, 2 Oct 2023 17:50:52 +0200 Subject: [PATCH 7/8] Sort matches before processing them --- crates/assistant/src/prompts.rs | 63 ++++++++++----------- crates/zed/src/languages/rust/embedding.scm | 3 + 2 files changed, 34 insertions(+), 32 deletions(-) diff --git a/crates/assistant/src/prompts.rs b/crates/assistant/src/prompts.rs index 2451369a184b18c312efa1a829b9c37c40e06579..bf041dff523d57d62cfbc3f312a350ad4766d160 100644 --- a/crates/assistant/src/prompts.rs +++ b/crates/assistant/src/prompts.rs @@ -1,8 +1,8 @@ use crate::codegen::CodegenKind; use language::{BufferSnapshot, OffsetRangeExt, ToOffset}; -use std::cmp; +use std::cmp::{self, Reverse}; +use std::fmt::Write; use std::ops::Range; -use std::{fmt::Write, iter}; fn summarize(buffer: &BufferSnapshot, selected_range: Range) -> String { #[derive(Debug)] @@ -12,59 +12,58 @@ fn summarize(buffer: &BufferSnapshot, selected_range: Range) -> S } let selected_range = selected_range.to_offset(buffer); - let mut matches = buffer.matches(0..buffer.len(), |grammar| { + let mut ts_matches = buffer.matches(0..buffer.len(), |grammar| { Some(&grammar.embedding_config.as_ref()?.query) }); - let configs = matches + let configs = ts_matches .grammars() .iter() .map(|g| g.embedding_config.as_ref().unwrap()) .collect::>(); - let mut matches = iter::from_fn(move || { - while let Some(mat) = matches.peek() { - let config = &configs[mat.grammar_index]; - if let Some(collapse) = mat.captures.iter().find_map(|cap| { - if Some(cap.index) == config.collapse_capture_ix { - Some(cap.node.byte_range()) + let mut matches = Vec::new(); + while let Some(mat) = ts_matches.peek() { + let config = &configs[mat.grammar_index]; + if let Some(collapse) = mat.captures.iter().find_map(|cap| { + if Some(cap.index) == config.collapse_capture_ix { + Some(cap.node.byte_range()) + } else { + None + } + }) { + let mut keep = Vec::new(); + for capture in mat.captures.iter() { + if Some(capture.index) == config.keep_capture_ix { + keep.push(capture.node.byte_range()); } else { - None - } - }) { - let mut keep = Vec::new(); - for capture in mat.captures.iter() { - if Some(capture.index) == config.keep_capture_ix { - keep.push(capture.node.byte_range()); - } else { - continue; - } + continue; } - matches.advance(); - return Some(Match { collapse, keep }); - } else { - matches.advance(); } + ts_matches.advance(); + matches.push(Match { collapse, keep }); + } else { + ts_matches.advance(); } - None - }) - .peekable(); + } + matches.sort_unstable_by_key(|mat| (mat.collapse.start, Reverse(mat.collapse.end))); + let mut matches = matches.into_iter().peekable(); let mut summary = String::new(); let mut offset = 0; let mut flushed_selection = false; - while let Some(mut mat) = matches.next() { + while let Some(mat) = matches.next() { // Keep extending the collapsed range if the next match surrounds // the current one. while let Some(next_mat) = matches.peek() { - if next_mat.collapse.start <= mat.collapse.start - && next_mat.collapse.end >= mat.collapse.end + if mat.collapse.start <= next_mat.collapse.start + && mat.collapse.end >= next_mat.collapse.end { - mat = matches.next().unwrap(); + matches.next().unwrap(); } else { break; } } - if offset >= mat.collapse.start { + if offset > mat.collapse.start { // Skip collapsed nodes that have already been summarized. offset = cmp::max(offset, mat.collapse.end); continue; diff --git a/crates/zed/src/languages/rust/embedding.scm b/crates/zed/src/languages/rust/embedding.scm index e4218382a9b1ceb7e087b0d9247d5a4e66b77236..c4ed7d20976fb9c56f39aec1c8a32bba5f405f15 100644 --- a/crates/zed/src/languages/rust/embedding.scm +++ b/crates/zed/src/languages/rust/embedding.scm @@ -2,6 +2,9 @@ [(line_comment) (attribute_item)]* @context . [ + (attribute_item) @collapse + (use_declaration) @collapse + (struct_item name: (_) @name) From 9f160537ef4ea8fec1a82c45c7c70e62973b24f3 Mon Sep 17 00:00:00 2001 From: KCaverly Date: Tue, 3 Oct 2023 11:56:45 +0300 Subject: [PATCH 8/8] move collapsed only matches outside item parent in embedding.scm --- .../semantic_index/src/semantic_index_tests.rs | 17 +++++++++++++++++ crates/zed/src/languages/rust/embedding.scm | 5 +++-- 2 files changed, 20 insertions(+), 2 deletions(-) diff --git a/crates/semantic_index/src/semantic_index_tests.rs b/crates/semantic_index/src/semantic_index_tests.rs index f2cae8a55701e910379d672225c19ac18a489897..182010ca8339e9cc8ec1ff06ac31741eb4fb78ae 100644 --- a/crates/semantic_index/src/semantic_index_tests.rs +++ b/crates/semantic_index/src/semantic_index_tests.rs @@ -305,6 +305,11 @@ async fn test_code_context_retrieval_rust() { todo!(); } } + + #[derive(Clone)] + struct D { + name: String + } " .unindent(); @@ -361,6 +366,15 @@ async fn test_code_context_retrieval_rust() { .unindent(), text.find("fn function_2").unwrap(), ), + ( + " + #[derive(Clone)] + struct D { + name: String + }" + .unindent(), + text.find("struct D").unwrap(), + ), ], ); } @@ -1422,6 +1436,9 @@ fn rust_lang() -> Arc { name: (_) @name) ] @item ) + + (attribute_item) @collapse + (use_declaration) @collapse "#, ) .unwrap(), diff --git a/crates/zed/src/languages/rust/embedding.scm b/crates/zed/src/languages/rust/embedding.scm index c4ed7d20976fb9c56f39aec1c8a32bba5f405f15..286b1d13571ad62964e3f38415fc4cbbb04e4e99 100644 --- a/crates/zed/src/languages/rust/embedding.scm +++ b/crates/zed/src/languages/rust/embedding.scm @@ -2,8 +2,6 @@ [(line_comment) (attribute_item)]* @context . [ - (attribute_item) @collapse - (use_declaration) @collapse (struct_item name: (_) @name) @@ -29,3 +27,6 @@ name: (_) @name) ] @item ) + +(attribute_item) @collapse +(use_declaration) @collapse