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