1mod declaration;
2mod declaration_scoring;
3mod excerpt;
4mod imports;
5mod outline;
6mod reference;
7mod syntax_index;
8pub mod text_similarity;
9
10use std::{path::Path, sync::Arc};
11
12use cloud_llm_client::predict_edits_v3;
13use collections::HashMap;
14use gpui::{App, AppContext as _, Entity, Task};
15use language::BufferSnapshot;
16use text::{Point, ToOffset as _};
17
18pub use declaration::*;
19pub use declaration_scoring::*;
20pub use excerpt::*;
21pub use imports::*;
22pub use reference::*;
23pub use syntax_index::*;
24
25pub use predict_edits_v3::Line;
26
27#[derive(Clone, Debug, PartialEq)]
28pub struct EditPredictionContextOptions {
29 pub use_imports: bool,
30 pub use_references: bool,
31 pub excerpt: EditPredictionExcerptOptions,
32 pub score: EditPredictionScoreOptions,
33}
34
35#[derive(Clone, Debug)]
36pub struct EditPredictionContext {
37 pub excerpt: EditPredictionExcerpt,
38 pub excerpt_text: EditPredictionExcerptText,
39 pub cursor_point: Point,
40 pub declarations: Vec<ScoredDeclaration>,
41}
42
43impl EditPredictionContext {
44 pub fn gather_context_in_background(
45 cursor_point: Point,
46 buffer: BufferSnapshot,
47 options: EditPredictionContextOptions,
48 syntax_index: Option<Entity<SyntaxIndex>>,
49 cx: &mut App,
50 ) -> Task<Option<Self>> {
51 let parent_abs_path = project::File::from_dyn(buffer.file()).and_then(|f| {
52 let mut path = f.worktree.read(cx).absolutize(&f.path);
53 if path.pop() { Some(path) } else { None }
54 });
55
56 if let Some(syntax_index) = syntax_index {
57 let index_state =
58 syntax_index.read_with(cx, |index, _cx| Arc::downgrade(index.state()));
59 cx.background_spawn(async move {
60 let parent_abs_path = parent_abs_path.as_deref();
61 let index_state = index_state.upgrade()?;
62 let index_state = index_state.lock().await;
63 Self::gather_context(
64 cursor_point,
65 &buffer,
66 parent_abs_path,
67 &options,
68 Some(&index_state),
69 )
70 })
71 } else {
72 cx.background_spawn(async move {
73 let parent_abs_path = parent_abs_path.as_deref();
74 Self::gather_context(cursor_point, &buffer, parent_abs_path, &options, None)
75 })
76 }
77 }
78
79 pub fn gather_context(
80 cursor_point: Point,
81 buffer: &BufferSnapshot,
82 parent_abs_path: Option<&Path>,
83 options: &EditPredictionContextOptions,
84 index_state: Option<&SyntaxIndexState>,
85 ) -> Option<Self> {
86 let imports = if options.use_imports {
87 Imports::gather(&buffer, parent_abs_path)
88 } else {
89 Imports::default()
90 };
91 Self::gather_context_with_references_fn(
92 cursor_point,
93 buffer,
94 &imports,
95 options,
96 index_state,
97 references_in_excerpt,
98 )
99 }
100
101 pub fn gather_context_with_references_fn(
102 cursor_point: Point,
103 buffer: &BufferSnapshot,
104 imports: &Imports,
105 options: &EditPredictionContextOptions,
106 index_state: Option<&SyntaxIndexState>,
107 get_references: impl FnOnce(
108 &EditPredictionExcerpt,
109 &EditPredictionExcerptText,
110 &BufferSnapshot,
111 ) -> HashMap<Identifier, Vec<Reference>>,
112 ) -> Option<Self> {
113 let excerpt = EditPredictionExcerpt::select_from_buffer(
114 cursor_point,
115 buffer,
116 &options.excerpt,
117 index_state,
118 )?;
119 let excerpt_text = excerpt.text(buffer);
120
121 let declarations = if options.use_references
122 && let Some(index_state) = index_state
123 {
124 let excerpt_occurrences =
125 text_similarity::Occurrences::within_string(&excerpt_text.body);
126
127 let adjacent_start = Point::new(cursor_point.row.saturating_sub(2), 0);
128 let adjacent_end = Point::new(cursor_point.row + 1, 0);
129 let adjacent_occurrences = text_similarity::Occurrences::within_string(
130 &buffer
131 .text_for_range(adjacent_start..adjacent_end)
132 .collect::<String>(),
133 );
134
135 let cursor_offset_in_file = cursor_point.to_offset(buffer);
136
137 let references = get_references(&excerpt, &excerpt_text, buffer);
138
139 scored_declarations(
140 &options.score,
141 &index_state,
142 &excerpt,
143 &excerpt_occurrences,
144 &adjacent_occurrences,
145 &imports,
146 references,
147 cursor_offset_in_file,
148 buffer,
149 )
150 } else {
151 vec![]
152 };
153
154 Some(Self {
155 excerpt,
156 excerpt_text,
157 cursor_point,
158 declarations,
159 })
160 }
161}
162
163#[cfg(test)]
164mod tests {
165 use super::*;
166 use std::sync::Arc;
167
168 use gpui::{Entity, TestAppContext};
169 use indoc::indoc;
170 use language::{Language, LanguageConfig, LanguageId, LanguageMatcher, tree_sitter_rust};
171 use project::{FakeFs, Project};
172 use serde_json::json;
173 use settings::SettingsStore;
174 use util::path;
175
176 use crate::{EditPredictionExcerptOptions, SyntaxIndex};
177
178 #[gpui::test]
179 async fn test_call_site(cx: &mut TestAppContext) {
180 let (project, index, _rust_lang_id) = init_test(cx).await;
181
182 let buffer = project
183 .update(cx, |project, cx| {
184 let project_path = project.find_project_path("c.rs", cx).unwrap();
185 project.open_buffer(project_path, cx)
186 })
187 .await
188 .unwrap();
189
190 cx.run_until_parked();
191
192 // first process_data call site
193 let cursor_point = language::Point::new(8, 21);
194 let buffer_snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot());
195
196 let context = cx
197 .update(|cx| {
198 EditPredictionContext::gather_context_in_background(
199 cursor_point,
200 buffer_snapshot,
201 EditPredictionContextOptions {
202 use_imports: true,
203 use_references: true,
204 excerpt: EditPredictionExcerptOptions {
205 max_bytes: 60,
206 min_bytes: 10,
207 target_before_cursor_over_total_bytes: 0.5,
208 },
209 score: EditPredictionScoreOptions {
210 omit_excerpt_overlaps: true,
211 },
212 },
213 Some(index.clone()),
214 cx,
215 )
216 })
217 .await
218 .unwrap();
219
220 let mut snippet_identifiers = context
221 .declarations
222 .iter()
223 .map(|snippet| snippet.identifier.name.as_ref())
224 .collect::<Vec<_>>();
225 snippet_identifiers.sort();
226 assert_eq!(snippet_identifiers, vec!["main", "process_data"]);
227 drop(buffer);
228 }
229
230 async fn init_test(
231 cx: &mut TestAppContext,
232 ) -> (Entity<Project>, Entity<SyntaxIndex>, LanguageId) {
233 cx.update(|cx| {
234 let settings_store = SettingsStore::test(cx);
235 cx.set_global(settings_store);
236 language::init(cx);
237 Project::init_settings(cx);
238 });
239
240 let fs = FakeFs::new(cx.executor());
241 fs.insert_tree(
242 path!("/root"),
243 json!({
244 "a.rs": indoc! {r#"
245 fn main() {
246 let x = 1;
247 let y = 2;
248 let z = add(x, y);
249 println!("Result: {}", z);
250 }
251
252 fn add(a: i32, b: i32) -> i32 {
253 a + b
254 }
255 "#},
256 "b.rs": indoc! {"
257 pub struct Config {
258 pub name: String,
259 pub value: i32,
260 }
261
262 impl Config {
263 pub fn new(name: String, value: i32) -> Self {
264 Config { name, value }
265 }
266 }
267 "},
268 "c.rs": indoc! {r#"
269 use std::collections::HashMap;
270
271 fn main() {
272 let args: Vec<String> = std::env::args().collect();
273 let data: Vec<i32> = args[1..]
274 .iter()
275 .filter_map(|s| s.parse().ok())
276 .collect();
277 let result = process_data(data);
278 println!("{:?}", result);
279 }
280
281 fn process_data(data: Vec<i32>) -> HashMap<i32, usize> {
282 let mut counts = HashMap::new();
283 for value in data {
284 *counts.entry(value).or_insert(0) += 1;
285 }
286 counts
287 }
288
289 #[cfg(test)]
290 mod tests {
291 use super::*;
292
293 #[test]
294 fn test_process_data() {
295 let data = vec![1, 2, 2, 3];
296 let result = process_data(data);
297 assert_eq!(result.get(&2), Some(&2));
298 }
299 }
300 "#}
301 }),
302 )
303 .await;
304 let project = Project::test(fs.clone(), [path!("/root").as_ref()], cx).await;
305 let language_registry = project.read_with(cx, |project, _| project.languages().clone());
306 let lang = rust_lang();
307 let lang_id = lang.id();
308 language_registry.add(Arc::new(lang));
309
310 let file_indexing_parallelism = 2;
311 let index = cx.new(|cx| SyntaxIndex::new(&project, file_indexing_parallelism, cx));
312 cx.run_until_parked();
313
314 (project, index, lang_id)
315 }
316
317 fn rust_lang() -> Language {
318 Language::new(
319 LanguageConfig {
320 name: "Rust".into(),
321 matcher: LanguageMatcher {
322 path_suffixes: vec!["rs".to_string()],
323 ..Default::default()
324 },
325 ..Default::default()
326 },
327 Some(tree_sitter_rust::LANGUAGE.into()),
328 )
329 .with_highlights_query(include_str!("../../languages/src/rust/highlights.scm"))
330 .unwrap()
331 .with_outline_query(include_str!("../../languages/src/rust/outline.scm"))
332 .unwrap()
333 }
334}