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