use std::path::PathBuf; use std::sync::atomic::AtomicBool; use std::sync::Arc; use anyhow::{anyhow, Result}; use async_trait::async_trait; use collections::HashMap; use derive_more::{Deref, Display}; use futures::future::{self, BoxFuture, Shared}; use futures::FutureExt; use fuzzy::StringMatchCandidate; use gpui::{AppContext, BackgroundExecutor, Task}; use heed::types::SerdeBincode; use heed::Database; use parking_lot::RwLock; use serde::{Deserialize, Serialize}; use util::ResultExt; use crate::IndexedDocsRegistry; #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Deref, Display)] pub struct ProviderId(pub Arc); /// The name of a package. #[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Deref, Display)] pub struct PackageName(Arc); impl From<&str> for PackageName { fn from(value: &str) -> Self { Self(value.into()) } } #[async_trait] pub trait IndexedDocsProvider { /// Returns the ID of this provider. fn id(&self) -> ProviderId; /// Returns the path to the database for this provider. fn database_path(&self) -> PathBuf; /// Returns a list of packages as suggestions to be included in the search /// results. /// /// This can be used to provide completions for known packages (e.g., from the /// local project or a registry) before a package has been indexed. async fn suggest_packages(&self) -> Result>; /// Indexes the package with the given name. async fn index(&self, package: PackageName, database: Arc) -> Result<()>; } /// A store for indexed docs. pub struct IndexedDocsStore { executor: BackgroundExecutor, database_future: Shared, Arc>>>, provider: Box, indexing_tasks_by_package: RwLock>>>>>, latest_errors_by_package: RwLock>>, } impl IndexedDocsStore { pub fn try_global(provider: ProviderId, cx: &AppContext) -> Result> { let registry = IndexedDocsRegistry::global(cx); registry .get_provider_store(provider.clone()) .ok_or_else(|| anyhow!("no indexed docs store found for {provider}")) } pub fn new( provider: Box, executor: BackgroundExecutor, ) -> Self { let database_future = executor .spawn({ let executor = executor.clone(); let database_path = provider.database_path(); async move { IndexedDocsDatabase::new(database_path, executor) } }) .then(|result| future::ready(result.map(Arc::new).map_err(Arc::new))) .boxed() .shared(); Self { executor, database_future, provider, indexing_tasks_by_package: RwLock::new(HashMap::default()), latest_errors_by_package: RwLock::new(HashMap::default()), } } pub fn latest_error_for_package(&self, package: &PackageName) -> Option> { self.latest_errors_by_package.read().get(package).cloned() } /// Returns whether the package with the given name is currently being indexed. pub fn is_indexing(&self, package: &PackageName) -> bool { self.indexing_tasks_by_package.read().contains_key(package) } pub async fn load(&self, key: String) -> Result { self.database_future .clone() .await .map_err(|err| anyhow!(err))? .load(key) .await } pub async fn load_many_by_prefix(&self, prefix: String) -> Result> { self.database_future .clone() .await .map_err(|err| anyhow!(err))? .load_many_by_prefix(prefix) .await } /// Returns whether any entries exist with the given prefix. pub async fn any_with_prefix(&self, prefix: String) -> Result { self.database_future .clone() .await .map_err(|err| anyhow!(err))? .any_with_prefix(prefix) .await } pub fn suggest_packages(self: Arc) -> Task>> { let this = self.clone(); self.executor .spawn(async move { this.provider.suggest_packages().await }) } pub fn index( self: Arc, package: PackageName, ) -> Shared>>> { if let Some(existing_task) = self.indexing_tasks_by_package.read().get(&package) { return existing_task.clone(); } let indexing_task = self .executor .spawn({ let this = self.clone(); let package = package.clone(); async move { let _finally = util::defer({ let this = this.clone(); let package = package.clone(); move || { this.indexing_tasks_by_package.write().remove(&package); } }); let index_task = { let package = package.clone(); async { let database = this .database_future .clone() .await .map_err(|err| anyhow!(err))?; this.provider.index(package, database).await } }; let result = index_task.await.map_err(Arc::new); match &result { Ok(_) => { this.latest_errors_by_package.write().remove(&package); } Err(err) => { this.latest_errors_by_package .write() .insert(package, err.to_string().into()); } } result } }) .shared(); self.indexing_tasks_by_package .write() .insert(package, 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 { 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(Debug, PartialEq, Eq, Clone, Display, Serialize, Deserialize)] pub struct MarkdownDocs(pub String); pub struct IndexedDocsDatabase { executor: BackgroundExecutor, env: heed::Env, entries: Database, SerdeBincode>, } impl IndexedDocsDatabase { 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, key: String) -> Task> { let env = self.env.clone(); let entries = self.entries; self.executor.spawn(async move { let txn = env.read_txn()?; entries .get(&txn, &key)? .ok_or_else(|| anyhow!("no docs found for {key}")) }) } pub fn load_many_by_prefix(&self, prefix: String) -> Task>> { let env = self.env.clone(); let entries = self.entries; self.executor.spawn(async move { let txn = env.read_txn()?; let results = entries .iter(&txn)? .filter_map(|entry| { let (key, value) = entry.ok()?; if key.starts_with(&prefix) { Some((key, value)) } else { None } }) .collect::>(); Ok(results) }) } /// Returns whether any entries exist with the given prefix. pub fn any_with_prefix(&self, prefix: String) -> Task> { let env = self.env.clone(); let entries = self.entries; self.executor.spawn(async move { let txn = env.read_txn()?; let any = entries .iter(&txn)? .any(|entry| entry.map_or(false, |(key, _value)| key.starts_with(&prefix))); Ok(any) }) } pub fn insert(&self, key: String, docs: String) -> Task> { let env = self.env.clone(); let entries = self.entries; self.executor.spawn(async move { let mut txn = env.write_txn()?; entries.put(&mut txn, &key, &MarkdownDocs(docs))?; txn.commit()?; Ok(()) }) } } impl extension::KeyValueStoreDelegate for IndexedDocsDatabase { fn insert(&self, key: String, docs: String) -> Task> { IndexedDocsDatabase::insert(&self, key, docs) } }