oak-team/oak-gpui
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

zeta2: Parse imports via Tree-sitter queries + improve zeta retrieval-stats (#39735)

Browse Source 
Release Notes:

- N/A

---------

Co-authored-by: Max <max@zed.dev>
Co-authored-by: Agus <agus@zed.dev>
Co-authored-by: Oleksiy <oleksiy@zed.dev>
This commit is contained in:
Michael Sloan
2025-10-08 12:04:06 -06:00
committed by GitHub
co-authored by Max Agus Oleksiy
parent 5fd187769d
commit bcef3b5010
35 changed files with 3076 additions and 508 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/edit_prediction_context/src/declaration.rs
+79 -2
View File
@@ -1,9 +1,11 @@
use language::LanguageId;
use language::{Language, LanguageId};
use project::ProjectEntryId;
use std::borrow::Cow;
use std::ops::Range;
use std::sync::Arc;
use std::{borrow::Cow, path::Path};
use text::{Bias, BufferId, Rope};
use util::paths::{path_ends_with, strip_path_suffix};
use util::rel_path::RelPath;
use crate::outline::OutlineDeclaration;
@@ -22,12 +24,14 @@ pub enum Declaration {
File {
project_entry_id: ProjectEntryId,
declaration: FileDeclaration,
cached_path: CachedDeclarationPath,
},
Buffer {
project_entry_id: ProjectEntryId,
buffer_id: BufferId,
rope: Rope,
declaration: BufferDeclaration,
cached_path: CachedDeclarationPath,
},
}
@@ -73,6 +77,13 @@ impl Declaration {
}
}
pub fn cached_path(&self) -> &CachedDeclarationPath {
match self {
Declaration::File { cached_path, .. } => cached_path,
Declaration::Buffer { cached_path, .. } => cached_path,
}
}
pub fn item_range(&self) -> Range<usize> {
match self {
Declaration::File { declaration, .. } => declaration.item_range.clone(),
@@ -235,3 +246,69 @@ impl BufferDeclaration {
}
}
}
#[derive(Debug, Clone)]
pub struct CachedDeclarationPath {
pub worktree_abs_path: Arc<Path>,
pub rel_path: Arc<RelPath>,
/// The relative path of the file, possibly stripped according to `import_path_strip_regex`.
pub rel_path_after_regex_stripping: Arc<RelPath>,
}
impl CachedDeclarationPath {
pub fn new(
worktree_abs_path: Arc<Path>,
path: &Arc<RelPath>,
language: Option<&Arc<Language>>,
) -> Self {
let rel_path = path.clone();
let rel_path_after_regex_stripping = if let Some(language) = language
&& let Some(strip_regex) = language.config().import_path_strip_regex.as_ref()
&& let Ok(stripped) = RelPath::unix(&Path::new(
strip_regex.replace_all(rel_path.as_unix_str(), "").as_ref(),
)) {
Arc::from(stripped)
} else {
rel_path.clone()
};
CachedDeclarationPath {
worktree_abs_path,
rel_path,
rel_path_after_regex_stripping,
}
}
#[cfg(test)]
pub fn new_for_test(worktree_abs_path: &str, rel_path: &str) -> Self {
let rel_path: Arc<RelPath> = util::rel_path::rel_path(rel_path).into();
CachedDeclarationPath {
worktree_abs_path: std::path::PathBuf::from(worktree_abs_path).into(),
rel_path_after_regex_stripping: rel_path.clone(),
rel_path,
}
}
pub fn ends_with_posix_path(&self, path: &Path) -> bool {
if path.as_os_str().len() <= self.rel_path_after_regex_stripping.as_unix_str().len() {
path_ends_with(self.rel_path_after_regex_stripping.as_std_path(), path)
} else {
if let Some(remaining) =
strip_path_suffix(path, self.rel_path_after_regex_stripping.as_std_path())
{
path_ends_with(&self.worktree_abs_path, remaining)
} else {
false
}
}
}
pub fn equals_absolute_path(&self, path: &Path) -> bool {
if let Some(remaining) =
strip_path_suffix(path, &self.rel_path_after_regex_stripping.as_std_path())
{
self.worktree_abs_path.as_ref() == remaining
} else {
false
}
}
}
crates/edit_prediction_context/src/declaration_scoring.rs
+338 -111
View File
@@ -1,15 +1,15 @@
use cloud_llm_client::predict_edits_v3::DeclarationScoreComponents;
use collections::HashMap;
use itertools::Itertools as _;
use language::BufferSnapshot;
use ordered_float::OrderedFloat;
use serde::Serialize;
use std::{cmp::Reverse, ops::Range};
use std::{cmp::Reverse, ops::Range, path::Path, sync::Arc};
use strum::EnumIter;
use text::{Point, ToPoint};
use crate::{
Declaration, EditPredictionExcerpt, Identifier,
CachedDeclarationPath, Declaration, EditPredictionExcerpt, Identifier,
imports::{Import, Imports, Module},
reference::{Reference, ReferenceRegion},
syntax_index::SyntaxIndexState,
text_similarity::{Occurrences, jaccard_similarity, weighted_overlap_coefficient},
@@ -17,12 +17,17 @@ use crate::{
const MAX_IDENTIFIER_DECLARATION_COUNT: usize = 16;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct EditPredictionScoreOptions {
pub omit_excerpt_overlaps: bool,
}
#[derive(Clone, Debug)]
pub struct ScoredDeclaration {
/// identifier used by the local reference
pub identifier: Identifier,
pub declaration: Declaration,
pub score_components: DeclarationScoreComponents,
pub scores: DeclarationScores,
pub components: DeclarationScoreComponents,
}
#[derive(EnumIter, Clone, Copy, PartialEq, Eq, Hash, Debug)]
@@ -31,12 +36,55 @@ pub enum DeclarationStyle {
Declaration,
}
#[derive(Clone, Debug, Serialize, Default)]
pub struct DeclarationScores {
pub signature: f32,
pub declaration: f32,
pub retrieval: f32,
}
impl ScoredDeclaration {
/// Returns the score for this declaration with the specified style.
pub fn score(&self, style: DeclarationStyle) -> f32 {
// TODO: handle truncation
// Score related to how likely this is the correct declaration, range 0 to 1
let retrieval = self.retrieval_score();
// Score related to the distance between the reference and cursor, range 0 to 1
let distance_score = if self.components.is_referenced_nearby {
1.0 / (1.0 + self.components.reference_line_distance as f32 / 10.0).powf(2.0)
} else {
// same score as ~14 lines away, rationale is to not overly penalize references from parent signatures
0.5
};
// For now instead of linear combination, the scores are just multiplied together.
let combined_score = 10.0 * retrieval * distance_score;
match style {
DeclarationStyle::Signature => self.scores.signature,
DeclarationStyle::Declaration => self.scores.declaration,
DeclarationStyle::Signature => {
combined_score * self.components.excerpt_vs_signature_weighted_overlap
}
DeclarationStyle::Declaration => {
2.0 * combined_score * self.components.excerpt_vs_item_weighted_overlap
}
}
}
pub fn retrieval_score(&self) -> f32 {
if self.components.is_same_file {
10.0 / self.components.same_file_declaration_count as f32
} else if self.components.path_import_match_count > 0 {
3.0
} else if self.components.wildcard_path_import_match_count > 0 {
1.0
} else if self.components.normalized_import_similarity > 0.0 {
self.components.normalized_import_similarity
} else if self.components.normalized_wildcard_import_similarity > 0.0 {
0.5 * self.components.normalized_wildcard_import_similarity
} else {
1.0 / self.components.declaration_count as f32
}
}
@@ -54,100 +102,215 @@ impl ScoredDeclaration {
}
pub fn score_density(&self, style: DeclarationStyle) -> f32 {
self.score(style) / (self.size(style)) as f32
self.score(style) / self.size(style) as f32
}
}
pub fn scored_declarations(
options: &EditPredictionScoreOptions,
index: &SyntaxIndexState,
excerpt: &EditPredictionExcerpt,
excerpt_occurrences: &Occurrences,
adjacent_occurrences: &Occurrences,
imports: &Imports,
identifier_to_references: HashMap<Identifier, Vec<Reference>>,
cursor_offset: usize,
current_buffer: &BufferSnapshot,
) -> Vec<ScoredDeclaration> {
let cursor_point = cursor_offset.to_point(&current_buffer);
let mut wildcard_import_occurrences = Vec::new();
let mut wildcard_import_paths = Vec::new();
for wildcard_import in imports.wildcard_modules.iter() {
match wildcard_import {
Module::Namespace(namespace) => {
wildcard_import_occurrences.push(namespace.occurrences())
}
Module::SourceExact(path) => wildcard_import_paths.push(path),
Module::SourceFuzzy(path) => {
wildcard_import_occurrences.push(Occurrences::from_path(&path))
}
}
}
let mut declarations = identifier_to_references
.into_iter()
.flat_map(|(identifier, references)| {
let declarations =
index.declarations_for_identifier::<MAX_IDENTIFIER_DECLARATION_COUNT>(&identifier);
let mut import_occurrences = Vec::new();
let mut import_paths = Vec::new();
let mut found_external_identifier: Option<&Identifier> = None;
if let Some(imports) = imports.identifier_to_imports.get(&identifier) {
// only use alias when it's the only import, could be generalized if some language
// has overlapping aliases
//
// TODO: when an aliased declaration is included in the prompt, should include the
// aliasing in the prompt.
//
// TODO: For SourceFuzzy consider having componentwise comparison that pays
// attention to ordering.
if let [
Import::Alias {
module,
external_identifier,
},
] = imports.as_slice()
{
match module {
Module::Namespace(namespace) => {
import_occurrences.push(namespace.occurrences())
}
Module::SourceExact(path) => import_paths.push(path),
Module::SourceFuzzy(path) => {
import_occurrences.push(Occurrences::from_path(&path))
}
}
found_external_identifier = Some(&external_identifier);
} else {
for import in imports {
match import {
Import::Direct { module } => match module {
Module::Namespace(namespace) => {
import_occurrences.push(namespace.occurrences())
}
Module::SourceExact(path) => import_paths.push(path),
Module::SourceFuzzy(path) => {
import_occurrences.push(Occurrences::from_path(&path))
}
},
Import::Alias { .. } => {}
}
}
}
}
let identifier_to_lookup = found_external_identifier.unwrap_or(&identifier);
// TODO: update this to be able to return more declarations? Especially if there is the
// ability to quickly filter a large list (based on imports)
let declarations = index
.declarations_for_identifier::<MAX_IDENTIFIER_DECLARATION_COUNT>(
&identifier_to_lookup,
);
let declaration_count = declarations.len();
declarations
.into_iter()
.filter_map(|(declaration_id, declaration)| match declaration {
if declaration_count == 0 {
return Vec::new();
}
// TODO: option to filter out other candidates when same file / import match
let mut checked_declarations = Vec::new();
for (declaration_id, declaration) in declarations {
match declaration {
Declaration::Buffer {
buffer_id,
declaration: buffer_declaration,
..
} => {
let is_same_file = buffer_id == &current_buffer.remote_id();
if is_same_file {
let overlaps_excerpt =
if buffer_id == &current_buffer.remote_id() {
let already_included_in_prompt =
range_intersection(&buffer_declaration.item_range, &excerpt.range)
.is_some();
if overlaps_excerpt
|| excerpt
.parent_declarations
.iter()
.any(|(excerpt_parent, _)| excerpt_parent == &declaration_id)
{
None
} else {
.is_some()
|| excerpt.parent_declarations.iter().any(
|(excerpt_parent, _)| excerpt_parent == &declaration_id,
);
if !options.omit_excerpt_overlaps || !already_included_in_prompt {
let declaration_line = buffer_declaration
.item_range
.start
.to_point(current_buffer)
.row;
Some((
true,
(cursor_point.row as i32 - declaration_line as i32)
.unsigned_abs(),
let declaration_line_distance = (cursor_point.row as i32
- declaration_line as i32)
.unsigned_abs();
checked_declarations.push(CheckedDeclaration {
declaration,
))
same_file_line_distance: Some(declaration_line_distance),
path_import_match_count: 0,
wildcard_path_import_match_count: 0,
});
}
continue;
} else {
Some((false, u32::MAX, declaration))
}
}
Declaration::File { .. } => {
// We can assume that a file declaration is in a different file,
// because the current one must be open
Some((false, u32::MAX, declaration))
}
})
.sorted_by_key(|&(_, distance, _)| distance)
.enumerate()
.map(
|(
declaration_line_distance_rank,
(is_same_file, declaration_line_distance, declaration),
)| {
let same_file_declaration_count = index.file_declaration_count(declaration);
Declaration::File { .. } => {}
}
let declaration_path = declaration.cached_path();
let path_import_match_count = import_paths
.iter()
.filter(|import_path| {
declaration_path_matches_import(&declaration_path, import_path)
})
.count();
let wildcard_path_import_match_count = wildcard_import_paths
.iter()
.filter(|import_path| {
declaration_path_matches_import(&declaration_path, import_path)
})
.count();
checked_declarations.push(CheckedDeclaration {
declaration,
same_file_line_distance: None,
path_import_match_count,
wildcard_path_import_match_count,
});
}
score_declaration(
&identifier,
&references,
declaration.clone(),
is_same_file,
declaration_line_distance,
declaration_line_distance_rank,
same_file_declaration_count,
declaration_count,
&excerpt_occurrences,
&adjacent_occurrences,
cursor_point,
current_buffer,
)
},
)
.collect::<Vec<_>>()
let mut max_import_similarity = 0.0;
let mut max_wildcard_import_similarity = 0.0;
let mut scored_declarations_for_identifier = checked_declarations
.into_iter()
.map(|checked_declaration| {
let same_file_declaration_count =
index.file_declaration_count(checked_declaration.declaration);
let declaration = score_declaration(
&identifier,
&references,
checked_declaration,
same_file_declaration_count,
declaration_count,
&excerpt_occurrences,
&adjacent_occurrences,
&import_occurrences,
&wildcard_import_occurrences,
cursor_point,
current_buffer,
);
if declaration.components.import_similarity > max_import_similarity {
max_import_similarity = declaration.components.import_similarity;
}
if declaration.components.wildcard_import_similarity
> max_wildcard_import_similarity
{
max_wildcard_import_similarity =
declaration.components.wildcard_import_similarity;
}
declaration
})
.collect::<Vec<_>>();
if max_import_similarity > 0.0 || max_wildcard_import_similarity > 0.0 {
for declaration in scored_declarations_for_identifier.iter_mut() {
if max_import_similarity > 0.0 {
declaration.components.max_import_similarity = max_import_similarity;
declaration.components.normalized_import_similarity =
declaration.components.import_similarity / max_import_similarity;
}
if max_wildcard_import_similarity > 0.0 {
declaration.components.normalized_wildcard_import_similarity =
declaration.components.wildcard_import_similarity
/ max_wildcard_import_similarity;
}
}
}
scored_declarations_for_identifier
})
.flatten()
.collect::<Vec<_>>();
declarations.sort_unstable_by_key(|declaration| {
@@ -160,6 +323,24 @@ pub fn scored_declarations(
declarations
}
struct CheckedDeclaration<'a> {
declaration: &'a Declaration,
same_file_line_distance: Option<u32>,
path_import_match_count: usize,
wildcard_path_import_match_count: usize,
}
fn declaration_path_matches_import(
declaration_path: &CachedDeclarationPath,
import_path: &Arc<Path>,
) -> bool {
if import_path.is_absolute() {
declaration_path.equals_absolute_path(import_path)
} else {
declaration_path.ends_with_posix_path(import_path)
}
}
fn range_intersection<T: Ord + Clone>(a: &Range<T>, b: &Range<T>) -> Option<Range<T>> {
let start = a.start.clone().max(b.start.clone());
let end = a.end.clone().min(b.end.clone());
@@ -173,17 +354,23 @@ fn range_intersection<T: Ord + Clone>(a: &Range<T>, b: &Range<T>) -> Option<Rang
fn score_declaration(
identifier: &Identifier,
references: &[Reference],
declaration: Declaration,
is_same_file: bool,
declaration_line_distance: u32,
declaration_line_distance_rank: usize,
checked_declaration: CheckedDeclaration,
same_file_declaration_count: usize,
declaration_count: usize,
excerpt_occurrences: &Occurrences,
adjacent_occurrences: &Occurrences,
import_occurrences: &[Occurrences],
wildcard_import_occurrences: &[Occurrences],
cursor: Point,
current_buffer: &BufferSnapshot,
) -> Option<ScoredDeclaration> {
) -> ScoredDeclaration {
let CheckedDeclaration {
declaration,
same_file_line_distance,
path_import_match_count,
wildcard_path_import_match_count,
} = checked_declaration;
let is_referenced_nearby = references
.iter()
.any(|r| r.region == ReferenceRegion::Nearby);
@@ -200,6 +387,9 @@ fn score_declaration(
.min()
.unwrap();
let is_same_file = same_file_line_distance.is_some();
let declaration_line_distance = same_file_line_distance.unwrap_or(u32::MAX);
let item_source_occurrences = Occurrences::within_string(&declaration.item_text().0);
let item_signature_occurrences = Occurrences::within_string(&declaration.signature_text().0);
let excerpt_vs_item_jaccard = jaccard_similarity(excerpt_occurrences, &item_source_occurrences);
@@ -219,6 +409,37 @@ fn score_declaration(
let adjacent_vs_signature_weighted_overlap =
weighted_overlap_coefficient(adjacent_occurrences, &item_signature_occurrences);
let mut import_similarity = 0f32;
let mut wildcard_import_similarity = 0f32;
if !import_occurrences.is_empty() || !wildcard_import_occurrences.is_empty() {
let cached_path = declaration.cached_path();
let path_occurrences = Occurrences::from_worktree_path(
cached_path
.worktree_abs_path
.file_name()
.map(|f| f.to_string_lossy()),
&cached_path.rel_path,
);
import_similarity = import_occurrences
.iter()
.map(|namespace_occurrences| {
OrderedFloat(jaccard_similarity(namespace_occurrences, &path_occurrences))
})
.max()
.map(|similarity| similarity.into_inner())
.unwrap_or_default();
// TODO: Consider something other than max
wildcard_import_similarity = wildcard_import_occurrences
.iter()
.map(|namespace_occurrences| {
OrderedFloat(jaccard_similarity(namespace_occurrences, &path_occurrences))
})
.max()
.map(|similarity| similarity.into_inner())
.unwrap_or_default();
}
// TODO: Consider adding declaration_file_count
let score_components = DeclarationScoreComponents {
is_same_file,
@@ -226,7 +447,6 @@ fn score_declaration(
is_referenced_in_breadcrumb,
reference_line_distance,
declaration_line_distance,
declaration_line_distance_rank,
reference_count,
same_file_declaration_count,
declaration_count,
@@ -238,52 +458,59 @@ fn score_declaration(
excerpt_vs_signature_weighted_overlap,
adjacent_vs_item_weighted_overlap,
adjacent_vs_signature_weighted_overlap,
path_import_match_count,
wildcard_path_import_match_count,
import_similarity,
max_import_similarity: 0.0,
normalized_import_similarity: 0.0,
wildcard_import_similarity,
normalized_wildcard_import_similarity: 0.0,
};
Some(ScoredDeclaration {
ScoredDeclaration {
identifier: identifier.clone(),
declaration: declaration,
scores: DeclarationScores::score(&score_components),
score_components,
})
declaration: declaration.clone(),
components: score_components,
}
}
#[derive(Clone, Debug, Serialize)]
pub struct DeclarationScores {
pub signature: f32,
pub declaration: f32,
pub retrieval: f32,
}
#[cfg(test)]
mod test {
use super::*;
#[test]
fn test_declaration_path_matches() {
let declaration_path =
CachedDeclarationPath::new_for_test("/home/user/project", "src/maths.ts");
assert!(declaration_path_matches_import(
&declaration_path,
&Path::new("maths.ts").into()
));
assert!(declaration_path_matches_import(
&declaration_path,
&Path::new("project/src/maths.ts").into()
));
assert!(declaration_path_matches_import(
&declaration_path,
&Path::new("user/project/src/maths.ts").into()
));
assert!(declaration_path_matches_import(
&declaration_path,
&Path::new("/home/user/project/src/maths.ts").into()
));
assert!(!declaration_path_matches_import(
&declaration_path,
&Path::new("other.ts").into()
));
impl DeclarationScores {
fn score(components: &DeclarationScoreComponents) -> DeclarationScores {
// TODO: handle truncation
// Score related to how likely this is the correct declaration, range 0 to 1
let retrieval = if components.is_same_file {
// TODO: use declaration_line_distance_rank
1.0 / components.same_file_declaration_count as f32
} else {
1.0 / components.declaration_count as f32
};
// Score related to the distance between the reference and cursor, range 0 to 1
let distance_score = if components.is_referenced_nearby {
1.0 / (1.0 + components.reference_line_distance as f32 / 10.0).powf(2.0)
} else {
// same score as ~14 lines away, rationale is to not overly penalize references from parent signatures
0.5
};
// For now instead of linear combination, the scores are just multiplied together.
let combined_score = 10.0 * retrieval * distance_score;
DeclarationScores {
signature: combined_score * components.excerpt_vs_signature_weighted_overlap,
// declaration score gets boosted both by being multiplied by 2 and by there being more
// weighted overlap.
declaration: 2.0 * combined_score * components.excerpt_vs_item_weighted_overlap,
retrieval,
}
assert!(!declaration_path_matches_import(
&declaration_path,
&Path::new("/home/user/project/src/other.ts").into()
));
}
}
crates/edit_prediction_context/src/edit_prediction_context.rs
+51 -13
View File
@@ -1,12 +1,13 @@
mod declaration;
mod declaration_scoring;
mod excerpt;
mod imports;
mod outline;
mod reference;
mod syntax_index;
pub mod text_similarity;
use std::sync::Arc;
use std::{path::Path, sync::Arc};
use collections::HashMap;
use gpui::{App, AppContext as _, Entity, Task};
@@ -16,9 +17,17 @@ use text::{Point, ToOffset as _};
pub use declaration::*;
pub use declaration_scoring::*;
pub use excerpt::*;
pub use imports::*;
pub use reference::*;
pub use syntax_index::*;
#[derive(Clone, Debug, PartialEq)]
pub struct EditPredictionContextOptions {
pub use_imports: bool,
pub excerpt: EditPredictionExcerptOptions,
pub score: EditPredictionScoreOptions,
}
#[derive(Clone, Debug)]
pub struct EditPredictionContext {
pub excerpt: EditPredictionExcerpt,
@@ -31,21 +40,34 @@ impl EditPredictionContext {
pub fn gather_context_in_background(
cursor_point: Point,
buffer: BufferSnapshot,
excerpt_options: EditPredictionExcerptOptions,
options: EditPredictionContextOptions,
syntax_index: Option<Entity<SyntaxIndex>>,
cx: &mut App,
) -> Task<Option<Self>> {
let parent_abs_path = project::File::from_dyn(buffer.file()).and_then(|f| {
let mut path = f.worktree.read(cx).absolutize(&f.path);
if path.pop() { Some(path) } else { None }
});
if let Some(syntax_index) = syntax_index {
let index_state =
syntax_index.read_with(cx, |index, _cx| Arc::downgrade(index.state()));
cx.background_spawn(async move {
let parent_abs_path = parent_abs_path.as_deref();
let index_state = index_state.upgrade()?;
let index_state = index_state.lock().await;
Self::gather_context(cursor_point, &buffer, &excerpt_options, Some(&index_state))
Self::gather_context(
cursor_point,
&buffer,
parent_abs_path,
&options,
Some(&index_state),
)
})
} else {
cx.background_spawn(async move {
Self::gather_context(cursor_point, &buffer, &excerpt_options, None)
let parent_abs_path = parent_abs_path.as_deref();
Self::gather_context(cursor_point, &buffer, parent_abs_path, &options, None)
})
}
}
@@ -53,13 +75,20 @@ impl EditPredictionContext {
pub fn gather_context(
cursor_point: Point,
buffer: &BufferSnapshot,
excerpt_options: &EditPredictionExcerptOptions,
parent_abs_path: Option<&Path>,
options: &EditPredictionContextOptions,
index_state: Option<&SyntaxIndexState>,
) -> Option<Self> {
let imports = if options.use_imports {
Imports::gather(&buffer, parent_abs_path)
} else {
Imports::default()
};
Self::gather_context_with_references_fn(
cursor_point,
buffer,
excerpt_options,
&imports,
options,
index_state,
references_in_excerpt,
)
@@ -68,7 +97,8 @@ impl EditPredictionContext {
pub fn gather_context_with_references_fn(
cursor_point: Point,
buffer: &BufferSnapshot,
excerpt_options: &EditPredictionExcerptOptions,
imports: &Imports,
options: &EditPredictionContextOptions,
index_state: Option<&SyntaxIndexState>,
get_references: impl FnOnce(
&EditPredictionExcerpt,
@@ -79,7 +109,7 @@ impl EditPredictionContext {
let excerpt = EditPredictionExcerpt::select_from_buffer(
cursor_point,
buffer,
excerpt_options,
&options.excerpt,
index_state,
)?;
let excerpt_text = excerpt.text(buffer);
@@ -101,10 +131,12 @@ impl EditPredictionContext {
let references = get_references(&excerpt, &excerpt_text, buffer);
scored_declarations(
&options.score,
&index_state,
&excerpt,
&excerpt_occurrences,
&adjacent_occurrences,
&imports,
references,
cursor_offset_in_file,
buffer,
@@ -160,12 +192,18 @@ mod tests {
EditPredictionContext::gather_context_in_background(
cursor_point,
buffer_snapshot,
EditPredictionExcerptOptions {
max_bytes: 60,
min_bytes: 10,
target_before_cursor_over_total_bytes: 0.5,
EditPredictionContextOptions {
use_imports: true,
excerpt: EditPredictionExcerptOptions {
max_bytes: 60,
min_bytes: 10,
target_before_cursor_over_total_bytes: 0.5,
},
score: EditPredictionScoreOptions {
omit_excerpt_overlaps: true,
},
},
Some(index),
Some(index.clone()),
cx,
)
})
crates/edit_prediction_context/src/imports.rs
+1319
View File
File diff suppressed because it is too large Load Diff
crates/edit_prediction_context/src/syntax_index.rs
+101 -73
View File
@@ -5,6 +5,7 @@ use futures::lock::Mutex;
use futures::{FutureExt as _, StreamExt, future};
use gpui::{App, AppContext as _, AsyncApp, Context, Entity, Task, WeakEntity};
use itertools::Itertools;
use language::{Buffer, BufferEvent};
use postage::stream::Stream as _;
use project::buffer_store::{BufferStore, BufferStoreEvent};
@@ -17,6 +18,7 @@ use std::sync::Arc;
use text::BufferId;
use util::{RangeExt as _, debug_panic, some_or_debug_panic};
use crate::CachedDeclarationPath;
use crate::declaration::{
BufferDeclaration, Declaration, DeclarationId, FileDeclaration, Identifier,
};
@@ -28,6 +30,8 @@ use crate::outline::declarations_in_buffer;
// `buffer_declarations_containing_range` assumes that the index is always immediately up to date.
//
// * Add a per language configuration for skipping indexing.
//
// * Handle tsx / ts / js referencing each-other
// Potential future improvements:
//
@@ -61,6 +65,7 @@ pub struct SyntaxIndex {
state: Arc<Mutex<SyntaxIndexState>>,
project: WeakEntity<Project>,
initial_file_indexing_done_rx: postage::watch::Receiver<bool>,
_file_indexing_task: Option<Task<()>>,
}
pub struct SyntaxIndexState {
@@ -70,7 +75,6 @@ pub struct SyntaxIndexState {
buffers: HashMap<BufferId, BufferState>,
dirty_files: HashMap<ProjectEntryId, ProjectPath>,
dirty_files_tx: mpsc::Sender<()>,
_file_indexing_task: Option<Task<()>>,
}
#[derive(Debug, Default)]
@@ -102,12 +106,12 @@ impl SyntaxIndex {
buffers: HashMap::default(),
dirty_files: HashMap::default(),
dirty_files_tx,
_file_indexing_task: None,
};
let this = Self {
let mut this = Self {
project: project.downgrade(),
state: Arc::new(Mutex::new(initial_state)),
initial_file_indexing_done_rx,
_file_indexing_task: None,
};
let worktree_store = project.read(cx).worktree_store();
@@ -116,75 +120,77 @@ impl SyntaxIndex {
.worktrees()
.map(|w| w.read(cx).snapshot())
.collect::<Vec<_>>();
if !initial_worktree_snapshots.is_empty() {
this.state.try_lock().unwrap()._file_indexing_task =
Some(cx.spawn(async move |this, cx| {
let snapshots_file_count = initial_worktree_snapshots
.iter()
.map(|worktree| worktree.file_count())
.sum::<usize>();
let chunk_size = snapshots_file_count.div_ceil(file_indexing_parallelism);
let chunk_count = snapshots_file_count.div_ceil(chunk_size);
let file_chunks = initial_worktree_snapshots
.iter()
.flat_map(|worktree| {
let worktree_id = worktree.id();
worktree.files(false, 0).map(move |entry| {
(
entry.id,
ProjectPath {
worktree_id,
path: entry.path.clone(),
},
)
})
this._file_indexing_task = Some(cx.spawn(async move |this, cx| {
let snapshots_file_count = initial_worktree_snapshots
.iter()
.map(|worktree| worktree.file_count())
.sum::<usize>();
if snapshots_file_count > 0 {
let chunk_size = snapshots_file_count.div_ceil(file_indexing_parallelism);
let chunk_count = snapshots_file_count.div_ceil(chunk_size);
let file_chunks = initial_worktree_snapshots
.iter()
.flat_map(|worktree| {
let worktree_id = worktree.id();
worktree.files(false, 0).map(move |entry| {
(
entry.id,
ProjectPath {
worktree_id,
path: entry.path.clone(),
},
)
})
.chunks(chunk_size);
})
.chunks(chunk_size);
let mut tasks = Vec::with_capacity(chunk_count);
for chunk in file_chunks.into_iter() {
tasks.push(Self::update_dirty_files(
&this,
chunk.into_iter().collect(),
cx.clone(),
));
}
futures::future::join_all(tasks).await;
let mut tasks = Vec::with_capacity(chunk_count);
for chunk in file_chunks.into_iter() {
tasks.push(Self::update_dirty_files(
&this,
chunk.into_iter().collect(),
cx.clone(),
));
}
futures::future::join_all(tasks).await;
log::info!("Finished initial file indexing");
}
log::info!("Finished initial file indexing");
*initial_file_indexing_done_tx.borrow_mut() = true;
*initial_file_indexing_done_tx.borrow_mut() = true;
let Ok(state) = this.read_with(cx, |this, _cx| this.state.clone()) else {
return;
};
while dirty_files_rx.next().await.is_some() {
let mut state = state.lock().await;
let was_underused = state.dirty_files.capacity() > 255
&& state.dirty_files.len() * 8 < state.dirty_files.capacity();
let dirty_files = state.dirty_files.drain().collect::<Vec<_>>();
if was_underused {
state.dirty_files.shrink_to_fit();
}
drop(state);
if dirty_files.is_empty() {
continue;
}
let Ok(state) = this.read_with(cx, |this, _cx| Arc::downgrade(&this.state)) else {
return;
};
while dirty_files_rx.next().await.is_some() {
let Some(state) = state.upgrade() else {
return;
};
let mut state = state.lock().await;
let was_underused = state.dirty_files.capacity() > 255
&& state.dirty_files.len() * 8 < state.dirty_files.capacity();
let dirty_files = state.dirty_files.drain().collect::<Vec<_>>();
if was_underused {
state.dirty_files.shrink_to_fit();
}
drop(state);
if dirty_files.is_empty() {
continue;
}
let chunk_size = dirty_files.len().div_ceil(file_indexing_parallelism);
let chunk_count = dirty_files.len().div_ceil(chunk_size);
let mut tasks = Vec::with_capacity(chunk_count);
let chunks = dirty_files.into_iter().chunks(chunk_size);
for chunk in chunks.into_iter() {
tasks.push(Self::update_dirty_files(
&this,
chunk.into_iter().collect(),
cx.clone(),
));
}
futures::future::join_all(tasks).await;
}
}));
}
let chunk_size = dirty_files.len().div_ceil(file_indexing_parallelism);
let chunk_count = dirty_files.len().div_ceil(chunk_size);
let mut tasks = Vec::with_capacity(chunk_count);
let chunks = dirty_files.into_iter().chunks(chunk_size);
for chunk in chunks.into_iter() {
tasks.push(Self::update_dirty_files(
&this,
chunk.into_iter().collect(),
cx.clone(),
));
}
futures::future::join_all(tasks).await;
}
}));
cx.subscribe(&worktree_store, Self::handle_worktree_store_event)
.detach();
@@ -364,7 +370,9 @@ impl SyntaxIndex {
cx: &mut Context<Self>,
) {
match event {
BufferEvent::Edited => self.update_buffer(buffer, cx),
BufferEvent::Edited |
// paths are cached and so should be updated
BufferEvent::FileHandleChanged => self.update_buffer(buffer, cx),
_ => {}
}
}
@@ -375,8 +383,16 @@ impl SyntaxIndex {
return;
}
let Some(project_entry_id) =
project::File::from_dyn(buffer.file()).and_then(|f| f.project_entry_id(cx))
let Some((project_entry_id, cached_path)) = project::File::from_dyn(buffer.file())
.and_then(|f| {
let project_entry_id = f.project_entry_id()?;
let cached_path = CachedDeclarationPath::new(
f.worktree.read(cx).abs_path(),
&f.path,
buffer.language(),
);
Some((project_entry_id, cached_path))
})
else {
return;
};
@@ -440,6 +456,7 @@ impl SyntaxIndex {
buffer_id,
declaration,
project_entry_id,
cached_path: cached_path.clone(),
});
new_ids.push(declaration_id);
@@ -507,13 +524,14 @@ impl SyntaxIndex {
let snapshot_task = worktree.update(cx, |worktree, cx| {
let load_task = worktree.load_file(&project_path.path, cx);
let worktree_abs_path = worktree.abs_path();
cx.spawn(async move |_this, cx| {
let loaded_file = load_task.await?;
let language = language.await?;
let buffer = cx.new(|cx| {
let mut buffer = Buffer::local(loaded_file.text, cx);
buffer.set_language(Some(language), cx);
buffer.set_language(Some(language.clone()), cx);
buffer
})?;
@@ -522,14 +540,22 @@ impl SyntaxIndex {
parse_status.changed().await?;
}
buffer.read_with(cx, |buffer, _cx| buffer.snapshot())
let cached_path = CachedDeclarationPath::new(
worktree_abs_path,
&project_path.path,
Some(&language),
);
let snapshot = buffer.read_with(cx, |buffer, _cx| buffer.snapshot())?;
anyhow::Ok((snapshot, cached_path))
})
});
let state = Arc::downgrade(&self.state);
cx.background_spawn(async move {
// TODO: How to handle errors?
let Ok(snapshot) = snapshot_task.await else {
let Ok((snapshot, cached_path)) = snapshot_task.await else {
return;
};
let rope = snapshot.as_rope();
@@ -567,6 +593,7 @@ impl SyntaxIndex {
let declaration_id = state.declarations.insert(Declaration::File {
project_entry_id: entry_id,
declaration,
cached_path: cached_path.clone(),
});
new_ids.push(declaration_id);
@@ -921,6 +948,7 @@ mod tests {
if let Declaration::File {
declaration,
project_entry_id: file,
..
} = declaration
{
assert_eq!(
crates/edit_prediction_context/src/text_similarity.rs
+44 -2
View File
@@ -1,9 +1,12 @@
use hashbrown::HashTable;
use regex::Regex;
use std::{
borrow::Cow,
hash::{Hash, Hasher as _},
path::Path,
sync::LazyLock,
};
use util::rel_path::RelPath;
use crate::reference::Reference;
@@ -45,19 +48,34 @@ impl Occurrences {
)
}
pub fn from_identifiers<'a>(identifiers: impl IntoIterator<Item = &'a str>) -> Self {
pub fn from_identifiers(identifiers: impl IntoIterator<Item = impl AsRef<str>>) -> Self {
let mut this = Self::default();
// TODO: Score matches that match case higher?
//
// TODO: Also include unsplit identifier?
for identifier in identifiers {
for identifier_part in split_identifier(identifier) {
for identifier_part in split_identifier(identifier.as_ref()) {
this.add_hash(fx_hash(&identifier_part.to_lowercase()));
}
}
this
}
pub fn from_worktree_path(worktree_name: Option<Cow<'_, str>>, rel_path: &RelPath) -> Self {
if let Some(worktree_name) = worktree_name {
Self::from_identifiers(
std::iter::once(worktree_name)
.chain(iter_path_without_extension(rel_path.as_std_path())),
)
} else {
Self::from_path(rel_path.as_std_path())
}
}
pub fn from_path(path: &Path) -> Self {
Self::from_identifiers(iter_path_without_extension(path))
}
fn add_hash(&mut self, hash: u64) {
self.table
.entry(
@@ -82,6 +100,15 @@ impl Occurrences {
}
}
fn iter_path_without_extension(path: &Path) -> impl Iterator<Item = Cow<'_, str>> {
let last_component: Option<Cow<'_, str>> = path.file_stem().map(|stem| stem.to_string_lossy());
let mut path_components = path.components();
path_components.next_back();
path_components
.map(|component| component.as_os_str().to_string_lossy())
.chain(last_component)
}
pub fn fx_hash<T: Hash + ?Sized>(data: &T) -> u64 {
let mut hasher = collections::FxHasher::default();
data.hash(&mut hasher);
@@ -269,4 +296,19 @@ mod test {
// the smaller set, 10.
assert_eq!(weighted_overlap_coefficient(&set_a, &set_b), 7.0 / 10.0);
}
#[test]
fn test_iter_path_without_extension() {
let mut iter = iter_path_without_extension(Path::new(""));
assert_eq!(iter.next(), None);
let iter = iter_path_without_extension(Path::new("foo"));
assert_eq!(iter.collect::<Vec<_>>(), ["foo"]);
let iter = iter_path_without_extension(Path::new("foo/bar.txt"));
assert_eq!(iter.collect::<Vec<_>>(), ["foo", "bar"]);
let iter = iter_path_without_extension(Path::new("foo/bar/baz.txt"));
assert_eq!(iter.collect::<Vec<_>>(), ["foo", "bar", "baz"]);
}
}