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