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