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 = cx.new(|cx| {
123 Buffer::local(input, cx).with_language_immediate(Arc::new(rust_lang()), cx)
124 });
125 buffer.read_with(cx, |buffer, _cx| {
126 let insertions = cursor_offset
127 .map(|offset| {
128 let point = buffer.offset_to_point(offset);
129 vec![(
130 predict_edits_v3::Point {
131 line: Line(point.row),
132 column: point.column,
133 },
134 "<|cursor|>",
135 )]
136 })
137 .unwrap_or_default();
138 let ranges: Vec<Range<Line>> = ranges
139 .into_iter()
140 .map(|range| {
141 let point_range = range.to_point(&buffer);
142 Line(point_range.start.row)..Line(point_range.end.row)
143 })
144 .collect();
145
146 let mut output = String::new();
147 cloud_zeta2_prompt::write_excerpts(
148 assemble_excerpts(&buffer.snapshot(), ranges).iter(),
149 &insertions,
150 Line(buffer.max_point().row),
151 true,
152 &mut output,
153 );
154 assert_eq!(output, expected_output);
155 });
156 }
157 }
158
159 fn rust_lang() -> Language {
160 Language::new(
161 LanguageConfig {
162 name: "Rust".into(),
163 matcher: LanguageMatcher {
164 path_suffixes: vec!["rs".to_string()],
165 ..Default::default()
166 },
167 ..Default::default()
168 },
169 Some(language::tree_sitter_rust::LANGUAGE.into()),
170 )
171 .with_outline_query(include_str!("../../languages/src/rust/outline.scm"))
172 .unwrap()
173 }
174}