Fix regression: Restore creation of multiple assist editors on ctrl-enter when selections span across multiple excerpts (#16190)

Release Notes:

- N/A

crates/assistant/src/inline_assistant.rs

@@ -45,7 +45,6 @@ use std::{
     task::{self, Poll},
     time::{Duration, Instant},
 };
-use text::ToOffset as _;
 use theme::ThemeSettings;
 use ui::{prelude::*, CheckboxWithLabel, IconButtonShape, Popover, Tooltip};
 use util::{RangeExt, ResultExt};
@@ -139,18 +138,23 @@ impl InlineAssistant {
         cx: &mut WindowContext,
     ) {
         let snapshot = editor.read(cx).buffer().read(cx).snapshot(cx);
-
         struct CodegenRange {
             transform_range: Range<Point>,
             selection_ranges: Vec<Range<Point>>,
             focus_assist: bool,
         }
 
-        let newest_selection = editor.read(cx).selections.newest::<Point>(cx);
+        let newest_selection_range = editor.read(cx).selections.newest::<Point>(cx).range();
         let mut codegen_ranges: Vec<CodegenRange> = Vec::new();
-        for selection in editor.read(cx).selections.all::<Point>(cx) {
+
-            let selection_is_newest = selection.id == newest_selection.id;
+        let selection_ranges = snapshot
-            let mut transform_range = selection.start..selection.end;
+            .split_ranges(editor.read(cx).selections.disjoint_anchor_ranges())
+            .map(|range| range.to_point(&snapshot))
+            .collect::<Vec<Range<Point>>>();
+
+        for selection_range in selection_ranges {
+            let selection_is_newest = newest_selection_range.contains_inclusive(&selection_range);
+            let mut transform_range = selection_range.start..selection_range.end;
 
             // Expand the transform range to start/end of lines.
             // If a non-empty selection ends at the start of the last line, clip at the end of the penultimate line.
@@ -159,7 +163,8 @@ impl InlineAssistant {
                 transform_range.end.row -= 1;
             }
             transform_range.end.column = snapshot.line_len(MultiBufferRow(transform_range.end.row));
-            let selection_range = selection.start..selection.end.min(transform_range.end);
+            let selection_range =
+                selection_range.start..selection_range.end.min(transform_range.end);
 
             // If we intersect the previous transform range,
             if let Some(CodegenRange {
@@ -2343,7 +2348,7 @@ impl Codegen {
         let language_name = language_name.as_deref();
         let start = buffer.point_to_buffer_offset(self.transform_range.start);
         let end = buffer.point_to_buffer_offset(self.transform_range.end);
-        let (buffer, range) = if let Some((start, end)) = start.zip(end) {
+        let (transform_buffer, transform_range) = if let Some((start, end)) = start.zip(end) {
             let (start_buffer, start_buffer_offset) = start;
             let (end_buffer, end_buffer_offset) = end;
             if start_buffer.remote_id() == end_buffer.remote_id() {
@@ -2358,16 +2363,26 @@ impl Codegen {
         let selected_ranges = self
             .selected_ranges
             .iter()
-            .map(|range| {
+            .filter_map(|selected_range| {
-                let start = range.start.text_anchor.to_offset(&buffer);
+                let start = buffer
-                let end = range.end.text_anchor.to_offset(&buffer);
+                    .point_to_buffer_offset(selected_range.start)
-                start..end
+                    .map(|(_, offset)| offset)?;
+                let end = buffer
+                    .point_to_buffer_offset(selected_range.end)
+                    .map(|(_, offset)| offset)?;
+                Some(start..end)
             })
             .collect::<Vec<_>>();
 
         let prompt = self
             .prompt_builder
-            .generate_content_prompt(user_prompt, language_name, buffer, range, selected_ranges)
+            .generate_content_prompt(
+                user_prompt,
+                language_name,
+                transform_buffer,
+                transform_range,
+                selected_ranges,
+            )
             .map_err(|e| anyhow::anyhow!("Failed to generate content prompt: {}", e))?;
 
         let mut messages = Vec::new();

crates/multi_buffer/src/multi_buffer.rs

@@ -3873,7 +3873,60 @@ impl MultiBufferSnapshot {
         }
     }
 
+    // Takes an iterator over anchor ranges and returns a new iterator over anchor ranges that don't
+    // span across excerpt boundaries.
+    pub fn split_ranges<'a, I>(&'a self, ranges: I) -> impl Iterator<Item = Range<Anchor>> + 'a
+    where
+        I: IntoIterator<Item = Range<Anchor>> + 'a,
+    {
+        let mut ranges = ranges.into_iter().map(|range| range.to_offset(self));
+        let mut cursor = self.excerpts.cursor::<usize>();
+        cursor.next(&());
+        let mut current_range = ranges.next();
+        iter::from_fn(move || {
+            let range = current_range.clone()?;
+            if range.start >= cursor.end(&()) {
+                cursor.seek_forward(&range.start, Bias::Right, &());
+                if range.start == self.len() {
+                    cursor.prev(&());
+                }
+            }
+
+            let excerpt = cursor.item()?;
+            let range_start_in_excerpt = cmp::max(range.start, *cursor.start());
+            let range_end_in_excerpt = if excerpt.has_trailing_newline {
+                cmp::min(range.end, cursor.end(&()) - 1)
+            } else {
+                cmp::min(range.end, cursor.end(&()))
+            };
+            let buffer_range = MultiBufferExcerpt::new(excerpt, *cursor.start())
+                .map_range_to_buffer(range_start_in_excerpt..range_end_in_excerpt);
+
+            let subrange_start_anchor = Anchor {
+                buffer_id: Some(excerpt.buffer_id),
+                excerpt_id: excerpt.id,
+                text_anchor: excerpt.buffer.anchor_before(buffer_range.start),
+            };
+            let subrange_end_anchor = Anchor {
+                buffer_id: Some(excerpt.buffer_id),
+                excerpt_id: excerpt.id,
+                text_anchor: excerpt.buffer.anchor_after(buffer_range.end),
+            };
+
+            if range.end > cursor.end(&()) {
+                cursor.next(&());
+            } else {
+                current_range = ranges.next();
+            }
+
+            Some(subrange_start_anchor..subrange_end_anchor)
+        })
+    }
+
     /// Returns excerpts overlapping the given ranges. If range spans multiple excerpts returns one range for each excerpt
+    ///
+    /// The ranges are specified in the coordinate space of the multibuffer, not the individual excerpted buffers.
+    /// Each returned excerpt's range is in the coordinate space of its source buffer.
     pub fn excerpts_in_ranges(
         &self,
         ranges: impl IntoIterator<Item = Range<Anchor>>,
@@ -6707,4 +6760,119 @@ mod tests {
 
         validate_excerpts(&excerpts, &expected_excerpts);
     }
+
+    #[gpui::test]
+    fn test_split_ranges(cx: &mut AppContext) {
+        let buffer_1 = cx.new_model(|cx| Buffer::local(sample_text(6, 6, 'a'), cx));
+        let buffer_2 = cx.new_model(|cx| Buffer::local(sample_text(6, 6, 'g'), cx));
+        let multibuffer = cx.new_model(|_| MultiBuffer::new(0, Capability::ReadWrite));
+        multibuffer.update(cx, |multibuffer, cx| {
+            multibuffer.push_excerpts(
+                buffer_1.clone(),
+                [ExcerptRange {
+                    context: 0..buffer_1.read(cx).len(),
+                    primary: None,
+                }],
+                cx,
+            );
+            multibuffer.push_excerpts(
+                buffer_2.clone(),
+                [ExcerptRange {
+                    context: 0..buffer_2.read(cx).len(),
+                    primary: None,
+                }],
+                cx,
+            );
+        });
+
+        let snapshot = multibuffer.read(cx).snapshot(cx);
+
+        let buffer_1_len = buffer_1.read(cx).len();
+        let buffer_2_len = buffer_2.read(cx).len();
+        let buffer_1_midpoint = buffer_1_len / 2;
+        let buffer_2_start = buffer_1_len + '\n'.len_utf8();
+        let buffer_2_midpoint = buffer_2_start + buffer_2_len / 2;
+        let total_len = buffer_2_start + buffer_2_len;
+
+        let input_ranges = [
+            0..buffer_1_midpoint,
+            buffer_1_midpoint..buffer_2_midpoint,
+            buffer_2_midpoint..total_len,
+        ]
+        .map(|range| snapshot.anchor_before(range.start)..snapshot.anchor_after(range.end));
+
+        let actual_ranges = snapshot
+            .split_ranges(input_ranges.into_iter())
+            .map(|range| range.to_offset(&snapshot))
+            .collect::<Vec<_>>();
+
+        let expected_ranges = vec![
+            0..buffer_1_midpoint,
+            buffer_1_midpoint..buffer_1_len,
+            buffer_2_start..buffer_2_midpoint,
+            buffer_2_midpoint..total_len,
+        ];
+
+        assert_eq!(actual_ranges, expected_ranges);
+    }
+
+    #[gpui::test]
+    fn test_split_ranges_single_range_spanning_three_excerpts(cx: &mut AppContext) {
+        let buffer_1 = cx.new_model(|cx| Buffer::local(sample_text(6, 6, 'a'), cx));
+        let buffer_2 = cx.new_model(|cx| Buffer::local(sample_text(6, 6, 'g'), cx));
+        let buffer_3 = cx.new_model(|cx| Buffer::local(sample_text(6, 6, 'm'), cx));
+        let multibuffer = cx.new_model(|_| MultiBuffer::new(0, Capability::ReadWrite));
+        multibuffer.update(cx, |multibuffer, cx| {
+            multibuffer.push_excerpts(
+                buffer_1.clone(),
+                [ExcerptRange {
+                    context: 0..buffer_1.read(cx).len(),
+                    primary: None,
+                }],
+                cx,
+            );
+            multibuffer.push_excerpts(
+                buffer_2.clone(),
+                [ExcerptRange {
+                    context: 0..buffer_2.read(cx).len(),
+                    primary: None,
+                }],
+                cx,
+            );
+            multibuffer.push_excerpts(
+                buffer_3.clone(),
+                [ExcerptRange {
+                    context: 0..buffer_3.read(cx).len(),
+                    primary: None,
+                }],
+                cx,
+            );
+        });
+
+        let snapshot = multibuffer.read(cx).snapshot(cx);
+
+        let buffer_1_len = buffer_1.read(cx).len();
+        let buffer_2_len = buffer_2.read(cx).len();
+        let buffer_3_len = buffer_3.read(cx).len();
+        let buffer_2_start = buffer_1_len + '\n'.len_utf8();
+        let buffer_3_start = buffer_2_start + buffer_2_len + '\n'.len_utf8();
+        let buffer_1_midpoint = buffer_1_len / 2;
+        let buffer_3_midpoint = buffer_3_start + buffer_3_len / 2;
+
+        let input_range =
+            snapshot.anchor_before(buffer_1_midpoint)..snapshot.anchor_after(buffer_3_midpoint);
+
+        let actual_ranges = snapshot
+            .split_ranges(std::iter::once(input_range))
+            .map(|range| range.to_offset(&snapshot))
+            .collect::<Vec<_>>();
+
+        let expected_ranges = vec![
+            buffer_1_midpoint..buffer_1_len,
+            buffer_2_start..buffer_2_start + buffer_2_len,
+            buffer_3_start..buffer_3_midpoint,
+        ];
+
+        assert_eq!(actual_ranges, expected_ranges);
+    }
 }