1use anyhow::Result;
2use cloud_zeta2_prompt::retrieval_prompt::SearchToolQuery;
3use collections::HashMap;
4use futures::{
5 StreamExt,
6 channel::mpsc::{self, UnboundedSender},
7};
8use gpui::{AppContext, AsyncApp, Entity};
9use language::{Anchor, Buffer, BufferSnapshot, OffsetRangeExt, Point, ToOffset, ToPoint};
10use project::{
11 Project, WorktreeSettings,
12 search::{SearchQuery, SearchResult},
13};
14use smol::channel;
15use std::ops::Range;
16use util::{
17 ResultExt as _,
18 paths::{PathMatcher, PathStyle},
19};
20use workspace::item::Settings as _;
21
22#[cfg(feature = "eval-support")]
23type CachedSearchResults = std::collections::BTreeMap<std::path::PathBuf, Vec<Range<usize>>>;
24
25pub async fn run_retrieval_searches(
26 queries: Vec<SearchToolQuery>,
27 project: Entity<Project>,
28 #[cfg(feature = "eval-support")] eval_cache: Option<std::sync::Arc<dyn crate::EvalCache>>,
29 cx: &mut AsyncApp,
30) -> Result<HashMap<Entity<Buffer>, Vec<Range<Anchor>>>> {
31 #[cfg(feature = "eval-support")]
32 let cache = if let Some(eval_cache) = eval_cache {
33 use crate::EvalCacheEntryKind;
34 use anyhow::Context;
35 use collections::FxHasher;
36 use std::hash::{Hash, Hasher};
37
38 let mut hasher = FxHasher::default();
39 project.read_with(cx, |project, cx| {
40 let mut worktrees = project.worktrees(cx);
41 let Some(worktree) = worktrees.next() else {
42 panic!("Expected a single worktree in eval project. Found none.");
43 };
44 assert!(
45 worktrees.next().is_none(),
46 "Expected a single worktree in eval project. Found more than one."
47 );
48 worktree.read(cx).abs_path().hash(&mut hasher);
49 })?;
50
51 queries.hash(&mut hasher);
52 let key = (EvalCacheEntryKind::Search, hasher.finish());
53
54 if let Some(cached_results) = eval_cache.read(key) {
55 let file_results = serde_json::from_str::<CachedSearchResults>(&cached_results)
56 .context("Failed to deserialize cached search results")?;
57 let mut results = HashMap::default();
58
59 for (path, ranges) in file_results {
60 let buffer = project
61 .update(cx, |project, cx| {
62 let project_path = project.find_project_path(path, cx).unwrap();
63 project.open_buffer(project_path, cx)
64 })?
65 .await?;
66 let snapshot = buffer.read_with(cx, |buffer, _| buffer.snapshot())?;
67 let mut ranges = ranges
68 .into_iter()
69 .map(|range| {
70 snapshot.anchor_before(range.start)..snapshot.anchor_after(range.end)
71 })
72 .collect();
73 merge_anchor_ranges(&mut ranges, &snapshot);
74 results.insert(buffer, ranges);
75 }
76
77 return Ok(results);
78 }
79
80 Some((eval_cache, serde_json::to_string_pretty(&queries)?, key))
81 } else {
82 None
83 };
84
85 let (exclude_matcher, path_style) = project.update(cx, |project, cx| {
86 let global_settings = WorktreeSettings::get_global(cx);
87 let exclude_patterns = global_settings
88 .file_scan_exclusions
89 .sources()
90 .iter()
91 .chain(global_settings.private_files.sources().iter());
92 let path_style = project.path_style(cx);
93 anyhow::Ok((PathMatcher::new(exclude_patterns, path_style)?, path_style))
94 })??;
95
96 let (results_tx, mut results_rx) = mpsc::unbounded();
97
98 for query in queries {
99 let exclude_matcher = exclude_matcher.clone();
100 let results_tx = results_tx.clone();
101 let project = project.clone();
102 cx.spawn(async move |cx| {
103 run_query(
104 query,
105 results_tx.clone(),
106 path_style,
107 exclude_matcher,
108 &project,
109 cx,
110 )
111 .await
112 .log_err();
113 })
114 .detach()
115 }
116 drop(results_tx);
117
118 #[cfg(feature = "eval-support")]
119 let cache = cache.clone();
120 cx.background_spawn(async move {
121 let mut results: HashMap<Entity<Buffer>, Vec<Range<Anchor>>> = HashMap::default();
122 let mut snapshots = HashMap::default();
123
124 let mut total_bytes = 0;
125 'outer: while let Some((buffer, snapshot, excerpts)) = results_rx.next().await {
126 snapshots.insert(buffer.entity_id(), snapshot);
127 let existing = results.entry(buffer).or_default();
128 existing.reserve(excerpts.len());
129
130 for (range, size) in excerpts {
131 // Blunt trimming of the results until we have a proper algorithmic filtering step
132 if (total_bytes + size) > MAX_RESULTS_LEN {
133 log::trace!("Combined results reached limit of {MAX_RESULTS_LEN}B");
134 break 'outer;
135 }
136 total_bytes += size;
137 existing.push(range);
138 }
139 }
140
141 #[cfg(feature = "eval-support")]
142 if let Some((cache, queries, key)) = cache {
143 let cached_results: CachedSearchResults = results
144 .iter()
145 .filter_map(|(buffer, ranges)| {
146 let snapshot = snapshots.get(&buffer.entity_id())?;
147 let path = snapshot.file().map(|f| f.path());
148 let mut ranges = ranges
149 .iter()
150 .map(|range| range.to_offset(&snapshot))
151 .collect::<Vec<_>>();
152 ranges.sort_unstable_by_key(|range| (range.start, range.end));
153
154 Some((path?.as_std_path().to_path_buf(), ranges))
155 })
156 .collect();
157 cache.write(
158 key,
159 &queries,
160 &serde_json::to_string_pretty(&cached_results)?,
161 );
162 }
163
164 for (buffer, ranges) in results.iter_mut() {
165 if let Some(snapshot) = snapshots.get(&buffer.entity_id()) {
166 merge_anchor_ranges(ranges, snapshot);
167 }
168 }
169
170 Ok(results)
171 })
172 .await
173}
174
175fn merge_anchor_ranges(ranges: &mut Vec<Range<Anchor>>, snapshot: &BufferSnapshot) {
176 ranges.sort_unstable_by(|a, b| {
177 a.start
178 .cmp(&b.start, snapshot)
179 .then(b.end.cmp(&b.end, snapshot))
180 });
181
182 let mut index = 1;
183 while index < ranges.len() {
184 if ranges[index - 1]
185 .end
186 .cmp(&ranges[index].start, snapshot)
187 .is_ge()
188 {
189 let removed = ranges.remove(index);
190 ranges[index - 1].end = removed.end;
191 } else {
192 index += 1;
193 }
194 }
195}
196
197const MAX_EXCERPT_LEN: usize = 768;
198const MAX_RESULTS_LEN: usize = MAX_EXCERPT_LEN * 5;
199
200struct SearchJob {
201 buffer: Entity<Buffer>,
202 snapshot: BufferSnapshot,
203 ranges: Vec<Range<usize>>,
204 query_ix: usize,
205 jobs_tx: channel::Sender<SearchJob>,
206}
207
208async fn run_query(
209 input_query: SearchToolQuery,
210 results_tx: UnboundedSender<(Entity<Buffer>, BufferSnapshot, Vec<(Range<Anchor>, usize)>)>,
211 path_style: PathStyle,
212 exclude_matcher: PathMatcher,
213 project: &Entity<Project>,
214 cx: &mut AsyncApp,
215) -> Result<()> {
216 let include_matcher = PathMatcher::new(vec![input_query.glob], path_style)?;
217
218 let make_search = |regex: &str| -> Result<SearchQuery> {
219 SearchQuery::regex(
220 regex,
221 false,
222 true,
223 false,
224 true,
225 include_matcher.clone(),
226 exclude_matcher.clone(),
227 true,
228 None,
229 )
230 };
231
232 if let Some(outer_syntax_regex) = input_query.syntax_node.first() {
233 let outer_syntax_query = make_search(outer_syntax_regex)?;
234 let nested_syntax_queries = input_query
235 .syntax_node
236 .into_iter()
237 .skip(1)
238 .map(|query| make_search(&query))
239 .collect::<Result<Vec<_>>>()?;
240 let content_query = input_query
241 .content
242 .map(|regex| make_search(®ex))
243 .transpose()?;
244
245 let (jobs_tx, jobs_rx) = channel::unbounded();
246
247 let outer_search_results_rx =
248 project.update(cx, |project, cx| project.search(outer_syntax_query, cx))?;
249
250 let outer_search_task = cx.spawn(async move |cx| {
251 futures::pin_mut!(outer_search_results_rx);
252 while let Some(SearchResult::Buffer { buffer, ranges }) =
253 outer_search_results_rx.next().await
254 {
255 buffer
256 .read_with(cx, |buffer, _| buffer.parsing_idle())?
257 .await;
258 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?;
259 let expanded_ranges: Vec<_> = ranges
260 .into_iter()
261 .filter_map(|range| expand_to_parent_range(&range, &snapshot))
262 .collect();
263 jobs_tx
264 .send(SearchJob {
265 buffer,
266 snapshot,
267 ranges: expanded_ranges,
268 query_ix: 0,
269 jobs_tx: jobs_tx.clone(),
270 })
271 .await?;
272 }
273 anyhow::Ok(())
274 });
275
276 let n_workers = cx.background_executor().num_cpus();
277 let search_job_task = cx.background_executor().scoped(|scope| {
278 for _ in 0..n_workers {
279 scope.spawn(async {
280 while let Ok(job) = jobs_rx.recv().await {
281 process_nested_search_job(
282 &results_tx,
283 &nested_syntax_queries,
284 &content_query,
285 job,
286 )
287 .await;
288 }
289 });
290 }
291 });
292
293 search_job_task.await;
294 outer_search_task.await?;
295 } else if let Some(content_regex) = &input_query.content {
296 let search_query = make_search(&content_regex)?;
297
298 let results_rx = project.update(cx, |project, cx| project.search(search_query, cx))?;
299 futures::pin_mut!(results_rx);
300
301 while let Some(SearchResult::Buffer { buffer, ranges }) = results_rx.next().await {
302 let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?;
303
304 let ranges = ranges
305 .into_iter()
306 .map(|range| {
307 let range = range.to_offset(&snapshot);
308 let range = expand_to_entire_lines(range, &snapshot);
309 let size = range.len();
310 let range =
311 snapshot.anchor_before(range.start)..snapshot.anchor_after(range.end);
312 (range, size)
313 })
314 .collect();
315
316 let send_result = results_tx.unbounded_send((buffer.clone(), snapshot.clone(), ranges));
317
318 if let Err(err) = send_result
319 && !err.is_disconnected()
320 {
321 log::error!("{err}");
322 }
323 }
324 } else {
325 log::warn!("Context gathering model produced a glob-only search");
326 }
327
328 anyhow::Ok(())
329}
330
331async fn process_nested_search_job(
332 results_tx: &UnboundedSender<(Entity<Buffer>, BufferSnapshot, Vec<(Range<Anchor>, usize)>)>,
333 queries: &Vec<SearchQuery>,
334 content_query: &Option<SearchQuery>,
335 job: SearchJob,
336) {
337 if let Some(search_query) = queries.get(job.query_ix) {
338 let mut subranges = Vec::new();
339 for range in job.ranges {
340 let start = range.start;
341 let search_results = search_query.search(&job.snapshot, Some(range)).await;
342 for subrange in search_results {
343 let subrange = start + subrange.start..start + subrange.end;
344 subranges.extend(expand_to_parent_range(&subrange, &job.snapshot));
345 }
346 }
347 job.jobs_tx
348 .send(SearchJob {
349 buffer: job.buffer,
350 snapshot: job.snapshot,
351 ranges: subranges,
352 query_ix: job.query_ix + 1,
353 jobs_tx: job.jobs_tx.clone(),
354 })
355 .await
356 .ok();
357 } else {
358 let ranges = if let Some(content_query) = content_query {
359 let mut subranges = Vec::new();
360 for range in job.ranges {
361 let start = range.start;
362 let search_results = content_query.search(&job.snapshot, Some(range)).await;
363 for subrange in search_results {
364 let subrange = start + subrange.start..start + subrange.end;
365 subranges.push(subrange);
366 }
367 }
368 subranges
369 } else {
370 job.ranges
371 };
372
373 let matches = ranges
374 .into_iter()
375 .map(|range| {
376 let snapshot = &job.snapshot;
377 let range = expand_to_entire_lines(range, snapshot);
378 let size = range.len();
379 let range = snapshot.anchor_before(range.start)..snapshot.anchor_after(range.end);
380 (range, size)
381 })
382 .collect();
383
384 let send_result = results_tx.unbounded_send((job.buffer, job.snapshot, matches));
385
386 if let Err(err) = send_result
387 && !err.is_disconnected()
388 {
389 log::error!("{err}");
390 }
391 }
392}
393
394fn expand_to_entire_lines(range: Range<usize>, snapshot: &BufferSnapshot) -> Range<usize> {
395 let mut point_range = range.to_point(snapshot);
396 point_range.start.column = 0;
397 if point_range.end.column > 0 {
398 point_range.end = snapshot.max_point().min(point_range.end + Point::new(1, 0));
399 }
400 point_range.to_offset(snapshot)
401}
402
403fn expand_to_parent_range<T: ToPoint + ToOffset>(
404 range: &Range<T>,
405 snapshot: &BufferSnapshot,
406) -> Option<Range<usize>> {
407 let mut line_range = range.to_point(&snapshot);
408 line_range.start.column = snapshot.indent_size_for_line(line_range.start.row).len;
409 line_range.end.column = snapshot.line_len(line_range.end.row);
410 // TODO skip result if matched line isn't the first node line?
411
412 let node = snapshot.syntax_ancestor(line_range)?;
413 Some(node.byte_range())
414}
415
416#[cfg(test)]
417mod tests {
418 use super::*;
419 use crate::merge_excerpts::merge_excerpts;
420 use cloud_zeta2_prompt::write_codeblock;
421 use edit_prediction_context::Line;
422 use gpui::TestAppContext;
423 use indoc::indoc;
424 use language::{Language, LanguageConfig, LanguageMatcher, tree_sitter_rust};
425 use pretty_assertions::assert_eq;
426 use project::FakeFs;
427 use serde_json::json;
428 use settings::SettingsStore;
429 use std::path::Path;
430 use util::path;
431
432 #[gpui::test]
433 async fn test_retrieval(cx: &mut TestAppContext) {
434 init_test(cx);
435 let fs = FakeFs::new(cx.executor());
436 fs.insert_tree(
437 path!("/root"),
438 json!({
439 "user.rs": indoc!{"
440 pub struct Organization {
441 owner: Arc<User>,
442 }
443
444 pub struct User {
445 first_name: String,
446 last_name: String,
447 }
448
449 impl Organization {
450 pub fn owner(&self) -> Arc<User> {
451 self.owner.clone()
452 }
453 }
454
455 impl User {
456 pub fn new(first_name: String, last_name: String) -> Self {
457 Self {
458 first_name,
459 last_name
460 }
461 }
462
463 pub fn first_name(&self) -> String {
464 self.first_name.clone()
465 }
466
467 pub fn last_name(&self) -> String {
468 self.last_name.clone()
469 }
470 }
471 "},
472 "main.rs": indoc!{r#"
473 fn main() {
474 let user = User::new(FIRST_NAME.clone(), "doe".into());
475 println!("user {:?}", user);
476 }
477 "#},
478 }),
479 )
480 .await;
481
482 let project = Project::test(fs, vec![Path::new(path!("/root"))], cx).await;
483 project.update(cx, |project, _cx| {
484 project.languages().add(rust_lang().into())
485 });
486
487 assert_results(
488 &project,
489 SearchToolQuery {
490 glob: "user.rs".into(),
491 syntax_node: vec!["impl\\s+User".into(), "pub\\s+fn\\s+first_name".into()],
492 content: None,
493 },
494 indoc! {r#"
495 `````root/user.rs
496 …
497 impl User {
498 …
499 pub fn first_name(&self) -> String {
500 self.first_name.clone()
501 }
502 …
503 `````
504 "#},
505 cx,
506 )
507 .await;
508
509 assert_results(
510 &project,
511 SearchToolQuery {
512 glob: "user.rs".into(),
513 syntax_node: vec!["impl\\s+User".into()],
514 content: Some("\\.clone".into()),
515 },
516 indoc! {r#"
517 `````root/user.rs
518 …
519 impl User {
520 …
521 pub fn first_name(&self) -> String {
522 self.first_name.clone()
523 …
524 pub fn last_name(&self) -> String {
525 self.last_name.clone()
526 …
527 `````
528 "#},
529 cx,
530 )
531 .await;
532
533 assert_results(
534 &project,
535 SearchToolQuery {
536 glob: "*.rs".into(),
537 syntax_node: vec![],
538 content: Some("\\.clone".into()),
539 },
540 indoc! {r#"
541 `````root/main.rs
542 fn main() {
543 let user = User::new(FIRST_NAME.clone(), "doe".into());
544 …
545 `````
546
547 `````root/user.rs
548 …
549 impl Organization {
550 pub fn owner(&self) -> Arc<User> {
551 self.owner.clone()
552 …
553 impl User {
554 …
555 pub fn first_name(&self) -> String {
556 self.first_name.clone()
557 …
558 pub fn last_name(&self) -> String {
559 self.last_name.clone()
560 …
561 `````
562 "#},
563 cx,
564 )
565 .await;
566 }
567
568 async fn assert_results(
569 project: &Entity<Project>,
570 query: SearchToolQuery,
571 expected_output: &str,
572 cx: &mut TestAppContext,
573 ) {
574 let results =
575 run_retrieval_searches(vec![query], project.clone(), None, &mut cx.to_async())
576 .await
577 .unwrap();
578
579 let mut results = results.into_iter().collect::<Vec<_>>();
580 results.sort_by_key(|results| {
581 results
582 .0
583 .read_with(cx, |buffer, _| buffer.file().unwrap().path().clone())
584 });
585
586 let mut output = String::new();
587 for (buffer, ranges) in results {
588 buffer.read_with(cx, |buffer, cx| {
589 let excerpts = ranges.into_iter().map(|range| {
590 let point_range = range.to_point(buffer);
591 if point_range.end.column > 0 {
592 Line(point_range.start.row)..Line(point_range.end.row + 1)
593 } else {
594 Line(point_range.start.row)..Line(point_range.end.row)
595 }
596 });
597
598 write_codeblock(
599 &buffer.file().unwrap().full_path(cx),
600 merge_excerpts(&buffer.snapshot(), excerpts).iter(),
601 &[],
602 Line(buffer.max_point().row),
603 false,
604 &mut output,
605 );
606 });
607 }
608 output.pop();
609
610 assert_eq!(output, expected_output);
611 }
612
613 fn rust_lang() -> Language {
614 Language::new(
615 LanguageConfig {
616 name: "Rust".into(),
617 matcher: LanguageMatcher {
618 path_suffixes: vec!["rs".to_string()],
619 ..Default::default()
620 },
621 ..Default::default()
622 },
623 Some(tree_sitter_rust::LANGUAGE.into()),
624 )
625 .with_outline_query(include_str!("../../languages/src/rust/outline.scm"))
626 .unwrap()
627 }
628
629 fn init_test(cx: &mut TestAppContext) {
630 cx.update(move |cx| {
631 let settings_store = SettingsStore::test(cx);
632 cx.set_global(settings_store);
633 zlog::init_test();
634 });
635 }
636}