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 paths::SUPPORT_DIR;
16use serde::{Deserialize, Serialize};
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 /// Returns whether the crate with the given name is currently being indexed.
75 pub fn is_indexing(&self, crate_name: &CrateName) -> bool {
76 self.indexing_tasks_by_crate.read().contains_key(crate_name)
77 }
78
79 pub async fn load(
80 &self,
81 crate_name: CrateName,
82 item_path: Option<String>,
83 ) -> Result<RustdocDatabaseEntry> {
84 self.database_future
85 .clone()
86 .await
87 .map_err(|err| anyhow!(err))?
88 .load(crate_name, item_path)
89 .await
90 }
91
92 pub fn index(
93 self: Arc<Self>,
94 crate_name: CrateName,
95 provider: Box<dyn RustdocProvider + Send + Sync + 'static>,
96 ) -> Shared<Task<Result<(), Arc<anyhow::Error>>>> {
97 if let Some(existing_task) = self.indexing_tasks_by_crate.read().get(&crate_name) {
98 return existing_task.clone();
99 }
100
101 let indexing_task = self
102 .executor
103 .spawn({
104 let this = self.clone();
105 let crate_name = crate_name.clone();
106 async move {
107 let _finally = util::defer({
108 let this = this.clone();
109 let crate_name = crate_name.clone();
110 move || {
111 this.indexing_tasks_by_crate.write().remove(&crate_name);
112 }
113 });
114
115 let index_task = async {
116 let database = this
117 .database_future
118 .clone()
119 .await
120 .map_err(|err| anyhow!(err))?;
121 let indexer = RustdocIndexer::new(database, provider);
122
123 indexer.index(crate_name.clone()).await
124 };
125
126 index_task.await.map_err(Arc::new)
127 }
128 })
129 .shared();
130
131 self.indexing_tasks_by_crate
132 .write()
133 .insert(crate_name, indexing_task.clone());
134
135 indexing_task
136 }
137
138 pub fn search(&self, query: String) -> Task<Vec<String>> {
139 let executor = self.executor.clone();
140 let database_future = self.database_future.clone();
141 self.executor.spawn(async move {
142 if query.is_empty() {
143 return Vec::new();
144 }
145
146 let Some(database) = database_future.await.map_err(|err| anyhow!(err)).log_err() else {
147 return Vec::new();
148 };
149
150 let Some(items) = database.keys().await.log_err() else {
151 return Vec::new();
152 };
153
154 let candidates = items
155 .iter()
156 .enumerate()
157 .map(|(ix, item_path)| StringMatchCandidate::new(ix, item_path.clone()))
158 .collect::<Vec<_>>();
159
160 let matches = fuzzy::match_strings(
161 &candidates,
162 &query,
163 false,
164 100,
165 &AtomicBool::default(),
166 executor,
167 )
168 .await;
169
170 matches
171 .into_iter()
172 .map(|mat| items[mat.candidate_id].clone())
173 .collect()
174 })
175 }
176}
177
178#[derive(Serialize, Deserialize)]
179pub enum RustdocDatabaseEntry {
180 Crate { docs: String },
181 Item { kind: RustdocItemKind, docs: String },
182}
183
184impl RustdocDatabaseEntry {
185 pub fn docs(&self) -> &str {
186 match self {
187 Self::Crate { docs } | Self::Item { docs, .. } => &docs,
188 }
189 }
190}
191
192pub(crate) struct RustdocDatabase {
193 executor: BackgroundExecutor,
194 env: heed::Env,
195 entries: Database<SerdeBincode<String>, SerdeBincode<RustdocDatabaseEntry>>,
196}
197
198impl RustdocDatabase {
199 pub fn new(path: PathBuf, executor: BackgroundExecutor) -> Result<Self> {
200 std::fs::create_dir_all(&path)?;
201
202 const ONE_GB_IN_BYTES: usize = 1024 * 1024 * 1024;
203 let env = unsafe {
204 heed::EnvOpenOptions::new()
205 .map_size(ONE_GB_IN_BYTES)
206 .max_dbs(1)
207 .open(path)?
208 };
209
210 let mut txn = env.write_txn()?;
211 let entries = env.create_database(&mut txn, Some("rustdoc_entries"))?;
212 txn.commit()?;
213
214 Ok(Self {
215 executor,
216 env,
217 entries,
218 })
219 }
220
221 pub fn keys(&self) -> Task<Result<Vec<String>>> {
222 let env = self.env.clone();
223 let entries = self.entries;
224
225 self.executor.spawn(async move {
226 let txn = env.read_txn()?;
227 let mut iter = entries.iter(&txn)?;
228 let mut keys = Vec::new();
229 while let Some((key, _value)) = iter.next().transpose()? {
230 keys.push(key);
231 }
232
233 Ok(keys)
234 })
235 }
236
237 pub fn load(
238 &self,
239 crate_name: CrateName,
240 item_path: Option<String>,
241 ) -> Task<Result<RustdocDatabaseEntry>> {
242 let env = self.env.clone();
243 let entries = self.entries;
244 let item_path = if let Some(item_path) = item_path {
245 format!("{crate_name}::{item_path}")
246 } else {
247 crate_name.to_string()
248 };
249
250 self.executor.spawn(async move {
251 let txn = env.read_txn()?;
252 entries
253 .get(&txn, &item_path)?
254 .ok_or_else(|| anyhow!("no docs found for {item_path}"))
255 })
256 }
257
258 pub fn insert(
259 &self,
260 crate_name: CrateName,
261 item: Option<&RustdocItem>,
262 docs: String,
263 ) -> Task<Result<()>> {
264 let env = self.env.clone();
265 let entries = self.entries;
266 let (item_path, entry) = if let Some(item) = item {
267 (
268 format!("{crate_name}::{}", item.display()),
269 RustdocDatabaseEntry::Item {
270 kind: item.kind,
271 docs,
272 },
273 )
274 } else {
275 (crate_name.to_string(), RustdocDatabaseEntry::Crate { docs })
276 };
277
278 self.executor.spawn(async move {
279 let mut txn = env.write_txn()?;
280 entries.put(&mut txn, &item_path, &entry)?;
281 txn.commit()?;
282 Ok(())
283 })
284 }
285}