inline_assistant: Respect tabs when selection first row is not indented (#14886)

Artur Rodrigues created 2 years ago

When using the inline assistant with a language such as Go that uses
tabs, if the user selects a block of text that is correctly formatted
and where the first line has no indentation, the `suggested_line_indent`
variable ends up with `IndentSize { len: 0, kind: Space }`. That's
because `suggested_line_indent` current relies on
`BufferSnapshot::suggested_indents` suggestion for the first line on the
selection, but since it is already correctly indented, there are no
suggestions and `MultiBufferSnapshot::indent_size_for_line` is used
instead.


https://github.com/zed-industries/zed/blob/2d96bba61fd9e60951ecfcf697707a974475c1b2/crates/assistant/src/inline_assistant.rs#L2124-L2128

In this patch, we also take a look at the rest of the selection and
detect tabs. If one is encountered, we assume that tabs should always be
used. I suppose this isn't perfect, especially if the original file had
a mix of spaces and tabs, however it seems better than the status quo.

I considered using `BufferSnapshot::language_indent_size_at`, but I
imagine tabs should be preserved even when a specific language isn't
being used.

See screenshot below of the original prompt with this patch.

Tests:

* New unit test
* I've also manually tested with a few other cases: selection where all
lines are indented and file that only use spaces.

Release Notes:

- Fixed 'inline_assistant: tabs are overwritten with space characters
when first line in selection has no indentation'
([#14885](https://github.com/zed-industries/zed/issues/14885)).

<img width="942" alt="image"
src="https://github.com/user-attachments/assets/f2c5d7e9-e8bc-400b-bd6f-09e4a89d22c1">

Change summary

crates/assistant/src/inline_assistant.rs | 75 ++++++++++++++++++++++++-
1 file changed, 72 insertions(+), 3 deletions(-)

Detailed changes

crates/assistant/src/inline_assistant.rs

@@ -26,7 +26,7 @@ use gpui::{
     ModelContext, Subscription, Task, TextStyle, UpdateGlobal, View, ViewContext, WeakView,
     WindowContext,
 };
-use language::{Buffer, Point, Selection, TransactionId};
+use language::{Buffer, IndentKind, Point, Selection, TransactionId};
 use language_model::{LanguageModelRequest, LanguageModelRequestMessage, Role};
 use multi_buffer::MultiBufferRow;
 use parking_lot::Mutex;
@@ -2087,12 +2087,26 @@ impl Codegen {
             .collect::<Rope>();
 
         let selection_start = range.start.to_point(&snapshot);
-        let suggested_line_indent = snapshot
-            .suggested_indents(selection_start.row..selection_start.row + 1, cx)
+
+        // Start with the indentation of the first line in the selection
+        let mut suggested_line_indent = snapshot
+            .suggested_indents(selection_start.row..=selection_start.row, cx)
             .into_values()
             .next()
             .unwrap_or_else(|| snapshot.indent_size_for_line(MultiBufferRow(selection_start.row)));
 
+        // If the first line in the selection does not have indentation, check the following lines
+        if suggested_line_indent.len == 0 && suggested_line_indent.kind == IndentKind::Space {
+            for row in selection_start.row..=range.end.to_point(&snapshot).row {
+                let line_indent = snapshot.indent_size_for_line(MultiBufferRow(row));
+                // Prefer tabs if a line in the selection uses tabs as indentation
+                if line_indent.kind == IndentKind::Tab {
+                    suggested_line_indent.kind = IndentKind::Tab;
+                    break;
+                }
+            }
+        }
+
         let telemetry = self.telemetry.clone();
         self.edit_position = range.start;
         self.diff = Diff::default();
@@ -2776,6 +2790,61 @@ mod tests {
         );
     }
 
+    #[gpui::test(iterations = 10)]
+    async fn test_autoindent_respects_tabs_in_selection(cx: &mut TestAppContext) {
+        cx.set_global(cx.update(SettingsStore::test));
+        cx.update(|cx| FakeCompletionProvider::setup_test(cx));
+        cx.update(language_settings::init);
+
+        let text = indoc! {"
+            func main() {
+            \tx := 0
+            \tfor i := 0; i < 10; i++ {
+            \t\tx++
+            \t}
+            }
+        "};
+        let buffer = cx.new_model(|cx| Buffer::local(text, cx));
+        let buffer = cx.new_model(|cx| MultiBuffer::singleton(buffer, cx));
+        let range = buffer.read_with(cx, |buffer, cx| {
+            let snapshot = buffer.snapshot(cx);
+            snapshot.anchor_before(Point::new(0, 0))..snapshot.anchor_after(Point::new(4, 2))
+        });
+        let codegen = cx.new_model(|cx| Codegen::new(buffer.clone(), range, None, None, cx));
+
+        let (chunks_tx, chunks_rx) = mpsc::unbounded();
+        codegen.update(cx, |codegen, cx| {
+            codegen.start(
+                String::new(),
+                future::ready(Ok(chunks_rx.map(|chunk| Ok(chunk)).boxed())),
+                cx,
+            )
+        });
+
+        let new_text = concat!(
+            "func main() {\n",
+            "\tx := 0\n",
+            "\tfor x < 10 {\n",
+            "\t\tx++\n",
+            "\t}", //
+        );
+        chunks_tx.unbounded_send(new_text.to_string()).unwrap();
+        drop(chunks_tx);
+        cx.background_executor.run_until_parked();
+
+        assert_eq!(
+            buffer.read_with(cx, |buffer, cx| buffer.snapshot(cx).text()),
+            indoc! {"
+                func main() {
+                \tx := 0
+                \tfor x < 10 {
+                \t\tx++
+                \t}
+                }
+            "}
+        );
+    }
+
     #[gpui::test]
     async fn test_strip_invalid_spans_from_codeblock() {
         assert_chunks("Lorem ipsum dolor", "Lorem ipsum dolor").await;