682 lines · 20.4 KB · Rust
RawDownload
use crate::editable_context::EditHistoryContextEntry;use anyhow::{Context as _, Result, bail};use gpui::{AppContext as _, AsyncApp, Entity};use language::{Buffer, Point, ToPoint as _};use project::Project;use std::{    cmp::Ordering,    collections::{HashMap, HashSet},    fs,    ops::Range,    path::{Path, PathBuf},    time::Instant,};use text::Anchor;use util::command::new_command; const BM25_CONTEXT_QUERY_LINE_COUNT: u32 = 20;const BM25_CONTEXT_EDIT_HISTORY_QUERY_ENTRY_COUNT: usize = 8;const BM25_CONTEXT_CHUNK_LINE_COUNT: usize = 40;const BM25_CONTEXT_CHUNK_OVERLAP_LINE_COUNT: usize = 10;const BM25_CONTEXT_CHUNK_COUNT: usize = 12;const BM25_CONTEXT_MAX_CHUNKS_PER_FILE: usize = 3;const BM25_CONTEXT_MAX_FILE_BYTES: u64 = 1_000_000;const BM25_K1: f64 = 1.2;const BM25_B: f64 = 0.75; pub(super) struct Bm25ContextCandidate {    pub path: PathBuf,    pub row_range: Range<u32>,    pub order: usize,} pub async fn collect_bm25_context(    project: Entity<Project>,    active_buffer: Entity<Buffer>,    cursor_position: Anchor,    edit_history: &[EditHistoryContextEntry],    next_order: usize,    cx: &mut AsyncApp,) -> Vec<Bm25ContextCandidate> {    let Some(query) = build_query(&project, &active_buffer, cursor_position, edit_history, cx)    else {        return Vec::new();    };     let result = cx        .background_spawn(async move { collect_bm25_context_from_disk(query, next_order).await })        .await;     match result {        Ok(context) => context,        Err(error) => {            log::debug!("failed to collect BM25 context: {error:#}");            Vec::new()        }    }} struct Bm25ContextQuery {    worktree_abs_path: PathBuf,    worktree_root_name: String,    active_path: String,    cursor_excerpt: String,    edit_history_excerpts: Vec<String>,} fn build_query(    project: &Entity<Project>,    active_buffer: &Entity<Buffer>,    cursor_position: Anchor,    edit_history: &[EditHistoryContextEntry],    cx: &mut AsyncApp,) -> Option<Bm25ContextQuery> {    let (worktree_abs_path, worktree_root_name, active_path, cursor_excerpt) = cx.update(|cx| {        let buffer = active_buffer.read(cx);        let file = buffer.file()?;        let project = project.read(cx);        if !project.is_local() {            return None;        }        let worktree = project.worktree_for_id(file.worktree_id(cx), cx)?;        let worktree = worktree.read(cx);        if !worktree.is_local() {            return None;        }         let snapshot = buffer.snapshot();        let range = expanded_anchor_range(&snapshot, cursor_position..cursor_position);        let cursor_excerpt = snapshot.text_for_range(range).collect::<String>();         Some((            worktree.abs_path(),            worktree.root_name().as_unix_str().to_string(),            file.path().as_unix_str().to_string(),            cursor_excerpt,        ))    })?;     let edit_history_excerpts = edit_history        .iter()        .take(BM25_CONTEXT_EDIT_HISTORY_QUERY_ENTRY_COUNT)        .map(|entry| {            entry.buffer.read_with(cx, |buffer, _cx| {                let snapshot = buffer.snapshot();                let range = expanded_anchor_range(&snapshot, entry.edited_range.clone());                snapshot.text_for_range(range).collect::<String>()            })        })        .collect();     Some(Bm25ContextQuery {        worktree_abs_path: worktree_abs_path.to_path_buf(),        worktree_root_name,        active_path,        cursor_excerpt,        edit_history_excerpts,    })} fn expanded_anchor_range(    snapshot: &language::BufferSnapshot,    range: Range<Anchor>,) -> Range<Anchor> {    let start = range.start.to_point(snapshot);    let end = range.end.to_point(snapshot);    let start_row = start.row.saturating_sub(BM25_CONTEXT_QUERY_LINE_COUNT);    let end_row = end        .row        .saturating_add(BM25_CONTEXT_QUERY_LINE_COUNT)        .min(snapshot.max_point().row);    let start = snapshot.anchor_before(Point::new(start_row, 0));    let end = snapshot.anchor_after(Point::new(end_row, snapshot.line_len(end_row)));    start..end} async fn collect_bm25_context_from_disk(    query: Bm25ContextQuery,    next_order: usize,) -> Result<Vec<Bm25ContextCandidate>> {    let query_terms = query_terms(&query);    if query_terms.is_empty() {        return Ok(Vec::new());    }     let started_at = Instant::now();    let index = Bm25Index::build(&query.worktree_abs_path).await?;    let elapsed = started_at.elapsed();    log::debug!(        "built BM25 context index: candidate_files:{}, indexed_files:{}, indexed_bytes:{}, chunks:{}, terms:{}, latency:{elapsed:?}",        index.stats.candidate_file_count,        index.stats.indexed_file_count,        index.stats.indexed_bytes,        index.stats.document_count,        index.stats.term_count,    );     let candidates = index.search(&query_terms, &query.worktree_root_name, next_order);    log::debug!("selected {} BM25 context chunks", candidates.len());    Ok(candidates)} fn query_terms(query: &Bm25ContextQuery) -> HashMap<String, f64> {    let mut terms = HashMap::new();    add_query_terms(&mut terms, &query.active_path, 3.0);    add_query_terms(&mut terms, &query.cursor_excerpt, 1.0);    for excerpt in &query.edit_history_excerpts {        add_query_terms(&mut terms, excerpt, 2.0);    }    terms} fn add_query_terms(terms: &mut HashMap<String, f64>, text: &str, weight: f64) {    for token in tokenize(text) {        *terms.entry(token).or_default() += weight;    }} struct Bm25Index {    documents: Vec<Document>,    document_frequencies: HashMap<String, usize>,    average_document_len: f64,    stats: Bm25IndexStats,} #[derive(Default)]struct Bm25IndexStats {    candidate_file_count: usize,    indexed_file_count: usize,    indexed_bytes: u64,    document_count: usize,    term_count: usize,} struct Document {    relative_path: PathBuf,    row_range: Range<u32>,    term_frequencies: HashMap<String, usize>,    len: usize,} struct ScoredDocument {    document_index: usize,    score: f64,} struct DocumentsForFile {    documents: Vec<Document>,    byte_len: u64,} impl Bm25Index {    async fn build(worktree_abs_path: &Path) -> Result<Self> {        let relative_paths = git_ls_files(worktree_abs_path).await?;        let mut stats = Bm25IndexStats {            candidate_file_count: relative_paths.len(),            ..Default::default()        };        let mut documents = Vec::new();        for relative_path in relative_paths {            let Some(documents_for_file) = documents_for_file(worktree_abs_path, relative_path)            else {                continue;            };             stats.indexed_file_count += 1;            stats.indexed_bytes += documents_for_file.byte_len;            documents.extend(documents_for_file.documents);        }         let mut document_frequencies = HashMap::new();        let mut total_document_len = 0;        for document in &documents {            total_document_len += document.len;            let mut seen_terms = HashSet::new();            for term in document.term_frequencies.keys() {                if seen_terms.insert(term) {                    *document_frequencies.entry(term.clone()).or_default() += 1;                }            }        }         let average_document_len = if documents.is_empty() {            0.0        } else {            total_document_len as f64 / documents.len() as f64        };        stats.document_count = documents.len();        stats.term_count = document_frequencies.len();         Ok(Self {            documents,            document_frequencies,            average_document_len,            stats,        })    }     fn search(        &self,        query_terms: &HashMap<String, f64>,        worktree_root_name: &str,        next_order: usize,    ) -> Vec<Bm25ContextCandidate> {        if self.documents.is_empty() || self.average_document_len == 0.0 {            return Vec::new();        }         let mut scored_documents = self            .documents            .iter()            .enumerate()            .filter_map(|(document_index, document)| {                let score = self.score_document(document, query_terms);                (score > 0.0).then_some(ScoredDocument {                    document_index,                    score,                })            })            .collect::<Vec<_>>();         scored_documents.sort_by(|left, right| {            right                .score                .partial_cmp(&left.score)                .unwrap_or(Ordering::Equal)                .then_with(|| {                    self.documents[left.document_index]                        .relative_path                        .cmp(&self.documents[right.document_index].relative_path)                })                .then_with(|| {                    self.documents[left.document_index]                        .row_range                        .start                        .cmp(&self.documents[right.document_index].row_range.start)                })        });         let mut selected_documents = Vec::new();        let mut chunks_per_file = HashMap::<PathBuf, usize>::new();        for scored_document in scored_documents {            let document = &self.documents[scored_document.document_index];            let chunk_count = chunks_per_file                .entry(document.relative_path.clone())                .or_default();            if *chunk_count >= BM25_CONTEXT_MAX_CHUNKS_PER_FILE {                continue;            }             *chunk_count += 1;            selected_documents.push(Bm25ContextCandidate {                path: Path::new(&format!(                    "{}/{}",                    worktree_root_name,                    document.relative_path.to_string_lossy()                ))                .into(),                row_range: document.row_range.clone(),                order: next_order + selected_documents.len(),            });             if selected_documents.len() >= BM25_CONTEXT_CHUNK_COUNT {                break;            }        }         selected_documents    }     fn score_document(&self, document: &Document, query_terms: &HashMap<String, f64>) -> f64 {        let document_count = self.documents.len() as f64;        let document_len = document.len as f64;        let mut score = 0.0;         for (term, query_weight) in query_terms {            let Some(term_frequency) = document.term_frequencies.get(term) else {                continue;            };            let document_frequency = self                .document_frequencies                .get(term)                .copied()                .unwrap_or_default() as f64;            if document_frequency == 0.0 {                continue;            }             let inverse_document_frequency =                ((document_count - document_frequency + 0.5) / (document_frequency + 0.5) + 1.0)                    .ln();            let term_frequency = *term_frequency as f64;            let denominator = term_frequency                + BM25_K1                    * (1.0 - BM25_B + BM25_B * document_len / self.average_document_len.max(1.0));            score += query_weight * inverse_document_frequency * term_frequency * (BM25_K1 + 1.0)                / denominator;        }         score    }} async fn git_ls_files(worktree_abs_path: &Path) -> Result<Vec<PathBuf>> {    let output = new_command("git")        .arg("ls-files")        .arg("-z")        .current_dir(worktree_abs_path)        .output()        .await        .with_context(|| {            format!(                "failed to run git ls-files in {}",                worktree_abs_path.display()            )        })?;     if !output.status.success() {        let stderr = String::from_utf8_lossy(&output.stderr);        bail!(            "git ls-files failed in {} with status {}: {}",            worktree_abs_path.display(),            output.status,            stderr.trim()        );    }     let output =        String::from_utf8(output.stdout).context("git ls-files output was not valid UTF-8")?;    Ok(output        .split('\0')        .filter(|path| !path.is_empty())        .map(PathBuf::from)        .collect())} fn documents_for_file(    worktree_abs_path: &Path,    relative_path: PathBuf,) -> Option<DocumentsForFile> {    let absolute_path = worktree_abs_path.join(&relative_path);    let metadata = fs::symlink_metadata(&absolute_path).ok()?;    if metadata.file_type().is_symlink()        || !metadata.is_file()        || metadata.len() > BM25_CONTEXT_MAX_FILE_BYTES    {        return None;    }     let text = fs::read_to_string(&absolute_path).ok()?;    if text.is_empty() {        return None;    }     let byte_len = metadata.len();    let lines = lines(&text);    let path_tokens = tokenize(&relative_path.to_string_lossy());     let documents = chunk_line_ranges(        &lines,        BM25_CONTEXT_CHUNK_LINE_COUNT,        BM25_CONTEXT_CHUNK_OVERLAP_LINE_COUNT,    )    .into_iter()    .filter_map(|row_range| {        let chunk_text = text_for_line_range(&text, row_range.clone());        let mut term_frequencies = HashMap::new();        add_term_frequencies(&mut term_frequencies, tokenize(&chunk_text), 1);        add_term_frequencies(&mut term_frequencies, path_tokens.clone(), 2);        let len = term_frequencies.values().sum();        if len == 0 {            return None;        }         Some(Document {            relative_path: relative_path.clone(),            row_range: row_range.start as u32..row_range.end as u32,            term_frequencies,            len,        })    })    .collect::<Vec<_>>();     (!documents.is_empty()).then_some(DocumentsForFile {        documents,        byte_len,    })} fn add_term_frequencies(    term_frequencies: &mut HashMap<String, usize>,    tokens: Vec<String>,    weight: usize,) {    for token in tokens {        *term_frequencies.entry(token).or_default() += weight;    }} fn chunk_line_ranges(    lines: &[&str],    target_line_count: usize,    overlap_line_count: usize,) -> Vec<Range<usize>> {    if lines.is_empty() || target_line_count == 0 {        return Vec::new();    }     let mut ranges = Vec::new();    let mut start = 0;    while start < lines.len() {        let ideal_end = start.saturating_add(target_line_count).min(lines.len());        let mut end = ideal_end;        if ideal_end < lines.len()            && let Some(boundary) =                empty_line_boundary_near(lines, start, ideal_end, overlap_line_count)        {            end = boundary;        }        if end <= start {            end = ideal_end;        }        if end <= start {            break;        }         ranges.push(start..end);        if end == lines.len() {            break;        }         let next_start = end.saturating_sub(overlap_line_count);        start = if next_start <= start { end } else { next_start };    }     ranges} fn empty_line_boundary_near(    lines: &[&str],    start: usize,

Lines 1–500 of 682.