1use cloud_llm_client::predict_edits_v3::Excerpt;
2use edit_prediction_context::Line;
3use language::{BufferSnapshot, Point};
4use std::ops::Range;
5
6pub fn assemble_excerpts(
7 buffer: &BufferSnapshot,
8 merged_line_ranges: impl IntoIterator<Item = Range<Line>>,
9) -> Vec<Excerpt> {
10 let mut output = Vec::new();
11
12 let outline_items = buffer.outline_items_as_points_containing(0..buffer.len(), false, None);
13 let mut outline_items = outline_items.into_iter().peekable();
14
15 for range in merged_line_ranges {
16 let point_range = Point::new(range.start.0, 0)..Point::new(range.end.0, 0);
17
18 while let Some(outline_item) = outline_items.peek() {
19 if outline_item.range.start >= point_range.start {
20 break;
21 }
22 if outline_item.range.end > point_range.start {
23 let mut point_range = outline_item.source_range_for_text.clone();
24 point_range.start.column = 0;
25 point_range.end.column = buffer.line_len(point_range.end.row);
26
27 output.push(Excerpt {
28 start_line: Line(point_range.start.row),
29 text: buffer
30 .text_for_range(point_range.clone())
31 .collect::<String>()
32 .into(),
33 })
34 }
35 outline_items.next();
36 }
37
38 output.push(Excerpt {
39 start_line: Line(point_range.start.row),
40 text: buffer
41 .text_for_range(point_range.clone())
42 .collect::<String>()
43 .into(),
44 })
45 }
46
47 output
48}
49
50#[cfg(test)]
51mod tests {
52 use std::sync::Arc;
53
54 use super::*;
55 use cloud_llm_client::predict_edits_v3;
56 use gpui::{TestAppContext, prelude::*};
57 use indoc::indoc;
58 use language::{Buffer, Language, LanguageConfig, LanguageMatcher, OffsetRangeExt};
59 use pretty_assertions::assert_eq;
60 use util::test::marked_text_ranges;
61
62 #[gpui::test]
63 fn test_rust(cx: &mut TestAppContext) {
64 let table = [
65 (
66 indoc! {r#"
67 struct User {
68 first_name: String,
69 « last_name: String,
70 ageˇ: u32,
71 » email: String,
72 create_at: Instant,
73 }
74
75 impl User {
76 pub fn first_name(&self) -> String {
77 self.first_name.clone()
78 }
79
80 pub fn full_name(&self) -> String {
81 « format!("{} {}", self.first_name, self.last_name)
82 » }
83 }
84 "#},
85 indoc! {r#"
86 1|struct User {
87 …
88 3| last_name: String,
89 4| age<|cursor|>: u32,
90 …
91 9|impl User {
92 …
93 14| pub fn full_name(&self) -> String {
94 15| format!("{} {}", self.first_name, self.last_name)
95 …
96 "#},
97 ),
98 (
99 indoc! {r#"
100 struct User {
101 first_name: String,
102 « last_name: String,
103 age: u32,
104 }
105 »"#
106 },
107 indoc! {r#"
108 1|struct User {
109 …
110 3| last_name: String,
111 4| age: u32,
112 5|}
113 "#},
114 ),
115 ];
116
117 for (input, expected_output) in table {
118 let input_without_ranges = input.replace(['«', '»'], "");
119 let input_without_caret = input.replace('ˇ', "");
120 let cursor_offset = input_without_ranges.find('ˇ');
121 let (input, ranges) = marked_text_ranges(&input_without_caret, false);
122 let buffer =
123 cx.new(|cx| Buffer::local(input, cx).with_language(Arc::new(rust_lang()), cx));
124 buffer.read_with(cx, |buffer, _cx| {
125 let insertions = cursor_offset
126 .map(|offset| {
127 let point = buffer.offset_to_point(offset);
128 vec![(
129 predict_edits_v3::Point {
130 line: Line(point.row),
131 column: point.column,
132 },
133 "<|cursor|>",
134 )]
135 })
136 .unwrap_or_default();
137 let ranges: Vec<Range<Line>> = ranges
138 .into_iter()
139 .map(|range| {
140 let point_range = range.to_point(&buffer);
141 Line(point_range.start.row)..Line(point_range.end.row)
142 })
143 .collect();
144
145 let mut output = String::new();
146 cloud_zeta2_prompt::write_excerpts(
147 assemble_excerpts(&buffer.snapshot(), ranges).iter(),
148 &insertions,
149 Line(buffer.max_point().row),
150 true,
151 &mut output,
152 );
153 assert_eq!(output, expected_output);
154 });
155 }
156 }
157
158 fn rust_lang() -> Language {
159 Language::new(
160 LanguageConfig {
161 name: "Rust".into(),
162 matcher: LanguageMatcher {
163 path_suffixes: vec!["rs".to_string()],
164 ..Default::default()
165 },
166 ..Default::default()
167 },
168 Some(language::tree_sitter_rust::LANGUAGE.into()),
169 )
170 .with_outline_query(include_str!("../../languages/src/rust/outline.scm"))
171 .unwrap()
172 }
173}