use crate::assemble_excerpts::assemble_excerpt_ranges;use anyhow::Result;use collections::HashMap;use futures::{FutureExt, StreamExt as _, channel::mpsc, future};use gpui::{ App, AppContext, AsyncApp, Context, Entity, EntityId, EventEmitter, Task, TaskExt, WeakEntity,};use language::{Anchor, Bias, Buffer, BufferSnapshot, OffsetRangeExt as _, Point, ToOffset as _};use project::{EditPredictionDefinition, Project, ProjectPath};use smallvec::SmallVec;use std::{ collections::hash_map, ops::Range, path::Path, sync::Arc, time::{Duration, Instant},};use util::paths::PathStyle;use util::rel_path::RelPath;use util::{RangeExt as _, ResultExt}; mod assemble_excerpts;mod bm25_context;#[cfg(test)]mod edit_prediction_context_tests;mod editable_context;#[cfg(test)]mod fake_definition_lsp;mod git_log_context; pub use editable_context::{ EditHistoryContextEntry, OracleTarget, collect_editable_context, limit_retrieved_context_to_bytes,}; pub use zeta_prompt::{ContextSource, RelatedExcerpt, RelatedFile}; const IDENTIFIER_LINE_COUNT: u32 = 3;const MAX_CONTEXT_IDENTIFIER_COUNT: usize = 32; pub struct RelatedExcerptStore { project: WeakEntity<Project>, related_buffers: Vec<RelatedBuffer>, cache: HashMap<Identifier, Arc<CacheEntry>>, update_tx: mpsc::UnboundedSender<(Entity<Buffer>, Anchor)>, identifier_line_count: u32,} struct RelatedBuffer { buffer: Entity<Buffer>, path: Arc<Path>, anchor_ranges: Vec<Range<Anchor>>, excerpt_orders: Vec<usize>, cached_file: Option<CachedRelatedFile>,} struct CachedRelatedFile { excerpts: Vec<RelatedExcerpt>, buffer_version: clock::Global,} pub enum RelatedExcerptStoreEvent { StartedRefresh, FinishedRefresh { cache_hit_count: usize, cache_miss_count: usize, mean_definition_latency: Duration, max_definition_latency: Duration, },} #[derive(Clone, Debug, PartialEq, Eq, Hash)]struct Identifier { pub name: String, pub range: Range<Anchor>,} enum DefinitionTask { CacheHit(Arc<CacheEntry>), CacheMiss { project: WeakEntity<Project>, task: Task<Result<Vec<EditPredictionDefinition>>>, },} #[derive(Debug)]struct CacheEntry { definitions: SmallVec<[CachedDefinition; 1]>,} #[derive(Clone, Debug)]struct CachedDefinition { path: ProjectPath, buffer: Entity<Buffer>, anchor_range: Range<Anchor>,} const DEBOUNCE_DURATION: Duration = Duration::from_millis(100); impl EventEmitter<RelatedExcerptStoreEvent> for RelatedExcerptStore {} impl RelatedExcerptStore { pub fn new(project: &Entity<Project>, cx: &mut Context<Self>) -> Self { let (update_tx, mut update_rx) = mpsc::unbounded::<(Entity<Buffer>, Anchor)>(); cx.spawn(async move |this, cx| { let executor = cx.background_executor().clone(); while let Some((mut buffer, mut position)) = update_rx.next().await { let mut timer = executor.timer(DEBOUNCE_DURATION).fuse(); loop { futures::select_biased! { next = update_rx.next() => { if let Some((new_buffer, new_position)) = next { buffer = new_buffer; position = new_position; timer = executor.timer(DEBOUNCE_DURATION).fuse(); } else { return anyhow::Ok(()); } } _ = timer => break, } } Self::fetch_excerpts(this.clone(), buffer, position, cx).await?; } anyhow::Ok(()) }) .detach_and_log_err(cx); RelatedExcerptStore { project: project.downgrade(), update_tx, related_buffers: Vec::new(), cache: Default::default(), identifier_line_count: IDENTIFIER_LINE_COUNT, } } pub fn set_identifier_line_count(&mut self, count: u32) { self.identifier_line_count = count; } pub fn refresh(&mut self, buffer: Entity<Buffer>, position: Anchor, _: &mut Context<Self>) { self.update_tx.unbounded_send((buffer, position)).ok(); } pub fn related_files(&mut self, cx: &App) -> Vec<RelatedFile> { self.related_buffers .iter_mut() .map(|related| related.related_file(cx)) .collect() } pub fn related_files_with_buffers( &mut self, cx: &App, ) -> impl Iterator<Item = (RelatedFile, Entity<Buffer>)> { self.related_buffers .iter_mut() .map(|related| (related.related_file(cx), related.buffer.clone())) } pub fn set_related_files(&mut self, files: Vec<RelatedFile>, cx: &App) { self.related_buffers = files .into_iter() .filter_map(|file| { let project = self.project.upgrade()?; let project = project.read(cx); let worktree = project.worktrees(cx).find(|wt| { let root_name = wt.read(cx).root_name().as_unix_str(); file.path .components() .next() .is_some_and(|c| c.as_os_str() == root_name) })?; let worktree = worktree.read(cx); let relative_path = file .path .strip_prefix(worktree.root_name().as_unix_str()) .ok()?; let relative_path = RelPath::new(relative_path, PathStyle::Unix).ok()?; let project_path = ProjectPath { worktree_id: worktree.id(), path: relative_path.into_owned().into(), }; let buffer = project.get_open_buffer(&project_path, cx)?; let snapshot = buffer.read(cx).snapshot(); let mut anchor_ranges = Vec::with_capacity(file.excerpts.len()); let mut excerpt_orders = Vec::with_capacity(file.excerpts.len()); for excerpt in &file.excerpts { let start = snapshot.anchor_before(Point::new(excerpt.row_range.start, 0)); let end_row = excerpt.row_range.end; let end_col = snapshot.line_len(end_row); let end = snapshot.anchor_after(Point::new(end_row, end_col)); anchor_ranges.push(start..end); excerpt_orders.push(excerpt.order); } Some(RelatedBuffer { buffer, path: file.path.clone(), anchor_ranges, excerpt_orders, cached_file: None, }) }) .collect(); } async fn fetch_excerpts( this: WeakEntity<Self>, buffer: Entity<Buffer>, position: Anchor, cx: &mut AsyncApp, ) -> Result<()> { let (project, snapshot, identifier_line_count) = this.read_with(cx, |this, cx| { ( this.project.upgrade(), buffer.read(cx).snapshot(), this.identifier_line_count, ) })?; let Some(project) = project else { return Ok(()); }; let file = snapshot.file().cloned(); let file_extension = file .as_ref() .and_then(|file| file.path().extension()) .unwrap_or("") .to_string(); if let Some(file) = &file { log::debug!("retrieving_context buffer:{}", file.path().as_unix_str()); } let (lsp_store, is_via_ssh) = project.read_with(cx, |project, _| { (project.lsp_store(), project.is_via_remote_server()) }); let lsp_names = lsp_store.update(cx, |lsp_store, cx| { buffer.update(cx, |buffer, cx| { lsp_store .running_language_servers_for_local_buffer(buffer, cx) .map(|(_, server)| server.name().to_string()) .collect::<Vec<_>>() }) }); this.update(cx, |_, cx| { cx.emit(RelatedExcerptStoreEvent::StartedRefresh); })?; let identifiers_with_ranks = cx .background_spawn(async move { let cursor_offset = position.to_offset(&snapshot); let identifiers = identifiers_for_position(&snapshot, position, identifier_line_count); // Compute byte distance from cursor to each identifier, then sort by // distance so we can assign ordinal ranks. Identifiers at the same // distance share the same rank. let mut identifiers_with_distance: Vec<(Identifier, usize)> = identifiers .into_iter() .map(|id| { let start = id.range.start.to_offset(&snapshot); let end = id.range.end.to_offset(&snapshot); let distance = if cursor_offset < start { start - cursor_offset } else if cursor_offset > end { cursor_offset - end } else { 0 }; (id, distance) }) .collect(); // Only the closest `MAX_CONTEXT_IDENTIFIER_COUNT` identifiers are // used below, so select that prefix instead of fully sorting. util::truncate_to_bottom_n_sorted_by( &mut identifiers_with_distance, MAX_CONTEXT_IDENTIFIER_COUNT, &|(_, a), (_, b)| a.cmp(b), ); let mut cursor_distances: HashMap<Identifier, usize> = HashMap::default(); let mut current_rank = 0; let mut previous_distance = None; for (identifier, distance) in &identifiers_with_distance { if previous_distance != Some(*distance) { current_rank = cursor_distances.len(); previous_distance = Some(*distance); } cursor_distances.insert(identifier.clone(), current_rank); } (identifiers_with_distance, cursor_distances) }) .await; let (identifiers_with_distance, cursor_distances) = identifiers_with_ranks; let async_cx = cx.clone(); let start_time = Instant::now(); let futures = this.update(cx, |this, cx| { identifiers_with_distance .into_iter() .map(|(identifier, _)| { let task = if let Some(entry) = this.cache.get(&identifier) { DefinitionTask::CacheHit(entry.clone()) } else { let project = this.project.clone(); let task = project .update(cx, |project, cx| { // tombi LSP for toml will open a scratch buffer with the JSON schema of // the toml file when a goto type definition is requested let include_type_definitions = !is_tombi_lsp_in_toml(project, &buffer, cx); project.edit_prediction_definitions( &buffer, identifier.range.start, include_type_definitions, cx, ) }) .unwrap_or_else(|_| Task::ready(Ok(Vec::new()))); DefinitionTask::CacheMiss { project, task } }; let cx = async_cx.clone(); async move { match task { DefinitionTask::CacheHit(cache_entry) => { Some((identifier, cache_entry, None)) } DefinitionTask::CacheMiss { project, task } => { let definition_locations = task.await.log_err().unwrap_or_default(); let duration = start_time.elapsed(); let definitions: SmallVec<[CachedDefinition; 1]> = future::join_all(definition_locations.into_iter().map( |definition| { let project = project.clone(); let mut cx = cx.clone(); async move { process_definition(definition, &project, &mut cx) .await } }, )) .await .into_iter() .flatten() .collect(); Some(( identifier, Arc::new(CacheEntry { definitions }), Some(duration), )) } } } }) .collect::<Vec<_>>() })?; let mut cache_hit_count = 0; let mut cache_miss_count = 0; let mut mean_definition_latency = Duration::ZERO; let mut max_definition_latency = Duration::ZERO; let mut new_cache = HashMap::default(); new_cache.reserve(futures.len()); for (identifier, entry, duration) in future::join_all(futures).await.into_iter().flatten() { new_cache.insert(identifier, entry); if let Some(duration) = duration { cache_miss_count += 1; mean_definition_latency += duration; max_definition_latency = max_definition_latency.max(duration); } else { cache_hit_count += 1; } } let lsp_fetch_latency_ms = start_time.elapsed().as_millis(); mean_definition_latency /= cache_miss_count.max(1) as u32; let (new_cache, related_buffers) = rebuild_related_files(&project, new_cache, &cursor_distances, cx).await?; let latency_ms = start_time.elapsed().as_millis(); let returned_excerpt_count = related_buffers .iter() .map(|related_buffer| related_buffer.anchor_ranges.len()) .sum::<usize>(); telemetry::event!( "Edit Prediction LSP Context Retrieved", lsp_names, file_extension, latency_ms, lsp_fetch_latency_ms, returned_excerpt_count, is_via_ssh ); if let Some(file) = &file { log::debug!( "finished retrieving context buffer:{}, latency:{:?}", file.path().as_unix_str(), start_time.elapsed() ); } this.update(cx, |this, cx| { this.cache = new_cache; this.related_buffers = related_buffers; cx.emit(RelatedExcerptStoreEvent::FinishedRefresh { cache_hit_count, cache_miss_count, mean_definition_latency, max_definition_latency, }); })?; anyhow::Ok(()) }} async fn rebuild_related_files( project: &Entity<Project>, mut new_entries: HashMap<Identifier, Arc<CacheEntry>>, cursor_distances: &HashMap<Identifier, usize>, cx: &mut AsyncApp,) -> Result<(HashMap<Identifier, Arc<CacheEntry>>, Vec<RelatedBuffer>)> { let mut snapshots = HashMap::default(); let mut worktree_root_names = HashMap::default(); for entry in new_entries.values() { for definition in entry.definitions.iter() { if let hash_map::Entry::Vacant(e) = snapshots.entry(definition.buffer.entity_id()) { definition .buffer .read_with(cx, |buffer, _| buffer.parsing_idle()) .await; e.insert( definition .buffer .read_with(cx, |buffer, _| buffer.snapshot()), ); } let worktree_id = definition.path.worktree_id; if let hash_map::Entry::Vacant(e) = worktree_root_names.entry(definition.path.worktree_id) { project.read_with(cx, |project, cx| { if let Some(worktree) = project.worktree_for_id(worktree_id, cx) { e.insert(worktree.read(cx).root_name().as_unix_str().to_string()); } }); } } } let cursor_distances = cursor_distances.clone(); Ok(cx .background_spawn(async move { let mut ranges_by_buffer = HashMap::<EntityId, (Entity<Buffer>, Vec<(Range<Point>, usize)>)>::default(); let mut paths_by_buffer = HashMap::default(); let mut min_rank_by_buffer = HashMap::<EntityId, usize>::default(); for (identifier, entry) in new_entries.iter_mut() { let rank = cursor_distances .get(identifier) .copied() .unwrap_or(usize::MAX); for definition in entry.definitions.iter() { let Some(snapshot) = snapshots.get(&definition.buffer.entity_id()) else { continue; }; paths_by_buffer.insert(definition.buffer.entity_id(), definition.path.clone()); let buffer_rank = min_rank_by_buffer .entry(definition.buffer.entity_id()) .or_insert(usize::MAX); *buffer_rank = (*buffer_rank).min(rank); ranges_by_buffer .entry(definition.buffer.entity_id()) .or_insert_with(|| (definition.buffer.clone(), Vec::new())) .1 .push((definition.anchor_range.to_point(snapshot), rank)); } } let mut related_buffers: Vec<RelatedBuffer> = ranges_by_buffer .into_iter() .filter_map(|(entity_id, (buffer, ranges))| { let snapshot = snapshots.get(&entity_id)?; let project_path = paths_by_buffer.get(&entity_id)?; let assembled = assemble_excerpt_ranges(snapshot, ranges); let root_name = worktree_root_names.get(&project_path.worktree_id)?; let path: Arc<Path> = Path::new(&format!( "{}/{}", root_name, project_path.path.as_unix_str()Lines 1–500 of 762. Load more
use crate::assemble_excerpts::assemble_excerpt_ranges; use anyhow::Result; use collections::HashMap; use futures::{FutureExt, StreamExt as _, channel::mpsc, future}; use gpui::{ App, AppContext, AsyncApp, Context, Entity, EntityId, EventEmitter, Task, TaskExt, WeakEntity, }; use language::{Anchor, Bias, Buffer, BufferSnapshot, OffsetRangeExt as _, Point, ToOffset as _}; use project::{EditPredictionDefinition, Project, ProjectPath}; use smallvec::SmallVec; use std::{ collections::hash_map, ops::Range, path::Path, sync::Arc, time::{Duration, Instant}, }; use util::paths::PathStyle; use util::rel_path::RelPath; use util::{RangeExt as _, ResultExt}; mod assemble_excerpts; mod bm25_context; #[cfg(test)] mod edit_prediction_context_tests; mod editable_context; #[cfg(test)] mod fake_definition_lsp; mod git_log_context; pub use editable_context::{ EditHistoryContextEntry, OracleTarget, collect_editable_context, limit_retrieved_context_to_bytes, }; pub use zeta_prompt::{ContextSource, RelatedExcerpt, RelatedFile}; const IDENTIFIER_LINE_COUNT: u32 = 3; const MAX_CONTEXT_IDENTIFIER_COUNT: usize = 32; pub struct RelatedExcerptStore { project: WeakEntity<Project>, related_buffers: Vec<RelatedBuffer>, cache: HashMap<Identifier, Arc<CacheEntry>>, update_tx: mpsc::UnboundedSender<(Entity<Buffer>, Anchor)>, identifier_line_count: u32, } struct RelatedBuffer { buffer: Entity<Buffer>, path: Arc<Path>, anchor_ranges: Vec<Range<Anchor>>, excerpt_orders: Vec<usize>, cached_file: Option<CachedRelatedFile>, } struct CachedRelatedFile { excerpts: Vec<RelatedExcerpt>, buffer_version: clock::Global, } pub enum RelatedExcerptStoreEvent { StartedRefresh, FinishedRefresh { cache_hit_count: usize, cache_miss_count: usize, mean_definition_latency: Duration, max_definition_latency: Duration, }, } #[derive(Clone, Debug, PartialEq, Eq, Hash)] struct Identifier { pub name: String, pub range: Range<Anchor>, } enum DefinitionTask { CacheHit(Arc<CacheEntry>), CacheMiss { project: WeakEntity<Project>, task: Task<Result<Vec<EditPredictionDefinition>>>, }, } #[derive(Debug)] struct CacheEntry { definitions: SmallVec<[CachedDefinition; 1]>, } #[derive(Clone, Debug)] struct CachedDefinition { path: ProjectPath, buffer: Entity<Buffer>, anchor_range: Range<Anchor>, } const DEBOUNCE_DURATION: Duration = Duration::from_millis(100); impl EventEmitter<RelatedExcerptStoreEvent> for RelatedExcerptStore {} impl RelatedExcerptStore { pub fn new(project: &Entity<Project>, cx: &mut Context<Self>) -> Self { let (update_tx, mut update_rx) = mpsc::unbounded::<(Entity<Buffer>, Anchor)>(); cx.spawn(async move |this, cx| { let executor = cx.background_executor().clone(); while let Some((mut buffer, mut position)) = update_rx.next().await { let mut timer = executor.timer(DEBOUNCE_DURATION).fuse(); loop { futures::select_biased! { next = update_rx.next() => { if let Some((new_buffer, new_position)) = next { buffer = new_buffer; position = new_position; timer = executor.timer(DEBOUNCE_DURATION).fuse(); } else { return anyhow::Ok(()); } } _ = timer => break, } } Self::fetch_excerpts(this.clone(), buffer, position, cx).await?; } anyhow::Ok(()) }) .detach_and_log_err(cx); RelatedExcerptStore { project: project.downgrade(), update_tx, related_buffers: Vec::new(), cache: Default::default(), identifier_line_count: IDENTIFIER_LINE_COUNT, } } pub fn set_identifier_line_count(&mut self, count: u32) { self.identifier_line_count = count; } pub fn refresh(&mut self, buffer: Entity<Buffer>, position: Anchor, _: &mut Context<Self>) { self.update_tx.unbounded_send((buffer, position)).ok(); } pub fn related_files(&mut self, cx: &App) -> Vec<RelatedFile> { self.related_buffers .iter_mut() .map(|related| related.related_file(cx)) .collect() } pub fn related_files_with_buffers( &mut self, cx: &App, ) -> impl Iterator<Item = (RelatedFile, Entity<Buffer>)> { self.related_buffers .iter_mut() .map(|related| (related.related_file(cx), related.buffer.clone())) } pub fn set_related_files(&mut self, files: Vec<RelatedFile>, cx: &App) { self.related_buffers = files .into_iter() .filter_map(|file| { let project = self.project.upgrade()?; let project = project.read(cx); let worktree = project.worktrees(cx).find(|wt| { let root_name = wt.read(cx).root_name().as_unix_str(); file.path .components() .next() .is_some_and(|c| c.as_os_str() == root_name) })?; let worktree = worktree.read(cx); let relative_path = file .path .strip_prefix(worktree.root_name().as_unix_str()) .ok()?; let relative_path = RelPath::new(relative_path, PathStyle::Unix).ok()?; let project_path = ProjectPath { worktree_id: worktree.id(), path: relative_path.into_owned().into(), }; let buffer = project.get_open_buffer(&project_path, cx)?; let snapshot = buffer.read(cx).snapshot(); let mut anchor_ranges = Vec::with_capacity(file.excerpts.len()); let mut excerpt_orders = Vec::with_capacity(file.excerpts.len()); for excerpt in &file.excerpts { let start = snapshot.anchor_before(Point::new(excerpt.row_range.start, 0)); let end_row = excerpt.row_range.end; let end_col = snapshot.line_len(end_row); let end = snapshot.anchor_after(Point::new(end_row, end_col)); anchor_ranges.push(start..end); excerpt_orders.push(excerpt.order); } Some(RelatedBuffer { buffer, path: file.path.clone(), anchor_ranges, excerpt_orders, cached_file: None, }) }) .collect(); } async fn fetch_excerpts( this: WeakEntity<Self>, buffer: Entity<Buffer>, position: Anchor, cx: &mut AsyncApp, ) -> Result<()> { let (project, snapshot, identifier_line_count) = this.read_with(cx, |this, cx| { ( this.project.upgrade(), buffer.read(cx).snapshot(), this.identifier_line_count, ) })?; let Some(project) = project else { return Ok(()); }; let file = snapshot.file().cloned(); let file_extension = file .as_ref() .and_then(|file| file.path().extension()) .unwrap_or("") .to_string(); if let Some(file) = &file { log::debug!("retrieving_context buffer:{}", file.path().as_unix_str()); } let (lsp_store, is_via_ssh) = project.read_with(cx, |project, _| { (project.lsp_store(), project.is_via_remote_server()) }); let lsp_names = lsp_store.update(cx, |lsp_store, cx| { buffer.update(cx, |buffer, cx| { lsp_store .running_language_servers_for_local_buffer(buffer, cx) .map(|(_, server)| server.name().to_string()) .collect::<Vec<_>>() }) }); this.update(cx, |_, cx| { cx.emit(RelatedExcerptStoreEvent::StartedRefresh); })?; let identifiers_with_ranks = cx .background_spawn(async move { let cursor_offset = position.to_offset(&snapshot); let identifiers = identifiers_for_position(&snapshot, position, identifier_line_count); // Compute byte distance from cursor to each identifier, then sort by // distance so we can assign ordinal ranks. Identifiers at the same // distance share the same rank. let mut identifiers_with_distance: Vec<(Identifier, usize)> = identifiers .into_iter() .map(|id| { let start = id.range.start.to_offset(&snapshot); let end = id.range.end.to_offset(&snapshot); let distance = if cursor_offset < start { start - cursor_offset } else if cursor_offset > end { cursor_offset - end } else { 0 }; (id, distance) }) .collect(); // Only the closest `MAX_CONTEXT_IDENTIFIER_COUNT` identifiers are // used below, so select that prefix instead of fully sorting. util::truncate_to_bottom_n_sorted_by( &mut identifiers_with_distance, MAX_CONTEXT_IDENTIFIER_COUNT, &|(_, a), (_, b)| a.cmp(b), ); let mut cursor_distances: HashMap<Identifier, usize> = HashMap::default(); let mut current_rank = 0; let mut previous_distance = None; for (identifier, distance) in &identifiers_with_distance { if previous_distance != Some(*distance) { current_rank = cursor_distances.len(); previous_distance = Some(*distance); } cursor_distances.insert(identifier.clone(), current_rank); } (identifiers_with_distance, cursor_distances) }) .await; let (identifiers_with_distance, cursor_distances) = identifiers_with_ranks; let async_cx = cx.clone(); let start_time = Instant::now(); let futures = this.update(cx, |this, cx| { identifiers_with_distance .into_iter() .map(|(identifier, _)| { let task = if let Some(entry) = this.cache.get(&identifier) { DefinitionTask::CacheHit(entry.clone()) } else { let project = this.project.clone(); let task = project .update(cx, |project, cx| { // tombi LSP for toml will open a scratch buffer with the JSON schema of // the toml file when a goto type definition is requested let include_type_definitions = !is_tombi_lsp_in_toml(project, &buffer, cx); project.edit_prediction_definitions( &buffer, identifier.range.start, include_type_definitions, cx, ) }) .unwrap_or_else(|_| Task::ready(Ok(Vec::new()))); DefinitionTask::CacheMiss { project, task } }; let cx = async_cx.clone(); async move { match task { DefinitionTask::CacheHit(cache_entry) => { Some((identifier, cache_entry, None)) } DefinitionTask::CacheMiss { project, task } => { let definition_locations = task.await.log_err().unwrap_or_default(); let duration = start_time.elapsed(); let definitions: SmallVec<[CachedDefinition; 1]> = future::join_all(definition_locations.into_iter().map( |definition| { let project = project.clone(); let mut cx = cx.clone(); async move { process_definition(definition, &project, &mut cx) .await } }, )) .await .into_iter() .flatten() .collect(); Some(( identifier, Arc::new(CacheEntry { definitions }), Some(duration), )) } } } }) .collect::<Vec<_>>() })?; let mut cache_hit_count = 0; let mut cache_miss_count = 0; let mut mean_definition_latency = Duration::ZERO; let mut max_definition_latency = Duration::ZERO; let mut new_cache = HashMap::default(); new_cache.reserve(futures.len()); for (identifier, entry, duration) in future::join_all(futures).await.into_iter().flatten() { new_cache.insert(identifier, entry); if let Some(duration) = duration { cache_miss_count += 1; mean_definition_latency += duration; max_definition_latency = max_definition_latency.max(duration); } else { cache_hit_count += 1; } } let lsp_fetch_latency_ms = start_time.elapsed().as_millis(); mean_definition_latency /= cache_miss_count.max(1) as u32; let (new_cache, related_buffers) = rebuild_related_files(&project, new_cache, &cursor_distances, cx).await?; let latency_ms = start_time.elapsed().as_millis(); let returned_excerpt_count = related_buffers .iter() .map(|related_buffer| related_buffer.anchor_ranges.len()) .sum::<usize>(); telemetry::event!( "Edit Prediction LSP Context Retrieved", lsp_names, file_extension, latency_ms, lsp_fetch_latency_ms, returned_excerpt_count, is_via_ssh ); if let Some(file) = &file { log::debug!( "finished retrieving context buffer:{}, latency:{:?}", file.path().as_unix_str(), start_time.elapsed() ); } this.update(cx, |this, cx| { this.cache = new_cache; this.related_buffers = related_buffers; cx.emit(RelatedExcerptStoreEvent::FinishedRefresh { cache_hit_count, cache_miss_count, mean_definition_latency, max_definition_latency, }); })?; anyhow::Ok(()) } } async fn rebuild_related_files( project: &Entity<Project>, mut new_entries: HashMap<Identifier, Arc<CacheEntry>>, cursor_distances: &HashMap<Identifier, usize>, cx: &mut AsyncApp, ) -> Result<(HashMap<Identifier, Arc<CacheEntry>>, Vec<RelatedBuffer>)> { let mut snapshots = HashMap::default(); let mut worktree_root_names = HashMap::default(); for entry in new_entries.values() { for definition in entry.definitions.iter() { if let hash_map::Entry::Vacant(e) = snapshots.entry(definition.buffer.entity_id()) { definition .buffer .read_with(cx, |buffer, _| buffer.parsing_idle()) .await; e.insert( definition .buffer .read_with(cx, |buffer, _| buffer.snapshot()), ); } let worktree_id = definition.path.worktree_id; if let hash_map::Entry::Vacant(e) = worktree_root_names.entry(definition.path.worktree_id) { project.read_with(cx, |project, cx| { if let Some(worktree) = project.worktree_for_id(worktree_id, cx) { e.insert(worktree.read(cx).root_name().as_unix_str().to_string()); } }); } } } let cursor_distances = cursor_distances.clone(); Ok(cx .background_spawn(async move { let mut ranges_by_buffer = HashMap::<EntityId, (Entity<Buffer>, Vec<(Range<Point>, usize)>)>::default(); let mut paths_by_buffer = HashMap::default(); let mut min_rank_by_buffer = HashMap::<EntityId, usize>::default(); for (identifier, entry) in new_entries.iter_mut() { let rank = cursor_distances .get(identifier) .copied() .unwrap_or(usize::MAX); for definition in entry.definitions.iter() { let Some(snapshot) = snapshots.get(&definition.buffer.entity_id()) else { continue; }; paths_by_buffer.insert(definition.buffer.entity_id(), definition.path.clone()); let buffer_rank = min_rank_by_buffer .entry(definition.buffer.entity_id()) .or_insert(usize::MAX); *buffer_rank = (*buffer_rank).min(rank); ranges_by_buffer .entry(definition.buffer.entity_id()) .or_insert_with(|| (definition.buffer.clone(), Vec::new())) .1 .push((definition.anchor_range.to_point(snapshot), rank)); } } let mut related_buffers: Vec<RelatedBuffer> = ranges_by_buffer .into_iter() .filter_map(|(entity_id, (buffer, ranges))| { let snapshot = snapshots.get(&entity_id)?; let project_path = paths_by_buffer.get(&entity_id)?; let assembled = assemble_excerpt_ranges(snapshot, ranges); let root_name = worktree_root_names.get(&project_path.worktree_id)?; let path: Arc<Path> = Path::new(&format!( "{}/{}", root_name, project_path.path.as_unix_str()