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