Auto indent on `insert_at_line_start` (#5837)

2 files changed, 126 insertions, 10 deletions

diff --git a/helix-term/src/commands.rs b/helix-term/src/commands.rs
index 99f50b04..99f27c00 100644
--- a/helix-term/src/commands.rs
+++ b/helix-term/src/commands.rs
@@ -2775,24 +2775,87 @@ fn last_picker(cx: &mut Context) {
     }));
 }
 
-// I inserts at the first nonwhitespace character of each line with a selection
+/// Fallback position to use for [`insert_with_indent`].
+enum IndentFallbackPos {
+    LineStart,
+    LineEnd,
+}
+
+// `I` inserts at the first nonwhitespace character of each line with a selection.
+// If the line is empty, automatically indent.
 fn insert_at_line_start(cx: &mut Context) {
-    goto_first_nonwhitespace(cx);
-    enter_insert_mode(cx);
+    insert_with_indent(cx, IndentFallbackPos::LineStart);
 }
 
-// A inserts at the end of each line with a selection
+// `A` inserts at the end of each line with a selection.
+// If the line is empty, automatically indent.
 fn insert_at_line_end(cx: &mut Context) {
+    insert_with_indent(cx, IndentFallbackPos::LineEnd);
+}
+
+// Enter insert mode and auto-indent the current line if it is empty.
+// If the line is not empty, move the cursor to the specified fallback position.
+fn insert_with_indent(cx: &mut Context, cursor_fallback: IndentFallbackPos) {
     enter_insert_mode(cx);
+
     let (view, doc) = current!(cx.editor);
 
-    let selection = doc.selection(view.id).clone().transform(|range| {
-        let text = doc.text().slice(..);
-        let line = range.cursor_line(text);
-        let pos = line_end_char_index(&text, line);
-        Range::new(pos, pos)
+    let text = doc.text().slice(..);
+    let contents = doc.text();
+    let selection = doc.selection(view.id);
+
+    let language_config = doc.language_config();
+    let syntax = doc.syntax();
+    let tab_width = doc.tab_width();
+
+    let mut ranges = SmallVec::with_capacity(selection.len());
+    let mut offs = 0;
+
+    let mut transaction = Transaction::change_by_selection(contents, selection, |range| {
+        let cursor_line = range.cursor_line(text);
+        let cursor_line_start = text.line_to_char(cursor_line);
+
+        if line_end_char_index(&text, cursor_line) == cursor_line_start {
+            // line is empty => auto indent
+            let line_end_index = cursor_line_start;
+
+            let indent = indent::indent_for_newline(
+                language_config,
+                syntax,
+                &doc.indent_style,
+                tab_width,
+                text,
+                cursor_line,
+                line_end_index,
+                cursor_line,
+            );
+
+            // calculate new selection ranges
+            let pos = offs + cursor_line_start;
+            let indent_width = indent.chars().count();
+            ranges.push(Range::point(pos + indent_width));
+            offs += indent_width;
+
+            (line_end_index, line_end_index, Some(indent.into()))
+        } else {
+            // move cursor to the fallback position
+            let pos = match cursor_fallback {
+                IndentFallbackPos::LineStart => {
+                    find_first_non_whitespace_char(text.line(cursor_line))
+                        .map(|ws_offset| ws_offset + cursor_line_start)
+                        .unwrap_or(cursor_line_start)
+                }
+                IndentFallbackPos::LineEnd => line_end_char_index(&text, cursor_line),
+            };
+
+            ranges.push(range.put_cursor(text, pos + offs, cx.editor.mode == Mode::Select));
+
+            (cursor_line_start, cursor_line_start, None)
+        }
     });
-    doc.set_selection(view.id, selection);
+
+    transaction = transaction.with_selection(Selection::new(ranges, selection.primary_index()));
+    doc.apply(&transaction, view.id);
 }
 
 // Creates an LspCallback that waits for formatting changes to be computed. When they're done,
diff --git a/helix-term/tests/test/commands.rs b/helix-term/tests/test/commands.rs
index 52b123c7..b13c37bc 100644
--- a/helix-term/tests/test/commands.rs
+++ b/helix-term/tests/test/commands.rs
@@ -426,3 +426,56 @@ async fn test_delete_char_forward() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_insert_with_indent() -> anyhow::Result<()> {
+    const INPUT: &str = "\
+#[f|]#n foo() {
+    if let Some(_) = None {
+
+    }
+\x20
+}
+
+fn bar() {
+
+}";
+
+    // insert_at_line_start
+    test((
+        INPUT,
+        ":lang rust<ret>%<A-s>I",
+        "\
+#[f|]#n foo() {
+    #(i|)#f let Some(_) = None {
+        #(\n|)#\
+\x20   #(}|)#
+#(\x20|)#
+#(}|)#
+#(\n|)#\
+#(f|)#n bar() {
+    #(\n|)#\
+#(}|)#",
+    ))
+    .await?;
+
+    // insert_at_line_end
+    test((
+        INPUT,
+        ":lang rust<ret>%<A-s>A",
+        "\
+fn foo() {#[\n|]#\
+\x20   if let Some(_) = None {#(\n|)#\
+\x20       #(\n|)#\
+\x20   }#(\n|)#\
+\x20#(\n|)#\
+}#(\n|)#\
+#(\n|)#\
+fn bar() {#(\n|)#\
+\x20   #(\n|)#\
+}#(|)#",
+    ))
+    .await?;
+
+    Ok(())
+}