Extend symbol ranges by their annotation range when suggesting edits (#15677)

Release Notes:

- N/A

---------

Co-authored-by: Nathan <nathan@zed.dev>

crates/assistant/src/context.rs

@@ -598,6 +598,10 @@ impl EditOperation {
                 buffer.update(&mut cx, |buffer, _| {
                     let outline_item = &outline.items[candidate.id];
                     let symbol_range = outline_item.range.to_point(buffer);
+                    let annotation_range = outline_item
+                        .annotation_range
+                        .as_ref()
+                        .map(|range| range.to_point(buffer));
                     let body_range = outline_item
                         .body_range
                         .as_ref()
@@ -606,23 +610,28 @@ impl EditOperation {
 
                     match kind {
                         EditOperationKind::PrependChild { .. } => {
-                            let position = buffer.anchor_after(body_range.start);
+                            let anchor = buffer.anchor_after(body_range.start);
-                            position..position
+                            anchor..anchor
                         }
                         EditOperationKind::AppendChild { .. } => {
-                            let position = buffer.anchor_before(body_range.end);
+                            let anchor = buffer.anchor_before(body_range.end);
-                            position..position
+                            anchor..anchor
                         }
                         EditOperationKind::InsertSiblingBefore { .. } => {
-                            let position = buffer.anchor_before(symbol_range.start);
+                            let anchor = buffer.anchor_before(
-                            position..position
+                                annotation_range.map_or(symbol_range.start, |annotation_range| {
+                                    annotation_range.start
+                                }),
+                            );
+                            anchor..anchor
                         }
                         EditOperationKind::InsertSiblingAfter { .. } => {
-                            let position = buffer.anchor_after(symbol_range.end);
+                            let anchor = buffer.anchor_after(symbol_range.end);
-                            position..position
+                            anchor..anchor
                         }
                         EditOperationKind::Update { .. } | EditOperationKind::Delete { .. } => {
-                            let start = Point::new(symbol_range.start.row, 0);
+                            let start = annotation_range.map_or(symbol_range.start, |range| range.start);
+                            let start = Point::new(start.row, 0);
                             let end = Point::new(
                                 symbol_range.end.row,
                                 buffer.line_len(symbol_range.end.row),

crates/language/src/buffer.rs

@@ -10,11 +10,11 @@ use crate::{
     markdown::parse_markdown,
     outline::OutlineItem,
     syntax_map::{
-        SyntaxLayer, SyntaxMap, SyntaxMapCapture, SyntaxMapCaptures, SyntaxMapMatches,
+        SyntaxLayer, SyntaxMap, SyntaxMapCapture, SyntaxMapCaptures, SyntaxMapMatch,
-        SyntaxSnapshot, ToTreeSitterPoint,
+        SyntaxMapMatches, SyntaxSnapshot, ToTreeSitterPoint,
     },
     task_context::RunnableRange,
-    LanguageScope, Outline, RunnableCapture, RunnableTag,
+    LanguageScope, Outline, OutlineConfig, RunnableCapture, RunnableTag,
 };
 use anyhow::{anyhow, Context, Result};
 use async_watch as watch;
@@ -2768,130 +2768,44 @@ impl BufferSnapshot {
             .collect::<Vec<_>>();
 
         let mut items = Vec::new();
+        let mut annotation_row_ranges: Vec<Range<u32>> = Vec::new();
         while let Some(mat) = matches.peek() {
             let config = &configs[mat.grammar_index];
-            let item_node = mat.captures.iter().find_map(|cap| {
+            if let Some(item) =
-                if cap.index == config.item_capture_ix {
+                self.next_outline_item(config, &mat, &range, include_extra_context, theme)
-                    Some(cap.node)
+            {
-                } else {
+                items.push(item);
-                    None
+            } else if let Some(capture) = mat
-                }
+                .captures
-            })?;
+                .iter()
-
+                .find(|capture| Some(capture.index) == config.annotation_capture_ix)
-            let item_range = item_node.byte_range();
+            {
-            if item_range.end < range.start || item_range.start > range.end {
+                let capture_range = capture.node.start_position()..capture.node.end_position();
-                matches.advance();
+                let mut capture_row_range =
-                continue;
+                    capture_range.start.row as u32..capture_range.end.row as u32;
-            }
+                if capture_range.end.row > capture_range.start.row && capture_range.end.column == 0
-
-            let mut open_index = None;
-            let mut close_index = None;
-
-            let mut buffer_ranges = Vec::new();
-            for capture in mat.captures {
-                let node_is_name;
-                if capture.index == config.name_capture_ix {
-                    node_is_name = true;
-                } else if Some(capture.index) == config.context_capture_ix
-                    || (Some(capture.index) == config.extra_context_capture_ix
-                        && include_extra_context)
                 {
-                    node_is_name = false;
+                    capture_row_range.end -= 1;
-                } else {
-                    if Some(capture.index) == config.open_capture_ix {
-                        open_index = Some(capture.node.end_byte());
-                    } else if Some(capture.index) == config.close_capture_ix {
-                        close_index = Some(capture.node.start_byte());
-                    }
-
-                    continue;
                 }
-
+                if let Some(last_row_range) = annotation_row_ranges.last_mut() {
-                let mut range = capture.node.start_byte()..capture.node.end_byte();
+                    if last_row_range.end >= capture_row_range.start.saturating_sub(1) {
-                let start = capture.node.start_position();
+                        last_row_range.end = capture_row_range.end;
-                if capture.node.end_position().row > start.row {
-                    range.end =
-                        range.start + self.line_len(start.row as u32) as usize - start.column;
-                }
-
-                if !range.is_empty() {
-                    buffer_ranges.push((range, node_is_name));
-                }
-            }
-
-            if buffer_ranges.is_empty() {
-                matches.advance();
-                continue;
-            }
-
-            let mut text = String::new();
-            let mut highlight_ranges = Vec::new();
-            let mut name_ranges = Vec::new();
-            let mut chunks = self.chunks(
-                buffer_ranges.first().unwrap().0.start..buffer_ranges.last().unwrap().0.end,
-                true,
-            );
-            let mut last_buffer_range_end = 0;
-            for (buffer_range, is_name) in buffer_ranges {
-                if !text.is_empty() && buffer_range.start > last_buffer_range_end {
-                    text.push(' ');
-                }
-                last_buffer_range_end = buffer_range.end;
-                if is_name {
-                    let mut start = text.len();
-                    let end = start + buffer_range.len();
-
-                    // When multiple names are captured, then the matcheable text
-                    // includes the whitespace in between the names.
-                    if !name_ranges.is_empty() {
-                        start -= 1;
-                    }
-
-                    name_ranges.push(start..end);
-                }
-
-                let mut offset = buffer_range.start;
-                chunks.seek(offset);
-                for mut chunk in chunks.by_ref() {
-                    if chunk.text.len() > buffer_range.end - offset {
-                        chunk.text = &chunk.text[0..(buffer_range.end - offset)];
-                        offset = buffer_range.end;
                     } else {
-                        offset += chunk.text.len();
+                        annotation_row_ranges.push(capture_row_range);
-                    }
-                    let style = chunk
-                        .syntax_highlight_id
-                        .zip(theme)
-                        .and_then(|(highlight, theme)| highlight.style(theme));
-                    if let Some(style) = style {
-                        let start = text.len();
-                        let end = start + chunk.text.len();
-                        highlight_ranges.push((start..end, style));
-                    }
-                    text.push_str(chunk.text);
-                    if offset >= buffer_range.end {
-                        break;
                     }
+                } else {
+                    annotation_row_ranges.push(capture_row_range);
                 }
             }
-
             matches.advance();
-
-            items.push(OutlineItem {
-                depth: 0, // We'll calculate the depth later
-                range: item_range,
-                text,
-                highlight_ranges,
-                name_ranges,
-                body_range: open_index.zip(close_index).map(|(start, end)| start..end),
-            });
         }
 
         items.sort_by_key(|item| (item.range.start, Reverse(item.range.end)));
 
         // Assign depths based on containment relationships and convert to anchors.
-        let mut item_ends_stack = Vec::<usize>::new();
+        let mut item_ends_stack = Vec::<Point>::new();
         let mut anchor_items = Vec::new();
+        let mut annotation_row_ranges = annotation_row_ranges.into_iter().peekable();
         for item in items {
             while let Some(last_end) = item_ends_stack.last().copied() {
                 if last_end < item.range.end {
@@ -2901,6 +2815,20 @@ impl BufferSnapshot {
                 }
             }
 
+            let mut annotation_row_range = None;
+            while let Some(next_annotation_row_range) = annotation_row_ranges.peek() {
+                let row_preceding_item = item.range.start.row.saturating_sub(1);
+                if next_annotation_row_range.end < row_preceding_item {
+                    annotation_row_ranges.next();
+                } else {
+                    if next_annotation_row_range.end == row_preceding_item {
+                        annotation_row_range = Some(next_annotation_row_range.clone());
+                        annotation_row_ranges.next();
+                    }
+                    break;
+                }
+            }
+
             anchor_items.push(OutlineItem {
                 depth: item_ends_stack.len(),
                 range: self.anchor_after(item.range.start)..self.anchor_before(item.range.end),
@@ -2910,6 +2838,13 @@ impl BufferSnapshot {
                 body_range: item.body_range.map(|body_range| {
                     self.anchor_after(body_range.start)..self.anchor_before(body_range.end)
                 }),
+                annotation_range: annotation_row_range.map(|annotation_range| {
+                    self.anchor_after(Point::new(annotation_range.start, 0))
+                        ..self.anchor_before(Point::new(
+                            annotation_range.end,
+                            self.line_len(annotation_range.end),
+                        ))
+                }),
             });
             item_ends_stack.push(item.range.end);
         }
@@ -2917,6 +2852,125 @@ impl BufferSnapshot {
         Some(anchor_items)
     }
 
+    fn next_outline_item(
+        &self,
+        config: &OutlineConfig,
+        mat: &SyntaxMapMatch,
+        range: &Range<usize>,
+        include_extra_context: bool,
+        theme: Option<&SyntaxTheme>,
+    ) -> Option<OutlineItem<Point>> {
+        let item_node = mat.captures.iter().find_map(|cap| {
+            if cap.index == config.item_capture_ix {
+                Some(cap.node)
+            } else {
+                None
+            }
+        })?;
+
+        let item_byte_range = item_node.byte_range();
+        if item_byte_range.end < range.start || item_byte_range.start > range.end {
+            return None;
+        }
+        let item_point_range = Point::from_ts_point(item_node.start_position())
+            ..Point::from_ts_point(item_node.end_position());
+
+        let mut open_point = None;
+        let mut close_point = None;
+        let mut buffer_ranges = Vec::new();
+        for capture in mat.captures {
+            let node_is_name;
+            if capture.index == config.name_capture_ix {
+                node_is_name = true;
+            } else if Some(capture.index) == config.context_capture_ix
+                || (Some(capture.index) == config.extra_context_capture_ix && include_extra_context)
+            {
+                node_is_name = false;
+            } else {
+                if Some(capture.index) == config.open_capture_ix {
+                    open_point = Some(Point::from_ts_point(capture.node.end_position()));
+                } else if Some(capture.index) == config.close_capture_ix {
+                    close_point = Some(Point::from_ts_point(capture.node.start_position()));
+                }
+
+                continue;
+            }
+
+            let mut range = capture.node.start_byte()..capture.node.end_byte();
+            let start = capture.node.start_position();
+            if capture.node.end_position().row > start.row {
+                range.end = range.start + self.line_len(start.row as u32) as usize - start.column;
+            }
+
+            if !range.is_empty() {
+                buffer_ranges.push((range, node_is_name));
+            }
+        }
+        if buffer_ranges.is_empty() {
+            return None;
+        }
+        let mut text = String::new();
+        let mut highlight_ranges = Vec::new();
+        let mut name_ranges = Vec::new();
+        let mut chunks = self.chunks(
+            buffer_ranges.first().unwrap().0.start..buffer_ranges.last().unwrap().0.end,
+            true,
+        );
+        let mut last_buffer_range_end = 0;
+        for (buffer_range, is_name) in buffer_ranges {
+            if !text.is_empty() && buffer_range.start > last_buffer_range_end {
+                text.push(' ');
+            }
+            last_buffer_range_end = buffer_range.end;
+            if is_name {
+                let mut start = text.len();
+                let end = start + buffer_range.len();
+
+                // When multiple names are captured, then the matcheable text
+                // includes the whitespace in between the names.
+                if !name_ranges.is_empty() {
+                    start -= 1;
+                }
+
+                name_ranges.push(start..end);
+            }
+
+            let mut offset = buffer_range.start;
+            chunks.seek(offset);
+            for mut chunk in chunks.by_ref() {
+                if chunk.text.len() > buffer_range.end - offset {
+                    chunk.text = &chunk.text[0..(buffer_range.end - offset)];
+                    offset = buffer_range.end;
+                } else {
+                    offset += chunk.text.len();
+                }
+                let style = chunk
+                    .syntax_highlight_id
+                    .zip(theme)
+                    .and_then(|(highlight, theme)| highlight.style(theme));
+                if let Some(style) = style {
+                    let start = text.len();
+                    let end = start + chunk.text.len();
+                    highlight_ranges.push((start..end, style));
+                }
+                text.push_str(chunk.text);
+                if offset >= buffer_range.end {
+                    break;
+                }
+            }
+        }
+
+        Some(OutlineItem {
+            depth: 0, // We'll calculate the depth later
+            range: item_point_range,
+            text,
+            highlight_ranges,
+            name_ranges,
+            body_range: open_point.zip(close_point).map(|(start, end)| start..end),
+            annotation_range: None,
+        })
+    }
+
     /// For each grammar in the language, runs the provided
     /// [tree_sitter::Query] against the given range.
     pub fn matches(

crates/language/src/buffer_tests.rs

@@ -775,6 +775,61 @@ async fn test_outline_with_extra_context(cx: &mut gpui::TestAppContext) {
     );
 }
 
+#[gpui::test]
+fn test_outline_annotations(cx: &mut AppContext) {
+    // Add this new test case
+    let text = r#"
+        /// This is a doc comment
+        /// that spans multiple lines
+        fn annotated_function() {
+            // This is not an annotation
+        }
+
+        // This is a single-line annotation
+        fn another_function() {}
+
+        fn unannotated_function() {}
+
+        // This comment is not an annotation
+
+        fn function_after_blank_line() {}
+    "#
+    .unindent();
+
+    let buffer =
+        cx.new_model(|cx| Buffer::local(text, cx).with_language(Arc::new(rust_lang()), cx));
+    let outline = buffer
+        .update(cx, |buffer, _| buffer.snapshot().outline(None))
+        .unwrap();
+
+    assert_eq!(
+        outline
+            .items
+            .into_iter()
+            .map(|item| (
+                item.text,
+                item.depth,
+                item.annotation_range
+                    .map(|range| { buffer.read(cx).text_for_range(range).collect::<String>() })
+            ))
+            .collect::<Vec<_>>(),
+        &[
+            (
+                "fn annotated_function".to_string(),
+                0,
+                Some("/// This is a doc comment\n/// that spans multiple lines".to_string())
+            ),
+            (
+                "fn another_function".to_string(),
+                0,
+                Some("// This is a single-line annotation".to_string())
+            ),
+            ("fn unannotated_function".to_string(), 0, None),
+            ("fn function_after_blank_line".to_string(), 0, None),
+        ]
+    );
+}
+
 #[gpui::test]
 async fn test_symbols_containing(cx: &mut gpui::TestAppContext) {
     let text = r#"
@@ -2603,6 +2658,8 @@ fn rust_lang() -> Language {
     .unwrap()
     .with_outline_query(
         r#"
+        (line_comment) @annotation
+
         (struct_item
             "struct" @context
             name: (_) @name) @item

crates/language/src/language.rs

@@ -864,6 +864,7 @@ pub struct OutlineConfig {
     pub extra_context_capture_ix: Option<u32>,
     pub open_capture_ix: Option<u32>,
     pub close_capture_ix: Option<u32>,
+    pub annotation_capture_ix: Option<u32>,
 }
 
 #[derive(Debug)]
@@ -1049,6 +1050,7 @@ impl Language {
         let mut extra_context_capture_ix = None;
         let mut open_capture_ix = None;
         let mut close_capture_ix = None;
+        let mut annotation_capture_ix = None;
         get_capture_indices(
             &query,
             &mut [
@@ -1058,6 +1060,7 @@ impl Language {
                 ("context.extra", &mut extra_context_capture_ix),
                 ("open", &mut open_capture_ix),
                 ("close", &mut close_capture_ix),
+                ("annotation", &mut annotation_capture_ix),
             ],
         );
         if let Some((item_capture_ix, name_capture_ix)) = item_capture_ix.zip(name_capture_ix) {
@@ -1069,6 +1072,7 @@ impl Language {
                 extra_context_capture_ix,
                 open_capture_ix,
                 close_capture_ix,
+                annotation_capture_ix,
             });
         }
         Ok(self)

crates/language/src/outline.rs

@@ -21,6 +21,7 @@ pub struct OutlineItem<T> {
     pub highlight_ranges: Vec<(Range<usize>, HighlightStyle)>,
     pub name_ranges: Vec<Range<usize>>,
     pub body_range: Option<Range<T>>,
+    pub annotation_range: Option<Range<T>>,
 }
 
 impl<T> Outline<T> {

crates/languages/src/rust/outline.scm

@@ -1,3 +1,6 @@
+(attribute_item) @annotation
+(line_comment) @annotation
+
 (struct_item
     (visibility_modifier)? @context
     "struct" @context

crates/multi_buffer/src/multi_buffer.rs

@@ -3646,6 +3646,12 @@ impl MultiBufferSnapshot {
                                     ..self.anchor_in_excerpt(*excerpt_id, body_range.end)?,
                             )
                         }),
+                        annotation_range: item.annotation_range.and_then(|annotation_range| {
+                            Some(
+                                self.anchor_in_excerpt(*excerpt_id, annotation_range.start)?
+                                    ..self.anchor_in_excerpt(*excerpt_id, annotation_range.end)?,
+                            )
+                        }),
                     })
                 })
                 .collect(),
@@ -3681,6 +3687,12 @@ impl MultiBufferSnapshot {
                                     ..self.anchor_in_excerpt(excerpt_id, body_range.end)?,
                             )
                         }),
+                        annotation_range: item.annotation_range.and_then(|body_range| {
+                            Some(
+                                self.anchor_in_excerpt(excerpt_id, body_range.start)?
+                                    ..self.anchor_in_excerpt(excerpt_id, body_range.end)?,
+                            )
+                        }),
                     })
                 })
                 .collect(),