Fix block-wise autoindent when editing adjacent ranges (#19521)

This fixes problems where auto-indent wasn't working correctly for
assistant edits.

Release Notes:

- Fixed a bug where auto-indent didn't work correctly when pasting with
multiple cursors on adjacent lines

Co-authored-by: Marshall <marshall@zed.dev>

Cargo.lock

@@ -6230,6 +6230,7 @@ dependencies = [
  "lsp",
  "parking_lot",
  "postage",
+ "pretty_assertions",
  "pulldown-cmark 0.12.1",
  "rand 0.8.5",
  "regex",

crates/assistant/src/patch.rs

@@ -717,7 +717,6 @@ mod tests {
         );
 
         // Ensure InsertBefore merges correctly with Update of the same text
-
         assert_edits(
             "
                 fn foo() {
@@ -782,6 +781,90 @@ mod tests {
             .unindent(),
             cx,
         );
+
+        // Correctly indent new text when replacing multiple adjacent indented blocks.
+        assert_edits(
+            "
+            impl Numbers {
+                fn one() {
+                    1
+                }
+
+                fn two() {
+                    2
+                }
+
+                fn three() {
+                    3
+                }
+            }
+            "
+            .unindent(),
+            vec![
+                AssistantEditKind::Update {
+                    old_text: "
+                        fn one() {
+                            1
+                        }
+                    "
+                    .unindent(),
+                    new_text: "
+                        fn one() {
+                            101
+                        }
+                    "
+                    .unindent(),
+                    description: "pick better number".into(),
+                },
+                AssistantEditKind::Update {
+                    old_text: "
+                        fn two() {
+                            2
+                        }
+                    "
+                    .unindent(),
+                    new_text: "
+                        fn two() {
+                            102
+                        }
+                    "
+                    .unindent(),
+                    description: "pick better number".into(),
+                },
+                AssistantEditKind::Update {
+                    old_text: "
+                        fn three() {
+                            3
+                        }
+                    "
+                    .unindent(),
+                    new_text: "
+                        fn three() {
+                            103
+                        }
+                    "
+                    .unindent(),
+                    description: "pick better number".into(),
+                },
+            ],
+            "
+                impl Numbers {
+                    fn one() {
+                        101
+                    }
+
+                    fn two() {
+                        102
+                    }
+
+                    fn three() {
+                        103
+                    }
+                }
+            "
+            .unindent(),
+            cx,
+        );
     }
 
     #[track_caller]

crates/language/Cargo.toml

@@ -71,6 +71,7 @@ env_logger.workspace = true
 gpui = { workspace = true, features = ["test-support"] }
 indoc.workspace = true
 lsp = { workspace = true, features = ["test-support"] }
+pretty_assertions.workspace = true
 rand.workspace = true
 settings = { workspace = true, features = ["test-support"] }
 text = { workspace = true, features = ["test-support"] }

crates/language/src/buffer.rs

@@ -442,7 +442,7 @@ struct AutoindentRequest {
     is_block_mode: bool,
 }
 
-#[derive(Clone)]
+#[derive(Debug, Clone)]
 struct AutoindentRequestEntry {
     /// A range of the buffer whose indentation should be adjusted.
     range: Range<Anchor>,
@@ -1420,24 +1420,17 @@ impl Buffer {
                     yield_now().await;
                 }
 
-                // In block mode, only compute indentation suggestions for the first line
-                // of each insertion. Otherwise, compute suggestions for every inserted line.
-                let new_edited_row_ranges = contiguous_ranges(
-                    row_ranges.iter().flat_map(|(range, _)| {
-                        if request.is_block_mode {
-                            range.start..range.start + 1
-                        } else {
-                            range.clone()
-                        }
-                    }),
-                    max_rows_between_yields,
-                );
-
                 // Compute new suggestions for each line, but only include them in the result
                 // if they differ from the old suggestion for that line.
                 let mut language_indent_sizes = language_indent_sizes_by_new_row.iter().peekable();
                 let mut language_indent_size = IndentSize::default();
-                for new_edited_row_range in new_edited_row_ranges {
+                for (row_range, original_indent_column) in row_ranges {
+                    let new_edited_row_range = if request.is_block_mode {
+                        row_range.start..row_range.start + 1
+                    } else {
+                        row_range.clone()
+                    };
+
                     let suggestions = snapshot
                         .suggest_autoindents(new_edited_row_range.clone())
                         .into_iter()
@@ -1471,22 +1464,9 @@ impl Buffer {
                             }
                         }
                     }
-                    yield_now().await;
-                }
 
-                // For each block of inserted text, adjust the indentation of the remaining
+                    if let (true, Some(original_indent_column)) =
-                // lines of the block by the same amount as the first line was adjusted.
+                        (request.is_block_mode, original_indent_column)
-                if request.is_block_mode {
-                    for (row_range, original_indent_column) in
-                        row_ranges
-                            .into_iter()
-                            .filter_map(|(range, original_indent_column)| {
-                                if range.len() > 1 {
-                                    Some((range, original_indent_column?))
-                                } else {
-                                    None
-                                }
-                            })
                     {
                         let new_indent = indent_sizes
                             .get(&row_range.start)
@@ -1511,6 +1491,8 @@ impl Buffer {
                             }
                         }
                     }
+
+                    yield_now().await;
                 }
             }
 

crates/language/src/buffer_tests.rs

@@ -1658,6 +1658,69 @@ fn test_autoindent_block_mode_without_original_indent_columns(cx: &mut AppContex
     });
 }
 
+#[gpui::test]
+fn test_autoindent_block_mode_multiple_adjacent_ranges(cx: &mut AppContext) {
+    init_settings(cx, |_| {});
+
+    cx.new_model(|cx| {
+        let (text, ranges_to_replace) = marked_text_ranges(
+            &"
+            mod numbers {
+                «fn one() {
+                    1
+                }
+            »
+                «fn two() {
+                    2
+                }
+            »
+                «fn three() {
+                    3
+                }
+            »}
+            "
+            .unindent(),
+            false,
+        );
+
+        let mut buffer = Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx);
+
+        buffer.edit(
+            [
+                (ranges_to_replace[0].clone(), "fn one() {\n    101\n}\n"),
+                (ranges_to_replace[1].clone(), "fn two() {\n    102\n}\n"),
+                (ranges_to_replace[2].clone(), "fn three() {\n    103\n}\n"),
+            ],
+            Some(AutoindentMode::Block {
+                original_indent_columns: vec![0, 0, 0],
+            }),
+            cx,
+        );
+
+        pretty_assertions::assert_eq!(
+            buffer.text(),
+            "
+            mod numbers {
+                fn one() {
+                    101
+                }
+
+                fn two() {
+                    102
+                }
+
+                fn three() {
+                    103
+                }
+            }
+            "
+            .unindent()
+        );
+
+        buffer
+    });
+}
+
 #[gpui::test]
 fn test_autoindent_language_without_indents_query(cx: &mut AppContext) {
     init_settings(cx, |_| {});