use std::path::PathBuf; use std::sync::atomic::AtomicBool; use std::sync::Arc; use anyhow::{anyhow, Result}; use collections::HashMap; use derive_more::{Deref, Display}; use futures::future::{self, BoxFuture, Shared}; use futures::FutureExt; use fuzzy::StringMatchCandidate; use gpui::{AppContext, BackgroundExecutor, Global, ReadGlobal, Task, UpdateGlobal}; use heed::types::SerdeBincode; use heed::Database; use parking_lot::RwLock; use paths::SUPPORT_DIR; use serde::{Deserialize, Serialize}; use util::ResultExt; use crate::indexer::{RustdocIndexer, RustdocProvider}; use crate::{RustdocItem, RustdocItemKind}; /// The name of a crate. #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Deref, Display)] pub struct CrateName(Arc); impl From<&str> for CrateName { fn from(value: &str) -> Self { Self(value.into()) } } struct GlobalRustdocStore(Arc); impl Global for GlobalRustdocStore {} pub struct RustdocStore { executor: BackgroundExecutor, database_future: Shared, Arc>>>, indexing_tasks_by_crate: RwLock>>>>>, } impl RustdocStore { pub fn global(cx: &AppContext) -> Arc { GlobalRustdocStore::global(cx).0.clone() } pub fn init_global(cx: &mut AppContext) { GlobalRustdocStore::set_global( cx, GlobalRustdocStore(Arc::new(Self::new(cx.background_executor().clone()))), ); } pub fn new(executor: BackgroundExecutor) -> Self { let database_future = executor .spawn({ let executor = executor.clone(); async move { RustdocDatabase::new(SUPPORT_DIR.join("docs/rust/rustdoc-db.0.mdb"), executor) } }) .then(|result| future::ready(result.map(Arc::new).map_err(Arc::new))) .boxed() .shared(); Self { executor, database_future, indexing_tasks_by_crate: RwLock::new(HashMap::default()), } } /// Returns whether the crate with the given name is currently being indexed. pub fn is_indexing(&self, crate_name: &CrateName) -> bool { self.indexing_tasks_by_crate.read().contains_key(crate_name) } pub async fn load( &self, crate_name: CrateName, item_path: Option, ) -> Result { self.database_future .clone() .await .map_err(|err| anyhow!(err))? .load(crate_name, item_path) .await } pub fn index( self: Arc, crate_name: CrateName, provider: Box, ) -> Shared>>> { if let Some(existing_task) = self.indexing_tasks_by_crate.read().get(&crate_name) { return existing_task.clone(); } let indexing_task = self .executor .spawn({ let this = self.clone(); let crate_name = crate_name.clone(); async move { let _finally = util::defer({ let this = this.clone(); let crate_name = crate_name.clone(); move || { this.indexing_tasks_by_crate.write().remove(&crate_name); } }); let index_task = async { let database = this .database_future .clone() .await .map_err(|err| anyhow!(err))?; let indexer = RustdocIndexer::new(database, provider); indexer.index(crate_name.clone()).await }; index_task.await.map_err(Arc::new) } }) .shared(); self.indexing_tasks_by_crate .write() .insert(crate_name, indexing_task.clone()); indexing_task } pub fn search(&self, query: String) -> Task> { let executor = self.executor.clone(); let database_future = self.database_future.clone(); self.executor.spawn(async move { if query.is_empty() { return Vec::new(); } let Some(database) = database_future.await.map_err(|err| anyhow!(err)).log_err() else { return Vec::new(); }; let Some(items) = database.keys().await.log_err() else { return Vec::new(); }; let candidates = items .iter() .enumerate() .map(|(ix, item_path)| StringMatchCandidate::new(ix, item_path.clone())) .collect::>(); let matches = fuzzy::match_strings( &candidates, &query, false, 100, &AtomicBool::default(), executor, ) .await; matches .into_iter() .map(|mat| items[mat.candidate_id].clone()) .collect() }) } } #[derive(Serialize, Deserialize)] pub enum RustdocDatabaseEntry { Crate { docs: String }, Item { kind: RustdocItemKind, docs: String }, } impl RustdocDatabaseEntry { pub fn docs(&self) -> &str { match self { Self::Crate { docs } | Self::Item { docs, .. } => &docs, } } } pub(crate) struct RustdocDatabase { executor: BackgroundExecutor, env: heed::Env, entries: Database, SerdeBincode>, } impl RustdocDatabase { pub fn new(path: PathBuf, executor: BackgroundExecutor) -> Result { std::fs::create_dir_all(&path)?; const ONE_GB_IN_BYTES: usize = 1024 * 1024 * 1024; let env = unsafe { heed::EnvOpenOptions::new() .map_size(ONE_GB_IN_BYTES) .max_dbs(1) .open(path)? }; let mut txn = env.write_txn()?; let entries = env.create_database(&mut txn, Some("rustdoc_entries"))?; txn.commit()?; Ok(Self { executor, env, entries, }) } pub fn keys(&self) -> Task>> { let env = self.env.clone(); let entries = self.entries; self.executor.spawn(async move { let txn = env.read_txn()?; let mut iter = entries.iter(&txn)?; let mut keys = Vec::new(); while let Some((key, _value)) = iter.next().transpose()? { keys.push(key); } Ok(keys) }) } pub fn load( &self, crate_name: CrateName, item_path: Option, ) -> Task> { let env = self.env.clone(); let entries = self.entries; let item_path = if let Some(item_path) = item_path { format!("{crate_name}::{item_path}") } else { crate_name.to_string() }; self.executor.spawn(async move { let txn = env.read_txn()?; entries .get(&txn, &item_path)? .ok_or_else(|| anyhow!("no docs found for {item_path}")) }) } pub fn insert( &self, crate_name: CrateName, item: Option<&RustdocItem>, docs: String, ) -> Task> { let env = self.env.clone(); let entries = self.entries; let (item_path, entry) = if let Some(item) = item { ( format!("{crate_name}::{}", item.display()), RustdocDatabaseEntry::Item { kind: item.kind, docs, }, ) } else { (crate_name.to_string(), RustdocDatabaseEntry::Crate { docs }) }; self.executor.spawn(async move { let mut txn = env.write_txn()?; entries.put(&mut txn, &item_path, &entry)?; txn.commit()?; Ok(()) }) } }