Fix more failure cases of assistant edits (#19653)

* Make `description` optional (since we describe it as optional in the
prompt, and we're currently not showing it)
* Fix fuzzy location bug that neglected the cost of deleting prefixes of
the query.
* Make auto-indent work for single-line edits. Previously, auto-indent
would not occur when overwriting a single line (without inserting or
deleting a newline)

Release Notes:

- N/A

assets/prompts/edit_workflow.hbs

@@ -88,7 +88,6 @@ origin: (f64, f64),
 
 <edit>
 <path>src/shapes/rectangle.rs</path>
-<description>Update the Rectangle's new function to take an origin parameter</description>
 <operation>update</operation>
 <old_text>
 fn new(width: f64, height: f64) -> Self {
@@ -117,7 +116,6 @@ pub struct Circle {
 
 <edit>
 <path>src/shapes/circle.rs</path>
-<description>Update the Circle's new function to take an origin parameter</description>
 <operation>update</operation>
 <old_text>
 fn new(radius: f64) -> Self {
@@ -134,7 +132,6 @@ fn new(origin: (f64, f64), radius: f64) -> Self {
 
 <edit>
 <path>src/shapes/rectangle.rs</path>
-<description>Add an import for the std::fmt module</description>
 <operation>insert_before</operation>
 <old_text>
 struct Rectangle {
@@ -147,7 +144,10 @@ use std::fmt;
 
 <edit>
 <path>src/shapes/rectangle.rs</path>
-<description>Add a Display implementation for Rectangle</description>
+<description>
+Add a manual Display implementation for Rectangle.
+Currently, this is the same as a derived Display implementation.
+</description>
 <operation>insert_after</operation>
 <old_text>
         Rectangle { width, height }
@@ -169,7 +169,6 @@ impl fmt::Display for Rectangle {
 
 <edit>
 <path>src/shapes/circle.rs</path>
-<description>Add an import for the `std::fmt` module</description>
 <operation>insert_before</operation>
 <old_text>
 struct Circle {
@@ -181,7 +180,6 @@ use std::fmt;
 
 <edit>
 <path>src/shapes/circle.rs</path>
-<description>Add a Display implementation for Circle</description>
 <operation>insert_after</operation>
 <old_text>
         Circle { radius }

crates/assistant/src/context/context_tests.rs

@@ -636,7 +636,7 @@ async fn test_workflow_step_parsing(cx: &mut TestAppContext) {
             kind: AssistantEditKind::InsertAfter {
                 old_text: "fn one".into(),
                 new_text: "fn two() {}".into(),
-                description: "add a `two` function".into(),
+                description: Some("add a `two` function".into()),
             },
         }]],
         cx,
@@ -690,7 +690,7 @@ async fn test_workflow_step_parsing(cx: &mut TestAppContext) {
             kind: AssistantEditKind::InsertAfter {
                 old_text: "fn zero".into(),
                 new_text: "fn two() {}".into(),
-                description: "add a `two` function".into(),
+                description: Some("add a `two` function".into()),
             },
         }]],
         cx,
@@ -754,7 +754,7 @@ async fn test_workflow_step_parsing(cx: &mut TestAppContext) {
             kind: AssistantEditKind::InsertAfter {
                 old_text: "fn zero".into(),
                 new_text: "fn two() {}".into(),
-                description: "add a `two` function".into(),
+                description: Some("add a `two` function".into()),
             },
         }]],
         cx,
@@ -798,7 +798,7 @@ async fn test_workflow_step_parsing(cx: &mut TestAppContext) {
             kind: AssistantEditKind::InsertAfter {
                 old_text: "fn zero".into(),
                 new_text: "fn two() {}".into(),
-                description: "add a `two` function".into(),
+                description: Some("add a `two` function".into()),
             },
         }]],
         cx,

crates/assistant/src/patch.rs

@@ -33,21 +33,21 @@ pub enum AssistantEditKind {
     Update {
         old_text: String,
         new_text: String,
-        description: String,
+        description: Option<String>,
     },
     Create {
         new_text: String,
-        description: String,
+        description: Option<String>,
     },
     InsertBefore {
         old_text: String,
         new_text: String,
-        description: String,
+        description: Option<String>,
     },
     InsertAfter {
         old_text: String,
         new_text: String,
-        description: String,
+        description: Option<String>,
     },
     Delete {
         old_text: String,
@@ -86,19 +86,37 @@ enum SearchDirection {
     Diagonal,
 }
 
-// A measure of the currently quality of an in-progress fuzzy search.
-//
-// Uses 60 bits to store a numeric cost, and 4 bits to store the preceding
-// operation in the search.
 #[derive(Copy, Clone, PartialEq, Eq, PartialOrd, Ord)]
 struct SearchState {
-    score: u32,
+    cost: u32,
     direction: SearchDirection,
 }
 
 impl SearchState {
-    fn new(score: u32, direction: SearchDirection) -> Self {
+    fn new(cost: u32, direction: SearchDirection) -> Self {
-        Self { score, direction }
+        Self { cost, direction }
+    }
+}
+
+struct SearchMatrix {
+    cols: usize,
+    data: Vec<SearchState>,
+}
+
+impl SearchMatrix {
+    fn new(rows: usize, cols: usize) -> Self {
+        SearchMatrix {
+            cols,
+            data: vec![SearchState::new(0, SearchDirection::Diagonal); rows * cols],
+        }
+    }
+
+    fn get(&self, row: usize, col: usize) -> SearchState {
+        self.data[row * self.cols + col]
+    }
+
+    fn set(&mut self, row: usize, col: usize, cost: SearchState) {
+        self.data[row * self.cols + col] = cost;
     }
 }
 
@@ -187,23 +205,23 @@ impl AssistantEdit {
             "update" => AssistantEditKind::Update {
                 old_text: old_text.ok_or_else(|| anyhow!("missing old_text"))?,
                 new_text: new_text.ok_or_else(|| anyhow!("missing new_text"))?,
-                description: description.ok_or_else(|| anyhow!("missing description"))?,
+                description,
             },
             "insert_before" => AssistantEditKind::InsertBefore {
                 old_text: old_text.ok_or_else(|| anyhow!("missing old_text"))?,
                 new_text: new_text.ok_or_else(|| anyhow!("missing new_text"))?,
-                description: description.ok_or_else(|| anyhow!("missing description"))?,
+                description,
             },
             "insert_after" => AssistantEditKind::InsertAfter {
                 old_text: old_text.ok_or_else(|| anyhow!("missing old_text"))?,
                 new_text: new_text.ok_or_else(|| anyhow!("missing new_text"))?,
-                description: description.ok_or_else(|| anyhow!("missing description"))?,
+                description,
             },
             "delete" => AssistantEditKind::Delete {
                 old_text: old_text.ok_or_else(|| anyhow!("missing old_text"))?,
             },
             "create" => AssistantEditKind::Create {
-                description: description.ok_or_else(|| anyhow!("missing description"))?,
+                description,
                 new_text: new_text.ok_or_else(|| anyhow!("missing new_text"))?,
             },
             _ => Err(anyhow!("unknown operation {operation:?}"))?,
@@ -264,7 +282,7 @@ impl AssistantEditKind {
                 ResolvedEdit {
                     range,
                     new_text,
-                    description: Some(description),
+                    description,
                 }
             }
             Self::Create {
@@ -272,7 +290,7 @@ impl AssistantEditKind {
                 description,
             } => ResolvedEdit {
                 range: text::Anchor::MIN..text::Anchor::MAX,
-                description: Some(description),
+                description,
                 new_text,
             },
             Self::InsertBefore {
@@ -285,7 +303,7 @@ impl AssistantEditKind {
                 ResolvedEdit {
                     range: range.start..range.start,
                     new_text,
-                    description: Some(description),
+                    description,
                 }
             }
             Self::InsertAfter {
@@ -298,7 +316,7 @@ impl AssistantEditKind {
                 ResolvedEdit {
                     range: range.end..range.end,
                     new_text,
-                    description: Some(description),
+                    description,
                 }
             }
             Self::Delete { old_text } => {
@@ -314,44 +332,29 @@ impl AssistantEditKind {
 
     fn resolve_location(buffer: &text::BufferSnapshot, search_query: &str) -> Range<text::Anchor> {
         const INSERTION_COST: u32 = 3;
+        const DELETION_COST: u32 = 10;
         const WHITESPACE_INSERTION_COST: u32 = 1;
-        const DELETION_COST: u32 = 3;
         const WHITESPACE_DELETION_COST: u32 = 1;
-        const EQUALITY_BONUS: u32 = 5;
-
-        struct Matrix {
-            cols: usize,
-            data: Vec<SearchState>,
-        }
-
-        impl Matrix {
-            fn new(rows: usize, cols: usize) -> Self {
-                Matrix {
-                    cols,
-                    data: vec![SearchState::new(0, SearchDirection::Diagonal); rows * cols],
-                }
-            }
-
-            fn get(&self, row: usize, col: usize) -> SearchState {
-                self.data[row * self.cols + col]
-            }
-
-            fn set(&mut self, row: usize, col: usize, cost: SearchState) {
-                self.data[row * self.cols + col] = cost;
-            }
-        }
 
         let buffer_len = buffer.len();
         let query_len = search_query.len();
-        let mut matrix = Matrix::new(query_len + 1, buffer_len + 1);
+        let mut matrix = SearchMatrix::new(query_len + 1, buffer_len + 1);
-
+        let mut leading_deletion_cost = 0_u32;
         for (row, query_byte) in search_query.bytes().enumerate() {
+            let deletion_cost = if query_byte.is_ascii_whitespace() {
+                WHITESPACE_DELETION_COST
+            } else {
+                DELETION_COST
+            };
+
+            leading_deletion_cost = leading_deletion_cost.saturating_add(deletion_cost);
+            matrix.set(
+                row + 1,
+                0,
+                SearchState::new(leading_deletion_cost, SearchDirection::Diagonal),
+            );
+
             for (col, buffer_byte) in buffer.bytes_in_range(0..buffer.len()).flatten().enumerate() {
-                let deletion_cost = if query_byte.is_ascii_whitespace() {
-                    WHITESPACE_DELETION_COST
-                } else {
-                    DELETION_COST
-                };
                 let insertion_cost = if buffer_byte.is_ascii_whitespace() {
                     WHITESPACE_INSERTION_COST
                 } else {
@@ -359,38 +362,35 @@ impl AssistantEditKind {
                 };
 
                 let up = SearchState::new(
-                    matrix.get(row, col + 1).score.saturating_sub(deletion_cost),
+                    matrix.get(row, col + 1).cost.saturating_add(deletion_cost),
                     SearchDirection::Up,
                 );
                 let left = SearchState::new(
-                    matrix
+                    matrix.get(row + 1, col).cost.saturating_add(insertion_cost),
-                        .get(row + 1, col)
-                        .score
-                        .saturating_sub(insertion_cost),
                     SearchDirection::Left,
                 );
                 let diagonal = SearchState::new(
                     if query_byte == *buffer_byte {
-                        matrix.get(row, col).score.saturating_add(EQUALITY_BONUS)
+                        matrix.get(row, col).cost
                     } else {
                         matrix
                             .get(row, col)
-                            .score
+                            .cost
-                            .saturating_sub(deletion_cost + insertion_cost)
+                            .saturating_add(deletion_cost + insertion_cost)
                     },
                     SearchDirection::Diagonal,
                 );
-                matrix.set(row + 1, col + 1, up.max(left).max(diagonal));
+                matrix.set(row + 1, col + 1, up.min(left).min(diagonal));
             }
         }
 
         // Traceback to find the best match
         let mut best_buffer_end = buffer_len;
-        let mut best_score = 0;
+        let mut best_cost = u32::MAX;
         for col in 1..=buffer_len {
-            let score = matrix.get(query_len, col).score;
+            let cost = matrix.get(query_len, col).cost;
-            if score > best_score {
+            if cost < best_cost {
-                best_score = score;
+                best_cost = cost;
                 best_buffer_end = col;
             }
         }
@@ -560,89 +560,84 @@ mod tests {
         language_settings::AllLanguageSettings, Language, LanguageConfig, LanguageMatcher,
     };
     use settings::SettingsStore;
-    use text::{OffsetRangeExt, Point};
     use ui::BorrowAppContext;
     use unindent::Unindent as _;
+    use util::test::{generate_marked_text, marked_text_ranges};
 
     #[gpui::test]
     fn test_resolve_location(cx: &mut AppContext) {
-        {
+        assert_location_resolution(
-            let buffer = cx.new_model(|cx| {
+            concat!(
-                Buffer::local(
+                "    Lorem\n",
-                    concat!(
+                "«    ipsum\n",
-                        "    Lorem\n",
+                "    dolor sit amet»\n",
-                        "    ipsum\n",
+                "    consecteur",
-                        "    dolor sit amet\n",
+            ),
-                        "    consecteur",
+            "ipsum\ndolor",
-                    ),
+            cx,
-                    cx,
+        );
-                )
-            });
-            let snapshot = buffer.read(cx).snapshot();
-            assert_eq!(
-                AssistantEditKind::resolve_location(&snapshot, "ipsum\ndolor").to_point(&snapshot),
-                Point::new(1, 0)..Point::new(2, 18)
-            );
-        }
 
-        {
+        assert_location_resolution(
-            let buffer = cx.new_model(|cx| {
+            &"
-                Buffer::local(
+            «fn foo1(a: usize) -> usize {
-                    concat!(
+                40
-                        "fn foo1(a: usize) -> usize {\n",
+            }»
-                        "    40\n",
-                        "}\n",
-                        "\n",
-                        "fn foo2(b: usize) -> usize {\n",
-                        "    42\n",
-                        "}\n",
-                    ),
-                    cx,
-                )
-            });
-            let snapshot = buffer.read(cx).snapshot();
-            assert_eq!(
-                AssistantEditKind::resolve_location(&snapshot, "fn foo1(b: usize) {\n40\n}")
-                    .to_point(&snapshot),
-                Point::new(0, 0)..Point::new(2, 1)
-            );
-        }
 
-        {
+            fn foo2(b: usize) -> usize {
-            let buffer = cx.new_model(|cx| {
+                42
-                Buffer::local(
+            }
-                    concat!(
+            "
-                        "fn main() {\n",
+            .unindent(),
-                        "    Foo\n",
+            "fn foo1(b: usize) {\n40\n}",
-                        "        .bar()\n",
+            cx,
-                        "        .baz()\n",
+        );
-                        "        .qux()\n",
+
-                        "}\n",
+        assert_location_resolution(
-                        "\n",
+            &"
-                        "fn foo2(b: usize) -> usize {\n",
+            fn main() {
-                        "    42\n",
+            «    Foo
-                        "}\n",
+                    .bar()
-                    ),
+                    .baz()
-                    cx,
+                    .qux()»
-                )
+            }
-            });
+
-            let snapshot = buffer.read(cx).snapshot();
+            fn foo2(b: usize) -> usize {
-            assert_eq!(
+                42
-                AssistantEditKind::resolve_location(&snapshot, "Foo.bar.baz.qux()")
+            }
-                    .to_point(&snapshot),
+            "
-                Point::new(1, 0)..Point::new(4, 14)
+            .unindent(),
-            );
+            "Foo.bar.baz.qux()",
-        }
+            cx,
+        );
+
+        assert_location_resolution(
+            &"
+            class Something {
+                one() { return 1; }
+            «    two() { return 2222; }
+                three() { return 333; }
+                four() { return 4444; }
+                five() { return 5555; }
+                six() { return 6666; }
+            »    seven() { return 7; }
+                eight() { return 8; }
+            }
+            "
+            .unindent(),
+            &"
+                two() { return 2222; }
+                four() { return 4444; }
+                five() { return 5555; }
+                six() { return 6666; }
+            "
+            .unindent(),
+            cx,
+        );
     }
 
     #[gpui::test]
     fn test_resolve_edits(cx: &mut AppContext) {
-        let settings_store = SettingsStore::test(cx);
+        init_test(cx);
-        cx.set_global(settings_store);
-        language::init(cx);
-        cx.update_global::<SettingsStore, _>(|settings, cx| {
-            settings.update_user_settings::<AllLanguageSettings>(cx, |_| {});
-        });
 
         assert_edits(
             "
@@ -675,7 +670,7 @@ mod tests {
                         last_name: String,
                     "
                     .unindent(),
-                    description: "".into(),
+                    description: None,
                 },
                 AssistantEditKind::Update {
                     old_text: "
@@ -690,7 +685,7 @@ mod tests {
                         }
                     "
                     .unindent(),
-                    description: "".into(),
+                    description: None,
                 },
             ],
             "
@@ -734,7 +729,7 @@ mod tests {
                             qux();
                         }"
                     .unindent(),
-                    description: "implement bar".into(),
+                    description: Some("implement bar".into()),
                 },
                 AssistantEditKind::Update {
                     old_text: "
@@ -747,7 +742,7 @@ mod tests {
                             bar();
                         }"
                     .unindent(),
-                    description: "call bar in foo".into(),
+                    description: Some("call bar in foo".into()),
                 },
                 AssistantEditKind::InsertAfter {
                     old_text: "
@@ -762,7 +757,7 @@ mod tests {
                         }
                     "
                     .unindent(),
-                    description: "implement qux".into(),
+                    description: Some("implement qux".into()),
                 },
             ],
             "
@@ -814,7 +809,7 @@ mod tests {
                         }
                     "
                     .unindent(),
-                    description: "pick better number".into(),
+                    description: None,
                 },
                 AssistantEditKind::Update {
                     old_text: "
@@ -829,7 +824,7 @@ mod tests {
                         }
                     "
                     .unindent(),
-                    description: "pick better number".into(),
+                    description: None,
                 },
                 AssistantEditKind::Update {
                     old_text: "
@@ -844,7 +839,7 @@ mod tests {
                         }
                     "
                     .unindent(),
-                    description: "pick better number".into(),
+                    description: None,
                 },
             ],
             "
@@ -865,6 +860,69 @@ mod tests {
             .unindent(),
             cx,
         );
+
+        assert_edits(
+            "
+            impl Person {
+                fn set_name(&mut self, name: String) {
+                    self.name = name;
+                }
+
+                fn name(&self) -> String {
+                    return self.name;
+                }
+            }
+            "
+            .unindent(),
+            vec![
+                AssistantEditKind::Update {
+                    old_text: "self.name = name;".unindent(),
+                    new_text: "self._name = name;".unindent(),
+                    description: None,
+                },
+                AssistantEditKind::Update {
+                    old_text: "return self.name;\n".unindent(),
+                    new_text: "return self._name;\n".unindent(),
+                    description: None,
+                },
+            ],
+            "
+                impl Person {
+                    fn set_name(&mut self, name: String) {
+                        self._name = name;
+                    }
+
+                    fn name(&self) -> String {
+                        return self._name;
+                    }
+                }
+            "
+            .unindent(),
+            cx,
+        );
+    }
+
+    fn init_test(cx: &mut AppContext) {
+        let settings_store = SettingsStore::test(cx);
+        cx.set_global(settings_store);
+        language::init(cx);
+        cx.update_global::<SettingsStore, _>(|settings, cx| {
+            settings.update_user_settings::<AllLanguageSettings>(cx, |_| {});
+        });
+    }
+
+    #[track_caller]
+    fn assert_location_resolution(
+        text_with_expected_range: &str,
+        query: &str,
+        cx: &mut AppContext,
+    ) {
+        let (text, _) = marked_text_ranges(text_with_expected_range, false);
+        let buffer = cx.new_model(|cx| Buffer::local(text.clone(), cx));
+        let snapshot = buffer.read(cx).snapshot();
+        let range = AssistantEditKind::resolve_location(&snapshot, query).to_offset(&snapshot);
+        let text_with_actual_range = generate_marked_text(&text, &[range], false);
+        pretty_assertions::assert_eq!(text_with_actual_range, text_with_expected_range);
     }
 
     #[track_caller]

crates/language/src/buffer.rs

@@ -1967,18 +1967,27 @@ impl Buffer {
                     let new_text_length = new_text.len();
                     let old_start = range.start.to_point(&before_edit);
                     let new_start = (delta + range.start as isize) as usize;
-                    delta += new_text_length as isize - (range.end as isize - range.start as isize);
+                    let range_len = range.end - range.start;
+                    delta += new_text_length as isize - range_len as isize;
 
+                    // Decide what range of the insertion to auto-indent, and whether
+                    // the first line of the insertion should be considered a newly-inserted line
+                    // or an edit to an existing line.
                     let mut range_of_insertion_to_indent = 0..new_text_length;
-                    let mut first_line_is_new = false;
+                    let mut first_line_is_new = true;
-                    let mut original_indent_column = None;
 
-                    // When inserting an entire line at the beginning of an existing line,
+                    let old_line_start = before_edit.indent_size_for_line(old_start.row).len;
-                    // treat the insertion as new.
+                    let old_line_end = before_edit.line_len(old_start.row);
-                    if new_text.contains('\n')
+
-                        && old_start.column <= before_edit.indent_size_for_line(old_start.row).len
+                    if old_start.column > old_line_start {
+                        first_line_is_new = false;
+                    }
+
+                    if !new_text.contains('\n')
+                        && (old_start.column + (range_len as u32) < old_line_end
+                            || old_line_end == old_line_start)
                     {
-                        first_line_is_new = true;
+                        first_line_is_new = false;
                     }
 
                     // When inserting text starting with a newline, avoid auto-indenting the
@@ -1988,7 +1997,7 @@ impl Buffer {
                         first_line_is_new = true;
                     }
 
-                    // Avoid auto-indenting after the insertion.
+                    let mut original_indent_column = None;
                     if let AutoindentMode::Block {
                         original_indent_columns,
                     } = &mode
@@ -2000,6 +2009,8 @@ impl Buffer {
                                 )
                                 .len
                             }));
+
+                        // Avoid auto-indenting the line after the edit.
                         if new_text[range_of_insertion_to_indent.clone()].ends_with('\n') {
                             range_of_insertion_to_indent.end -= 1;
                         }

crates/language/src/buffer_tests.rs

@@ -1241,7 +1241,6 @@ fn test_autoindent_does_not_adjust_lines_with_unchanged_suggestion(cx: &mut AppC
             Some(AutoindentMode::EachLine),
             cx,
         );
-
         assert_eq!(
             buffer.text(),
             "
@@ -1256,6 +1255,74 @@ fn test_autoindent_does_not_adjust_lines_with_unchanged_suggestion(cx: &mut AppC
             "
             .unindent()
         );
+
+        // Insert a newline after the open brace. It is auto-indented
+        buffer.edit_via_marked_text(
+            &"
+            fn a() {«
+            »
+            c
+                .f
+                .g();
+            d
+                .f
+                .g();
+            }
+            "
+            .unindent(),
+            Some(AutoindentMode::EachLine),
+            cx,
+        );
+        assert_eq!(
+            buffer.text(),
+            "
+            fn a() {
+                ˇ
+            c
+                .f
+                .g();
+            d
+                .f
+                .g();
+            }
+            "
+            .unindent()
+            .replace("ˇ", "")
+        );
+
+        // Manually outdent the line. It stays outdented.
+        buffer.edit_via_marked_text(
+            &"
+            fn a() {
+            «»
+            c
+                .f
+                .g();
+            d
+                .f
+                .g();
+            }
+            "
+            .unindent(),
+            Some(AutoindentMode::EachLine),
+            cx,
+        );
+        assert_eq!(
+            buffer.text(),
+            "
+            fn a() {
+
+            c
+                .f
+                .g();
+            d
+                .f
+                .g();
+            }
+            "
+            .unindent()
+        );
+
         buffer
     });
 

crates/languages/Cargo.toml

@@ -10,7 +10,7 @@ workspace = true
 
 [features]
 test-support = [
-    "tree-sitter"
+    "load-grammars"
 ]
 load-grammars = [
     "tree-sitter-bash",
@@ -82,3 +82,8 @@ text.workspace = true
 theme = { workspace = true, features = ["test-support"] }
 unindent.workspace = true
 workspace = { workspace = true, features = ["test-support"] }
+tree-sitter-typescript.workspace = true
+tree-sitter-python.workspace = true
+tree-sitter-go.workspace = true
+tree-sitter-c.workspace = true
+tree-sitter-css.workspace = true

crates/languages/src/lib.rs

@@ -288,15 +288,6 @@ fn load_config(name: &str) -> LanguageConfig {
         .with_context(|| format!("failed to load config.toml for language {name:?}"))
         .unwrap();
 
-    #[cfg(not(feature = "load-grammars"))]
-    {
-        config = LanguageConfig {
-            name: config.name,
-            matcher: config.matcher,
-            ..Default::default()
-        }
-    }
-
     config
 }