Automatically expand context for inline assistant to nearest block (#27282)

Conrad Irwin and andrew j created 9 months ago

Release Notes:

- Inline assistant will now expand empty selections to the block under
the cursor.

Co-authored-by: andrew j <[e-mail]>

Change summary

crates/assistant/src/inline_assistant.rs  | 35 +++++++++++++++++++++++-
crates/assistant2/src/inline_assistant.rs | 35 +++++++++++++++++++++++-
crates/language/src/language.rs           |  6 ++++
3 files changed, 72 insertions(+), 4 deletions(-)

Detailed changes

crates/assistant/src/inline_assistant.rs

@@ -232,7 +232,7 @@ impl InlineAssistant {
     ) {
         let (snapshot, initial_selections) = editor.update(cx, |editor, cx| {
             (
-                editor.buffer().read(cx).snapshot(cx),
+                editor.snapshot(window, cx),
                 editor.selections.all::<Point>(cx),
             )
         });
@@ -246,7 +246,37 @@ impl InlineAssistant {
                 if selection.end.column == 0 {
                     selection.end.row -= 1;
                 }
-                selection.end.column = snapshot.line_len(MultiBufferRow(selection.end.row));
+                selection.end.column = snapshot
+                    .buffer_snapshot
+                    .line_len(MultiBufferRow(selection.end.row));
+            } else if let Some(fold) =
+                snapshot.crease_for_buffer_row(MultiBufferRow(selection.end.row))
+            {
+                selection.start = fold.range().start;
+                selection.end = fold.range().end;
+                if MultiBufferRow(selection.end.row) < snapshot.buffer_snapshot.max_row() {
+                    let chars = snapshot
+                        .buffer_snapshot
+                        .chars_at(Point::new(selection.end.row + 1, 0));
+
+                    for c in chars {
+                        if c == '\n' {
+                            break;
+                        }
+                        if c.is_whitespace() {
+                            continue;
+                        }
+                        if snapshot
+                            .language_at(selection.end)
+                            .is_some_and(|language| language.config().brackets.is_closing_brace(c))
+                        {
+                            selection.end.row += 1;
+                            selection.end.column = snapshot
+                                .buffer_snapshot
+                                .line_len(MultiBufferRow(selection.end.row));
+                        }
+                    }
+                }
             }
 
             if let Some(prev_selection) = selections.last_mut() {
@@ -262,6 +292,7 @@ impl InlineAssistant {
             }
             selections.push(selection);
         }
+        let snapshot = &snapshot.buffer_snapshot;
         let newest_selection = newest_selection.unwrap();
 
         let mut codegen_ranges = Vec::new();

crates/assistant2/src/inline_assistant.rs

@@ -324,7 +324,7 @@ impl InlineAssistant {
     ) {
         let (snapshot, initial_selections) = editor.update(cx, |editor, cx| {
             (
-                editor.buffer().read(cx).snapshot(cx),
+                editor.snapshot(window, cx),
                 editor.selections.all::<Point>(cx),
             )
         });
@@ -338,7 +338,37 @@ impl InlineAssistant {
                 if selection.end.column == 0 {
                     selection.end.row -= 1;
                 }
-                selection.end.column = snapshot.line_len(MultiBufferRow(selection.end.row));
+                selection.end.column = snapshot
+                    .buffer_snapshot
+                    .line_len(MultiBufferRow(selection.end.row));
+            } else if let Some(fold) =
+                snapshot.crease_for_buffer_row(MultiBufferRow(selection.end.row))
+            {
+                selection.start = fold.range().start;
+                selection.end = fold.range().end;
+                if MultiBufferRow(selection.end.row) < snapshot.buffer_snapshot.max_row() {
+                    let chars = snapshot
+                        .buffer_snapshot
+                        .chars_at(Point::new(selection.end.row + 1, 0));
+
+                    for c in chars {
+                        if c == '\n' {
+                            break;
+                        }
+                        if c.is_whitespace() {
+                            continue;
+                        }
+                        if snapshot
+                            .language_at(selection.end)
+                            .is_some_and(|language| language.config().brackets.is_closing_brace(c))
+                        {
+                            selection.end.row += 1;
+                            selection.end.column = snapshot
+                                .buffer_snapshot
+                                .line_len(MultiBufferRow(selection.end.row));
+                        }
+                    }
+                }
             }
 
             if let Some(prev_selection) = selections.last_mut() {
@@ -354,6 +384,7 @@ impl InlineAssistant {
             }
             selections.push(selection);
         }
+        let snapshot = &snapshot.buffer_snapshot;
         let newest_selection = newest_selection.unwrap();
 
         let mut codegen_ranges = Vec::new();

crates/language/src/language.rs

@@ -887,6 +887,12 @@ pub struct BracketPairConfig {
     pub disabled_scopes_by_bracket_ix: Vec<Vec<String>>,
 }
 
+impl BracketPairConfig {
+    pub fn is_closing_brace(&self, c: char) -> bool {
+        self.pairs.iter().any(|pair| pair.end.starts_with(c))
+    }
+}
+
 fn bracket_pair_config_json_schema(gen: &mut SchemaGenerator) -> Schema {
     Option::<Vec<BracketPairContent>>::json_schema(gen)
 }