287 lines · 8.3 KB · Rust
RawDownload
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 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<str>); impl From<&str> for CrateName {    fn from(value: &str) -> Self {        Self(value.into())    }} struct GlobalRustdocStore(Arc<RustdocStore>); impl Global for GlobalRustdocStore {} pub struct RustdocStore {    executor: BackgroundExecutor,    database_future: Shared<BoxFuture<'static, Result<Arc<RustdocDatabase>, Arc<anyhow::Error>>>>,    indexing_tasks_by_crate:        RwLock<HashMap<CrateName, Shared<Task<Result<(), Arc<anyhow::Error>>>>>>,} impl RustdocStore {    pub fn global(cx: &AppContext) -> Arc<Self> {        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(                        paths::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<String>,    ) -> Result<RustdocDatabaseEntry> {        self.database_future            .clone()            .await            .map_err(|err| anyhow!(err))?            .load(crate_name, item_path)            .await    }     pub fn index(        self: Arc<Self>,        crate_name: CrateName,        provider: Box<dyn RustdocProvider + Send + Sync + 'static>,    ) -> Shared<Task<Result<(), Arc<anyhow::Error>>>> {        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<Vec<String>> {        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::<Vec<_>>();             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<String>, SerdeBincode<RustdocDatabaseEntry>>,} impl RustdocDatabase {    pub fn new(path: PathBuf, executor: BackgroundExecutor) -> Result<Self> {        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<Result<Vec<String>>> {        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<String>,    ) -> Task<Result<RustdocDatabaseEntry>> {        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<Result<()>> {        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(())        })    }}