store.rs

  1use std::path::PathBuf;
  2use std::sync::atomic::AtomicBool;
  3use std::sync::Arc;
  4
  5use anyhow::{anyhow, Result};
  6use futures::future::{self, BoxFuture, Shared};
  7use futures::FutureExt;
  8use fuzzy::StringMatchCandidate;
  9use gpui::{AppContext, BackgroundExecutor, Global, ReadGlobal, Task, UpdateGlobal};
 10use heed::types::SerdeBincode;
 11use heed::Database;
 12use serde::{Deserialize, Serialize};
 13use util::paths::SUPPORT_DIR;
 14use util::ResultExt;
 15
 16use crate::indexer::{RustdocIndexer, RustdocProvider};
 17use crate::{RustdocItem, RustdocItemKind};
 18
 19struct GlobalRustdocStore(Arc<RustdocStore>);
 20
 21impl Global for GlobalRustdocStore {}
 22
 23pub struct RustdocStore {
 24    executor: BackgroundExecutor,
 25    database_future: Shared<BoxFuture<'static, Result<Arc<RustdocDatabase>, Arc<anyhow::Error>>>>,
 26}
 27
 28impl RustdocStore {
 29    pub fn global(cx: &AppContext) -> Arc<Self> {
 30        GlobalRustdocStore::global(cx).0.clone()
 31    }
 32
 33    pub fn init_global(cx: &mut AppContext) {
 34        GlobalRustdocStore::set_global(
 35            cx,
 36            GlobalRustdocStore(Arc::new(Self::new(cx.background_executor().clone()))),
 37        );
 38    }
 39
 40    pub fn new(executor: BackgroundExecutor) -> Self {
 41        let database_future = executor
 42            .spawn({
 43                let executor = executor.clone();
 44                async move {
 45                    RustdocDatabase::new(SUPPORT_DIR.join("docs/rust/rustdoc-db.0.mdb"), executor)
 46                }
 47            })
 48            .then(|result| future::ready(result.map(Arc::new).map_err(Arc::new)))
 49            .boxed()
 50            .shared();
 51
 52        Self {
 53            executor,
 54            database_future,
 55        }
 56    }
 57
 58    pub async fn load(
 59        &self,
 60        crate_name: String,
 61        item_path: Option<String>,
 62    ) -> Result<RustdocDatabaseEntry> {
 63        self.database_future
 64            .clone()
 65            .await
 66            .map_err(|err| anyhow!(err))?
 67            .load(crate_name, item_path)
 68            .await
 69    }
 70
 71    pub fn index(
 72        &self,
 73        crate_name: String,
 74        provider: Box<dyn RustdocProvider + Send + Sync + 'static>,
 75    ) -> Task<Result<()>> {
 76        let database_future = self.database_future.clone();
 77        self.executor.spawn(async move {
 78            let database = database_future.await.map_err(|err| anyhow!(err))?;
 79            let indexer = RustdocIndexer::new(database, provider);
 80
 81            indexer.index(crate_name.clone()).await
 82        })
 83    }
 84
 85    pub fn search(&self, query: String) -> Task<Vec<String>> {
 86        let executor = self.executor.clone();
 87        let database_future = self.database_future.clone();
 88        self.executor.spawn(async move {
 89            if query.is_empty() {
 90                return Vec::new();
 91            }
 92
 93            let Some(database) = database_future.await.map_err(|err| anyhow!(err)).log_err() else {
 94                return Vec::new();
 95            };
 96
 97            let Some(items) = database.keys().await.log_err() else {
 98                return Vec::new();
 99            };
100
101            let candidates = items
102                .iter()
103                .enumerate()
104                .map(|(ix, item_path)| StringMatchCandidate::new(ix, item_path.clone()))
105                .collect::<Vec<_>>();
106
107            let matches = fuzzy::match_strings(
108                &candidates,
109                &query,
110                false,
111                100,
112                &AtomicBool::default(),
113                executor,
114            )
115            .await;
116
117            matches
118                .into_iter()
119                .map(|mat| items[mat.candidate_id].clone())
120                .collect()
121        })
122    }
123}
124
125#[derive(Serialize, Deserialize)]
126pub enum RustdocDatabaseEntry {
127    Crate { docs: String },
128    Item { kind: RustdocItemKind, docs: String },
129}
130
131impl RustdocDatabaseEntry {
132    pub fn docs(&self) -> &str {
133        match self {
134            Self::Crate { docs } | Self::Item { docs, .. } => &docs,
135        }
136    }
137}
138
139pub(crate) struct RustdocDatabase {
140    executor: BackgroundExecutor,
141    env: heed::Env,
142    entries: Database<SerdeBincode<String>, SerdeBincode<RustdocDatabaseEntry>>,
143}
144
145impl RustdocDatabase {
146    pub fn new(path: PathBuf, executor: BackgroundExecutor) -> Result<Self> {
147        std::fs::create_dir_all(&path)?;
148
149        const ONE_GB_IN_BYTES: usize = 1024 * 1024 * 1024;
150        let env = unsafe {
151            heed::EnvOpenOptions::new()
152                .map_size(ONE_GB_IN_BYTES)
153                .max_dbs(1)
154                .open(path)?
155        };
156
157        let mut txn = env.write_txn()?;
158        let entries = env.create_database(&mut txn, Some("rustdoc_entries"))?;
159        txn.commit()?;
160
161        Ok(Self {
162            executor,
163            env,
164            entries,
165        })
166    }
167
168    pub fn keys(&self) -> Task<Result<Vec<String>>> {
169        let env = self.env.clone();
170        let entries = self.entries;
171
172        self.executor.spawn(async move {
173            let txn = env.read_txn()?;
174            let mut iter = entries.iter(&txn)?;
175            let mut keys = Vec::new();
176            while let Some((key, _value)) = iter.next().transpose()? {
177                keys.push(key);
178            }
179
180            Ok(keys)
181        })
182    }
183
184    pub fn load(
185        &self,
186        crate_name: String,
187        item_path: Option<String>,
188    ) -> Task<Result<RustdocDatabaseEntry>> {
189        let env = self.env.clone();
190        let entries = self.entries;
191        let item_path = if let Some(item_path) = item_path {
192            format!("{crate_name}::{item_path}")
193        } else {
194            crate_name
195        };
196
197        self.executor.spawn(async move {
198            let txn = env.read_txn()?;
199            entries
200                .get(&txn, &item_path)?
201                .ok_or_else(|| anyhow!("no docs found for {item_path}"))
202        })
203    }
204
205    pub fn insert(
206        &self,
207        crate_name: String,
208        item: Option<&RustdocItem>,
209        docs: String,
210    ) -> Task<Result<()>> {
211        let env = self.env.clone();
212        let entries = self.entries;
213        let (item_path, entry) = if let Some(item) = item {
214            (
215                format!("{crate_name}::{}", item.display()),
216                RustdocDatabaseEntry::Item {
217                    kind: item.kind,
218                    docs,
219                },
220            )
221        } else {
222            (crate_name, RustdocDatabaseEntry::Crate { docs })
223        };
224
225        self.executor.spawn(async move {
226            let mut txn = env.write_txn()?;
227            entries.put(&mut txn, &item_path, &entry)?;
228            txn.commit()?;
229            Ok(())
230        })
231    }
232}