use collections::{HashMap, HashSet};use gpui::{App, AppContext as _, AsyncApp, Entity, EntityId};use language::{Buffer, BufferSnapshot, Point, ToPoint as _};use project::{Project, ProjectPath};use std::{    ops::Range,    path::{Path, PathBuf},    sync::Arc,};use text::Anchor;use util::{paths::PathStyle, rel_path::RelPath};use zeta_prompt::{ContextSource, RelatedExcerpt, RelatedFile, multi_region::is_good_block_start}; use crate::{    bm25_context::{Bm25ContextCandidate, collect_bm25_context},    git_log_context::build_git_log_index,}; /// This module contains collectors for editable context:/// excerpts or full files that are likely to be edited.const CURSOR_CONTEXT_LINE_COUNT: u32 = 20;const EDIT_HISTORY_CONTEXT_LINE_COUNT: u32 = 20;const GIT_LOG_CONTEXT_LINE_COUNT: u32 = 10000;const GIT_LOG_CONTEXT_FILE_COUNT: usize = 10;const ORACLE_SNIPPET_MIN_CONTEXT_LINE_COUNT: u32 = 10;const ORACLE_SNIPPET_MAX_CONTEXT_LINE_COUNT: u32 = 40;/// How far excerpt boundaries may be nudged to land on a natural block/// boundary, mirroring `zeta_prompt::multi_region`'s marker placement.const BOUNDARY_SNAP_LINE_COUNT: u32 = 5;/// Maximum number of rows between two excerpts of the same buffer that get/// bridged into one contiguous excerpt instead of rendering an elision/// marker between them.const BRIDGED_GAP_LINE_COUNT: u32 = 3; type RangesByBuffer = HashMap<EntityId, (Entity<Buffer>, Vec<EditableContextRange>)>; #[derive(Clone)]pub struct EditHistoryContextEntry {    pub buffer: Entity<Buffer>,    pub edited_range: Range<Anchor>,} /// A file known (from expected patches) to be edited next, used by the/// oracle context sources when generating training data.#[derive(Clone, Debug)]pub struct OracleTarget {    pub path: Arc<Path>,    /// 0-based, end-exclusive row ranges of the expected edit hunks.    /// Only used by `ContextSource::OracleSnippet`.    pub row_ranges: Vec<Range<u32>>,} struct EditableContextRange {    range: Range<Anchor>,    order: usize,    context_source: ContextSource,} struct ResolvedEditableContextRange {    range: Range<Point>,    order: usize,    context_source: ContextSource,} pub async fn collect_editable_context(    project: Entity<Project>,    active_buffer: Entity<Buffer>,    cursor_position: Anchor,    edit_history: Vec<EditHistoryContextEntry>,    oracle_targets: Vec<OracleTarget>,    context_sources: Vec<ContextSource>,    cx: &mut AsyncApp,) -> anyhow::Result<Vec<RelatedFile>> {    let mut ranges_by_buffer = RangesByBuffer::default();     if context_sources.contains(&ContextSource::CursorExcerpt) {        collect_cursor_excerpt_context(            &mut ranges_by_buffer,            active_buffer.clone(),            cursor_position,            cx,        );    }    if context_sources.contains(&ContextSource::CurrentFile) {        collect_current_file_context(&mut ranges_by_buffer, active_buffer.clone(), cx);    }    if context_sources.contains(&ContextSource::EditHistory) {        collect_edit_history_context(&mut ranges_by_buffer, &edit_history, cx);    }    if context_sources.contains(&ContextSource::EditHistoryFile) {        collect_edit_history_file_context(&mut ranges_by_buffer, &edit_history, cx);    }    if context_sources.contains(&ContextSource::GitLog) {        collect_git_log_context(            &mut ranges_by_buffer,            project.clone(),            active_buffer.clone(),            cx,        )        .await;    }     // Collected before bm25 so that, under a related-files byte budget, the    // small snippets containing the expected edits are never trimmed away in    // favor of bm25 excerpts.    if context_sources.contains(&ContextSource::OracleSnippet) {        collect_oracle_snippet_context(&mut ranges_by_buffer, project.clone(), &oracle_targets, cx)            .await;    }     if context_sources.contains(&ContextSource::Bm25) {        collect_bm25_context_ranges(            &mut ranges_by_buffer,            project.clone(),            active_buffer,            cursor_position,            &edit_history,            cx,        )        .await;    }     if context_sources.contains(&ContextSource::OracleFile) {        collect_oracle_file_context(&mut ranges_by_buffer, project.clone(), &oracle_targets, cx)            .await;    }     Ok(cx.update(|cx| {        let project = project.read(cx);        let mut related_files = ranges_by_buffer            .into_values()            .filter_map(|(buffer, ranges)| related_file_for_ranges(&project, &buffer, ranges, cx))            .collect::<Vec<_>>();        related_files.sort_by_key(|file| {            file.excerpts                .iter()                .map(|excerpt| excerpt.order)                .min()                .unwrap_or(usize::MAX)        });        related_files    }))} pub fn limit_retrieved_context_to_bytes(    related_files: &[RelatedFile],    max_bytes: usize,) -> Vec<RelatedFile> {    struct ExcerptCandidate {        file_index: usize,        excerpt_index: usize,        order: usize,    }     let mut candidates = related_files        .iter()        .enumerate()        .flat_map(|(file_index, file)| {            file.excerpts                .iter()                .enumerate()                .map(move |(excerpt_index, excerpt)| ExcerptCandidate {                    file_index,                    excerpt_index,                    order: excerpt.order,                })        })        .collect::<Vec<_>>();    candidates.sort_by_key(|candidate| {        (            candidate.order,            candidate.file_index,            candidate.excerpt_index,        )    });     let mut selected_excerpts = related_files        .iter()        .map(|file| vec![false; file.excerpts.len()])        .collect::<Vec<_>>();    let mut covered_ranges_by_file = vec![Vec::<Range<u32>>::new(); related_files.len()];    let mut selected_bytes: usize = 0;     for candidate in candidates {        let file = &related_files[candidate.file_index];        let excerpt = &file.excerpts[candidate.excerpt_index];        let added_bytes =            uncovered_excerpt_bytes(excerpt, &covered_ranges_by_file[candidate.file_index]);        if added_bytes == 0 || selected_bytes.saturating_add(added_bytes) > max_bytes {            continue;        }         selected_bytes += added_bytes;        selected_excerpts[candidate.file_index][candidate.excerpt_index] = true;        push_covered_range(            &mut covered_ranges_by_file[candidate.file_index],            excerpt.row_range.clone(),        );    }     related_files        .iter()        .enumerate()        .filter_map(|(file_index, file)| {            let excerpts = file                .excerpts                .iter()                .enumerate()                .filter_map(|(excerpt_index, excerpt)| {                    selected_excerpts[file_index][excerpt_index].then(|| excerpt.clone())                })                .collect::<Vec<_>>();            if excerpts.is_empty() {                return None;            }             Some(RelatedFile {                path: file.path.clone(),                max_row: file.max_row,                excerpts,                in_open_source_repo: file.in_open_source_repo,            })        })        .collect()} fn uncovered_excerpt_bytes(excerpt: &RelatedExcerpt, covered_ranges: &[Range<u32>]) -> usize {    let mut bytes = 0;     for (row, line) in (excerpt.row_range.start..).zip(excerpt.text.split_inclusive('\n')) {        if row >= excerpt.row_range.end {            break;        }        if !covered_ranges            .iter()            .any(|covered_range| covered_range.contains(&row))        {            bytes += line.len();        }    }     bytes} fn push_covered_range(covered_ranges: &mut Vec<Range<u32>>, range: Range<u32>) {    covered_ranges.push(range);    covered_ranges.sort_by_key(|range| (range.start, range.end));     let mut merged_ranges: Vec<Range<u32>> = Vec::new();    for range in covered_ranges.drain(..) {        if let Some(last_range) = merged_ranges.last_mut()            && range.start <= last_range.end        {            last_range.end = last_range.end.max(range.end);            continue;        }         merged_ranges.push(range);    }     *covered_ranges = merged_ranges;} fn collect_cursor_excerpt_context(    ranges_by_buffer: &mut RangesByBuffer,    active_buffer: Entity<Buffer>,    cursor_position: Anchor,    cx: &mut AsyncApp,) {    let cursor_range = active_buffer.read_with(cx, |buffer, _cx| {        let snapshot = buffer.snapshot();        expanded_anchor_range(            &snapshot,            cursor_position..cursor_position,            CURSOR_CONTEXT_LINE_COUNT,        )    });     push_context_range(        ranges_by_buffer,        active_buffer,        cursor_range,        0,        ContextSource::CursorExcerpt,    );} fn collect_current_file_context(    ranges_by_buffer: &mut RangesByBuffer,    active_buffer: Entity<Buffer>,    cx: &mut AsyncApp,) {    collect_full_buffer_context(        ranges_by_buffer,        active_buffer,        0,        ContextSource::CurrentFile,        cx,    );} fn collect_edit_history_context(    ranges_by_buffer: &mut RangesByBuffer,    edit_history: &[EditHistoryContextEntry],    cx: &mut AsyncApp,) {    for (index, entry) in edit_history.iter().enumerate() {        let edit_history_range = entry.buffer.read_with(cx, |buffer, _cx| {            expanded_anchor_range(                &buffer.snapshot(),                entry.edited_range.clone(),                EDIT_HISTORY_CONTEXT_LINE_COUNT,            )        });         push_context_range(            ranges_by_buffer,            entry.buffer.clone(),            edit_history_range,            index + 1,            ContextSource::EditHistory,        );    }} fn collect_edit_history_file_context(    ranges_by_buffer: &mut RangesByBuffer,    edit_history: &[EditHistoryContextEntry],    cx: &mut AsyncApp,) {    let next_order = next_context_order(ranges_by_buffer);    let mut seen_buffers = HashSet::default();    let mut index = 0;     for entry in edit_history {        if !seen_buffers.insert(entry.buffer.entity_id()) {            continue;        }         collect_full_buffer_context(            ranges_by_buffer,            entry.buffer.clone(),            next_order + index,            ContextSource::EditHistoryFile,            cx,        );        index += 1;    }} async fn collect_bm25_context_ranges(    ranges_by_buffer: &mut RangesByBuffer,    project: Entity<Project>,    active_buffer: Entity<Buffer>,    cursor_position: Anchor,    edit_history: &[EditHistoryContextEntry],    cx: &mut AsyncApp,) {    let next_order = next_context_order(ranges_by_buffer);    let candidates = collect_bm25_context(        project.clone(),        active_buffer,        cursor_position,        edit_history,        next_order,        cx,    )    .await;     for candidate in candidates {        collect_bm25_candidate_context(ranges_by_buffer, &project, candidate, cx).await;    }} async fn collect_bm25_candidate_context(    ranges_by_buffer: &mut RangesByBuffer,    project: &Entity<Project>,    candidate: Bm25ContextCandidate,    cx: &mut AsyncApp,) {    let buffer = match open_buffer_for_path(project, &candidate.path, cx).await {        Ok(Some(buffer)) => buffer,        Ok(None) => {            log::debug!(                "failed to find BM25 context path: {}",                candidate.path.display()            );            return;        }        Err(error) => {            log::debug!(                "failed to open BM25 context path {}: {error:#}",                candidate.path.display()            );            return;        }    };     let Some(range) = buffer.read_with(cx, |buffer, _cx| {        anchor_range_for_row_range(&buffer.snapshot(), candidate.row_range.clone())    }) else {        return;    };     push_context_range(        ranges_by_buffer,        buffer,        range,        candidate.order,        ContextSource::Bm25,    );} fn anchor_range_for_row_range(    snapshot: &BufferSnapshot,    row_range: Range<u32>,) -> Option<Range<Anchor>> {    if row_range.start >= row_range.end || row_range.start > snapshot.max_point().row {        return None;    }     let max_point = snapshot.max_point();    let start = snapshot.anchor_before(Point::new(row_range.start, 0));    let end_point = if row_range.end > max_point.row {        max_point    } else {        Point::new(row_range.end, 0)    };    let end = snapshot.anchor_after(end_point);    Some(start..end)} async fn collect_oracle_file_context(    ranges_by_buffer: &mut RangesByBuffer,    project: Entity<Project>,    oracle_targets: &[OracleTarget],    cx: &mut AsyncApp,) {    let next_order = next_context_order(ranges_by_buffer);    let mut seen_buffers = HashSet::default();    let mut index = 0;     for target in oracle_targets {        let buffer = match open_buffer_for_path(&project, &target.path, cx).await {            Ok(Some(buffer)) => buffer,            Ok(None) => {                log::debug!("failed to find oracle file path: {}", target.path.display());                continue;            }            Err(error) => {                log::debug!(                    "failed to open oracle file path {}: {error:#}",                    target.path.display()                );                continue;            }        };         if !seen_buffers.insert(buffer.entity_id()) {            continue;        }         collect_full_buffer_context(            ranges_by_buffer,            buffer,            next_order + index,            ContextSource::OracleFile,            cx,        );        index += 1;    }} async fn collect_oracle_snippet_context(    ranges_by_buffer: &mut RangesByBuffer,    project: Entity<Project>,    oracle_targets: &[OracleTarget],    cx: &mut AsyncApp,) {    let next_order = next_context_order(ranges_by_buffer);    let mut index = 0;     for target in oracle_targets {        if target.row_ranges.is_empty() {            continue;        }         let buffer = match open_buffer_for_path(&project, &target.path, cx).await {            Ok(Some(buffer)) => buffer,            Ok(None) => {                log::debug!(                    "failed to find oracle snippet path: {}",                    target.path.display()                );                continue;            }            Err(error) => {                log::debug!(                    "failed to open oracle snippet path {}: {error:#}",                    target.path.display()

Lines 1–500 of 1238. Load more