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