ZIm/crates/rustdoc/src/store.rs

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<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(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(())
        })
    }
}