use crate::{    embedding::{EmbeddingProvider, TextToEmbed},    summary_index::FileSummary,    worktree_index::{WorktreeIndex, WorktreeIndexHandle},};use anyhow::{anyhow, Context as _, Result};use collections::HashMap;use fs::Fs;use futures::FutureExt;use gpui::{    App, AppContext as _, Context, Entity, EntityId, EventEmitter, Subscription, Task, WeakEntity,};use language::LanguageRegistry;use log;use project::{Project, Worktree, WorktreeId};use serde::{Deserialize, Serialize};use smol::channel;use std::{    cmp::Ordering,    future::Future,    num::NonZeroUsize,    ops::{Range, RangeInclusive},    path::{Path, PathBuf},    sync::Arc,};use util::ResultExt; #[derive(Debug)]pub struct SearchResult {    pub worktree: Entity<Worktree>,    pub path: Arc<Path>,    pub range: Range<usize>,    pub score: f32,    pub query_index: usize,} #[derive(Debug, PartialEq, Eq)]pub struct LoadedSearchResult {    pub path: Arc<Path>,    pub full_path: PathBuf,    pub excerpt_content: String,    pub row_range: RangeInclusive<u32>,    pub query_index: usize,} pub struct WorktreeSearchResult {    pub worktree_id: WorktreeId,    pub path: Arc<Path>,    pub range: Range<usize>,    pub query_index: usize,    pub score: f32,} #[derive(Copy, Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]pub enum Status {    Idle,    Loading,    Scanning { remaining_count: NonZeroUsize },} pub struct ProjectIndex {    db_connection: heed::Env,    project: WeakEntity<Project>,    worktree_indices: HashMap<EntityId, WorktreeIndexHandle>,    language_registry: Arc<LanguageRegistry>,    fs: Arc<dyn Fs>,    last_status: Status,    status_tx: channel::Sender<()>,    embedding_provider: Arc<dyn EmbeddingProvider>,    _maintain_status: Task<()>,    _subscription: Subscription,} impl ProjectIndex {    pub fn new(        project: Entity<Project>,        db_connection: heed::Env,        embedding_provider: Arc<dyn EmbeddingProvider>,        cx: &mut Context<Self>,    ) -> Self {        let language_registry = project.read(cx).languages().clone();        let fs = project.read(cx).fs().clone();        let (status_tx, status_rx) = channel::unbounded();        let mut this = ProjectIndex {            db_connection,            project: project.downgrade(),            worktree_indices: HashMap::default(),            language_registry,            fs,            status_tx,            last_status: Status::Idle,            embedding_provider,            _subscription: cx.subscribe(&project, Self::handle_project_event),            _maintain_status: cx.spawn(|this, mut cx| async move {                while status_rx.recv().await.is_ok() {                    if this                        .update(&mut cx, |this, cx| this.update_status(cx))                        .is_err()                    {                        break;                    }                }            }),        };        this.update_worktree_indices(cx);        this    }     pub fn status(&self) -> Status {        self.last_status    }     pub fn project(&self) -> WeakEntity<Project> {        self.project.clone()    }     pub fn fs(&self) -> Arc<dyn Fs> {        self.fs.clone()    }     fn handle_project_event(        &mut self,        _: Entity<Project>,        event: &project::Event,        cx: &mut Context<Self>,    ) {        match event {            project::Event::WorktreeAdded(_) | project::Event::WorktreeRemoved(_) => {                self.update_worktree_indices(cx);            }            _ => {}        }    }     fn update_worktree_indices(&mut self, cx: &mut Context<Self>) {        let Some(project) = self.project.upgrade() else {            return;        };         let worktrees = project            .read(cx)            .visible_worktrees(cx)            .filter_map(|worktree| {                if worktree.read(cx).is_local() {                    Some((worktree.entity_id(), worktree))                } else {                    None                }            })            .collect::<HashMap<_, _>>();         self.worktree_indices            .retain(|worktree_id, _| worktrees.contains_key(worktree_id));        for (worktree_id, worktree) in worktrees {            self.worktree_indices.entry(worktree_id).or_insert_with(|| {                let worktree_index = WorktreeIndex::load(                    worktree.clone(),                    self.db_connection.clone(),                    self.language_registry.clone(),                    self.fs.clone(),                    self.status_tx.clone(),                    self.embedding_provider.clone(),                    cx,                );                 let load_worktree = cx.spawn(|this, mut cx| async move {                    let result = match worktree_index.await {                        Ok(worktree_index) => {                            this.update(&mut cx, |this, _| {                                this.worktree_indices.insert(                                    worktree_id,                                    WorktreeIndexHandle::Loaded {                                        index: worktree_index.clone(),                                    },                                );                            })?;                            Ok(worktree_index)                        }                        Err(error) => {                            this.update(&mut cx, |this, _cx| {                                this.worktree_indices.remove(&worktree_id)                            })?;                            Err(Arc::new(error))                        }                    };                     this.update(&mut cx, |this, cx| this.update_status(cx))?;                     result                });                 WorktreeIndexHandle::Loading {                    index: load_worktree.shared(),                }            });        }         self.update_status(cx);    }     fn update_status(&mut self, cx: &mut Context<Self>) {        let mut indexing_count = 0;        let mut any_loading = false;         for index in self.worktree_indices.values_mut() {            match index {                WorktreeIndexHandle::Loading { .. } => {                    any_loading = true;                    break;                }                WorktreeIndexHandle::Loaded { index, .. } => {                    indexing_count += index.read(cx).entry_ids_being_indexed().len();                }            }        }         let status = if any_loading {            Status::Loading        } else if let Some(remaining_count) = NonZeroUsize::new(indexing_count) {            Status::Scanning { remaining_count }        } else {            Status::Idle        };         if status != self.last_status {            self.last_status = status;            cx.emit(status);        }    }     pub fn search(        &self,        queries: Vec<String>,        limit: usize,        cx: &App,    ) -> Task<Result<Vec<SearchResult>>> {        let (chunks_tx, chunks_rx) = channel::bounded(1024);        let mut worktree_scan_tasks = Vec::new();        for worktree_index in self.worktree_indices.values() {            let worktree_index = worktree_index.clone();            let chunks_tx = chunks_tx.clone();            worktree_scan_tasks.push(cx.spawn(|cx| async move {                let index = match worktree_index {                    WorktreeIndexHandle::Loading { index } => {                        index.clone().await.map_err(|error| anyhow!(error))?                    }                    WorktreeIndexHandle::Loaded { index } => index.clone(),                };                 index                    .read_with(&cx, |index, cx| {                        let worktree_id = index.worktree().read(cx).id();                        let db_connection = index.db_connection().clone();                        let db = *index.embedding_index().db();                        cx.background_spawn(async move {                            let txn = db_connection                                .read_txn()                                .context("failed to create read transaction")?;                            let db_entries = db.iter(&txn).context("failed to iterate database")?;                            for db_entry in db_entries {                                let (_key, db_embedded_file) = db_entry?;                                for chunk in db_embedded_file.chunks {                                    chunks_tx                                        .send((worktree_id, db_embedded_file.path.clone(), chunk))                                        .await?;                                }                            }                            anyhow::Ok(())                        })                    })?                    .await            }));        }        drop(chunks_tx);         let project = self.project.clone();        let embedding_provider = self.embedding_provider.clone();        cx.spawn(|cx| async move {            #[cfg(debug_assertions)]            let embedding_query_start = std::time::Instant::now();            log::info!("Searching for {queries:?}");            let queries: Vec<TextToEmbed> = queries                .iter()                .map(|s| TextToEmbed::new(s.as_str()))                .collect();             let query_embeddings = embedding_provider.embed(&queries[..]).await?;            if query_embeddings.len() != queries.len() {                return Err(anyhow!(                    "The number of query embeddings does not match the number of queries"                ));            }             let mut results_by_worker = Vec::new();            for _ in 0..cx.background_executor().num_cpus() {                results_by_worker.push(Vec::<WorktreeSearchResult>::new());            }             #[cfg(debug_assertions)]            let search_start = std::time::Instant::now();            cx.background_executor()                .scoped(|cx| {                    for results in results_by_worker.iter_mut() {                        cx.spawn(async {                            while let Ok((worktree_id, path, chunk)) = chunks_rx.recv().await {                                let (score, query_index) =                                    chunk.embedding.similarity(&query_embeddings);                                 let ix = match results.binary_search_by(|probe| {                                    score.partial_cmp(&probe.score).unwrap_or(Ordering::Equal)                                }) {                                    Ok(ix) | Err(ix) => ix,                                };                                if ix < limit {                                    results.insert(                                        ix,                                        WorktreeSearchResult {                                            worktree_id,                                            path: path.clone(),                                            range: chunk.chunk.range.clone(),                                            query_index,                                            score,                                        },                                    );                                    if results.len() > limit {                                        results.pop();                                    }                                }                            }                        });                    }                })                .await;             for scan_task in futures::future::join_all(worktree_scan_tasks).await {                scan_task.log_err();            }             project.read_with(&cx, |project, cx| {                let mut search_results = Vec::with_capacity(results_by_worker.len() * limit);                for worker_results in results_by_worker {                    search_results.extend(worker_results.into_iter().filter_map(|result| {                        Some(SearchResult {                            worktree: project.worktree_for_id(result.worktree_id, cx)?,                            path: result.path,                            range: result.range,                            score: result.score,                            query_index: result.query_index,                        })                    }));                }                search_results.sort_unstable_by(|a, b| {                    b.score.partial_cmp(&a.score).unwrap_or(Ordering::Equal)                });                search_results.truncate(limit);                 #[cfg(debug_assertions)]                {                    let search_elapsed = search_start.elapsed();                    log::debug!(                        "searched {} entries in {:?}",                        search_results.len(),                        search_elapsed                    );                    let embedding_query_elapsed = embedding_query_start.elapsed();                    log::debug!("embedding query took {:?}", embedding_query_elapsed);                }                 search_results            })        })    }     #[cfg(test)]    pub fn path_count(&self, cx: &App) -> Result<u64> {        let mut result = 0;        for worktree_index in self.worktree_indices.values() {            if let WorktreeIndexHandle::Loaded { index, .. } = worktree_index {                result += index.read(cx).path_count()?;            }        }        Ok(result)    }     pub(crate) fn worktree_index(        &self,        worktree_id: WorktreeId,        cx: &App,    ) -> Option<Entity<WorktreeIndex>> {        for index in self.worktree_indices.values() {            if let WorktreeIndexHandle::Loaded { index, .. } = index {                if index.read(cx).worktree().read(cx).id() == worktree_id {                    return Some(index.clone());                }            }        }        None    }     pub(crate) fn worktree_indices(&self, cx: &App) -> Vec<Entity<WorktreeIndex>> {        let mut result = self            .worktree_indices            .values()            .filter_map(|index| {                if let WorktreeIndexHandle::Loaded { index, .. } = index {                    Some(index.clone())                } else {                    None                }            })            .collect::<Vec<_>>();        result.sort_by_key(|index| index.read(cx).worktree().read(cx).id());        result    }     pub fn all_summaries(&self, cx: &App) -> Task<Result<Vec<FileSummary>>> {        let (summaries_tx, summaries_rx) = channel::bounded(1024);        let mut worktree_scan_tasks = Vec::new();        for worktree_index in self.worktree_indices.values() {            let worktree_index = worktree_index.clone();            let summaries_tx: channel::Sender<(String, String)> = summaries_tx.clone();            worktree_scan_tasks.push(cx.spawn(|cx| async move {                let index = match worktree_index {                    WorktreeIndexHandle::Loading { index } => {                        index.clone().await.map_err(|error| anyhow!(error))?                    }                    WorktreeIndexHandle::Loaded { index } => index.clone(),                };                 index                    .read_with(&cx, |index, cx| {                        let db_connection = index.db_connection().clone();                        let summary_index = index.summary_index();                        let file_digest_db = summary_index.file_digest_db();                        let summary_db = summary_index.summary_db();                         cx.background_spawn(async move {                            let txn = db_connection                                .read_txn()                                .context("failed to create db read transaction")?;                            let db_entries = file_digest_db                                .iter(&txn)                                .context("failed to iterate database")?;                            for db_entry in db_entries {                                let (file_path, db_file) = db_entry?;                                 match summary_db.get(&txn, &db_file.digest) {                                    Ok(opt_summary) => {                                        // Currently, we only use summaries we already have. If the file hasn't been                                        // summarized yet, then we skip it and don't include it in the inferred context.                                        // If we want to do just-in-time summarization, this would be the place to do it!                                        if let Some(summary) = opt_summary {                                            summaries_tx                                                .send((file_path.to_string(), summary.to_string()))                                                .await?;                                        } else {                                            log::warn!("No summary found for {:?}", &db_file);                                        }                                    }                                    Err(err) => {                                        log::error!(                                            "Error reading from summary database: {:?}",                                            err                                        );                                    }                                }                            }                            anyhow::Ok(())                        })                    })?                    .await            }));        }        drop(summaries_tx);         let project = self.project.clone();        cx.spawn(|cx| async move {            let mut results_by_worker = Vec::new();            for _ in 0..cx.background_executor().num_cpus() {                results_by_worker.push(Vec::<FileSummary>::new());            }             cx.background_executor()                .scoped(|cx| {                    for results in results_by_worker.iter_mut() {                        cx.spawn(async {                            while let Ok((filename, summary)) = summaries_rx.recv().await {                                results.push(FileSummary { filename, summary });                            }                        });                    }                })                .await;             for scan_task in futures::future::join_all(worktree_scan_tasks).await {                scan_task.log_err();            }             project.read_with(&cx, |_project, _cx| {                results_by_worker.into_iter().flatten().collect()

Lines 1–500 of 552. Load more