ZIm/crates/indexed_docs/src/store.rs

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<str>);

impl ProviderId {
    pub fn rustdoc() -> Self {
        Self("rustdoc".into())
    }
}

/// The name of a package.
#[derive(Debug, PartialEq, Eq, PartialOrd, Ord, Hash, Clone, Deref, Display)]
pub struct PackageName(Arc<str>);

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;

    /// Indexes the package with the given name.
    async fn index(&self, package: PackageName, database: Arc<IndexedDocsDatabase>) -> Result<()>;
}

/// A store for indexed docs.
pub struct IndexedDocsStore {
    executor: BackgroundExecutor,
    database_future:
        Shared<BoxFuture<'static, Result<Arc<IndexedDocsDatabase>, Arc<anyhow::Error>>>>,
    provider: Box<dyn IndexedDocsProvider + Send + Sync + 'static>,
    indexing_tasks_by_package:
        RwLock<HashMap<PackageName, Shared<Task<Result<(), Arc<anyhow::Error>>>>>>,
}

impl IndexedDocsStore {
    pub fn try_global(provider: ProviderId, cx: &AppContext) -> Result<Arc<Self>> {
        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<dyn IndexedDocsProvider + Send + Sync + 'static>,
        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()),
        }
    }

    /// 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<MarkdownDocs> {
        self.database_future
            .clone()
            .await
            .map_err(|err| anyhow!(err))?
            .load(key)
            .await
    }

    pub fn index(
        self: Arc<Self>,
        package: PackageName,
    ) -> Shared<Task<Result<(), Arc<anyhow::Error>>>> {
        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 = async {
                        let database = this
                            .database_future
                            .clone()
                            .await
                            .map_err(|err| anyhow!(err))?;
                        this.provider.index(package, database).await
                    };

                    index_task.await.map_err(Arc::new)
                }
            })
            .shared();

        self.indexing_tasks_by_package
            .write()
            .insert(package, 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 {
            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(Debug, PartialEq, Eq, Clone, Display, Serialize, Deserialize)]
pub struct MarkdownDocs(pub String);

pub struct IndexedDocsDatabase {
    executor: BackgroundExecutor,
    env: heed::Env,
    entries: Database<SerdeBincode<String>, SerdeBincode<MarkdownDocs>>,
}

impl IndexedDocsDatabase {
    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, key: String) -> Task<Result<MarkdownDocs>> {
        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 insert(&self, key: String, docs: String) -> Task<Result<()>> {
        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(())
        })
    }
}